Migration for AssetsStorage

contracts/src/storage/AssetsStorage.sol

@@ -6,8 +6,6 @@ import {TokenInfo, ParaID} from "../Types.sol";
 
 library AssetsStorage {
     struct Layout {
-        // Token registry by token address
-        mapping(address token => TokenInfo) tokenRegistry;
         address assetHubAgent;
         ParaID assetHubParaID;
         // XCM fee charged by AssetHub for registering a token (DOT)
@@ -16,11 +14,13 @@ library AssetsStorage {
         uint128 assetHubReserveTransferFee;
         // Extra fee for registering a token, to discourage spamming (Ether)
         uint256 registerTokenFee;
+        // Token registry by token address
+        mapping(address token => TokenInfo) tokenRegistry;
         // Token registry by tokenID
         mapping(bytes32 tokenID => TokenInfo) tokenRegistryByID;
     }
 
-    bytes32 internal constant SLOT = keccak256("org.snowbridge.storage.assets");
+    bytes32 internal constant SLOT = keccak256("org.snowbridge.storage.assets.v2");
 
     function layout() internal pure returns (Layout storage $) {
         bytes32 slot = SLOT;

contracts/src/storage/LegacyAssetsStorage.sol

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2023 Snowfork <hello@snowfork.com>
+pragma solidity 0.8.23;
+
+import {ParaID} from "../Types.sol";
+
+library LegacyAssetsStorage {
+    struct TokenInfoLegacy {
+        bool isRegistered;
+        bytes31 __padding;
+    }
+
+    struct Layout {
+        // Legacy token registry by token address
+        mapping(address token => TokenInfoLegacy) tokenRegistry;
+        address assetHubAgent;
+        ParaID assetHubParaID;
+        // XCM fee charged by AssetHub for registering a token (DOT)
+        uint128 assetHubCreateAssetFee;
+        // XCM fee charged by AssetHub for receiving a token from the Gateway (DOT)
+        uint128 assetHubReserveTransferFee;
+        // Extra fee for registering a token, to discourage spamming (Ether)
+        uint256 registerTokenFee;
+    }
+
+    bytes32 internal constant SLOT = keccak256("org.snowbridge.storage.assets");
+
+    function layout() internal pure returns (Layout storage $) {
+        bytes32 slot = SLOT;
+        assembly {
+            $.slot := slot
+        }
+    }
+}

contracts/src/upgrades/rococo/GatewayWithAssetStorageV2.sol

@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2023 Snowfork <hello@snowfork.com>
+pragma solidity 0.8.23;
+
+import "../../Gateway.sol";
+
+import {AssetsStorage} from "../../storage/AssetsStorage.sol";
+import {LegacyAssetsStorage} from "../../storage/LegacyAssetsStorage.sol";
+
+contract GatewayWithAssetStorageV2 is Gateway {
+    constructor(
+        address beefyClient,
+        address agentExecutor,
+        ParaID bridgeHubParaID,
+        bytes32 bridgeHubAgentID,
+        uint8 foreignTokenDecimals
+    ) Gateway(beefyClient, agentExecutor, bridgeHubParaID, bridgeHubAgentID, foreignTokenDecimals) {}
+
+    function initialize(bytes memory data) external override {
+        // Prevent initialization of storage in implementation contract
+        if (ERC1967.load() == address(0)) {
+            revert Unauthorized();
+        }
+
+        address[] memory tokens = abi.decode(data, (address[]));
+
+        LegacyAssetsStorage.Layout storage $ = LegacyAssetsStorage.layout();
+        AssetsStorage.Layout storage $v2 = AssetsStorage.layout();
+
+        $v2.assetHubAgent = $.assetHubAgent;
+        $v2.assetHubParaID = $.assetHubParaID;
+        $v2.assetHubCreateAssetFee = $.assetHubCreateAssetFee;
+        $v2.assetHubReserveTransferFee = $.assetHubReserveTransferFee;
+        $v2.registerTokenFee = $.registerTokenFee;
+        for (uint256 i = 0; i < tokens.length; i++) {
+            $v2.tokenRegistry[tokens[i]].isRegistered = $.tokenRegistry[tokens[i]].isRegistered;
+        }
+    }
+}