Impl enumeration for shards

Analog-Labs · Nov 4, 2024 · 824ccdc · 824ccdc
1 parent cc9cdfc
commit 824ccdc
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 13 deletions.
diff --git a/src/Gateway.sol b/src/Gateway.sol
@@ -8,6 +8,7 @@ import {BranchlessMath} from "./utils/BranchlessMath.sol";
 import {GasUtils} from "./utils/GasUtils.sol";
 import {ERC1967} from "./utils/ERC1967.sol";
 import {UFloat9x56, UFloatMath} from "./utils/Float9x56.sol";
+import {ShardStore} from "./storage/Shards.sol";
 import {IGateway} from "./interfaces/IGateway.sol";
 import {IUpgradable} from "./interfaces/IUpgradable.sol";
 import {IGmpReceiver} from "./interfaces/IGmpReceiver.sol";
@@ -58,6 +59,7 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
     using PrimitiveUtils for address;
     using BranchlessMath for uint256;
     using UFloatMath for UFloat9x56;
+    using ShardStore for ShardStore.MainStorage;
 
     uint8 internal constant SHARD_ACTIVE = (1 << 0); // Shard active bitflag
     uint8 internal constant SHARD_Y_PARITY = (1 << 1); // Pubkey y parity bitflag
@@ -175,8 +177,9 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
         return _messages[id];
     }
 
-    function keyInfo(bytes32 id) external view returns (KeyInfo memory) {
-        return _shards[id];
+    function keyInfo(bytes32 id) external view returns (ShardStore.KeyInfo memory) {
+        ShardStore.MainStorage storage store = ShardStore.getMainStorage();
+        return store.get(id);
     }
 
     function networkId() external view returns (uint16) {

diff --git a/src/storage/Shards.sol b/src/storage/Shards.sol
@@ -63,22 +63,12 @@ library ShardStore {
     error ShardAlreadyRegistered(ShardID id);
     error ShardNotExists(ShardID id);
 
-    function _getMainStorage() internal pure returns (MainStorage storage $) {
+    function getMainStorage() internal pure returns (MainStorage storage $) {
         assembly {
             $.slot := _EIP7201_NAMESPACE
         }
     }
 
-    // /**
-    //  * @dev Derive an ERC-7201 slot from a string (namespace).
-    //  */
-    // function erc7201Slot(string memory namespace) internal pure returns (bytes32 slot) {
-    //     assembly ("memory-safe") {
-    //         mstore(0x00, sub(keccak256(add(namespace, 0x20), mload(namespace)), 1))
-    //         slot := and(keccak256(0x00, 0x20), not(0xff))
-    //     }
-    // }
-
     function asPtr(KeyInfo storage keyInfo) internal pure returns (StoragePtr ptr) {
         assembly {
             ptr := keyInfo.slot
@@ -163,6 +153,19 @@ library ShardStore {
         return _getKeyInfo(ptr);
     }
 
+    /**
+     * @dev Returns the value associated with `key`. O(1).
+     *
+     * Requirements:
+     *
+     * - `key` must be in the map.
+     */
+    function get(MainStorage storage store, bytes32 key) internal view returns (KeyInfo storage) {
+        StoragePtr ptr = store.shards.get(key);
+        require(ptr.isNull() == false, "ShardStore: key not found");
+        return _getKeyInfo(ptr);
+    }
+
     //     /**
     //      * @dev Return the entire set in an array
     //      *