Replace EnumerableSet with OZ EnumerableMap (#56)

Co-authored-by: David Craven <[email protected]>

Author: haider-rs

src/GasUtils.sol

@@ -3,13 +3,13 @@
 
 pragma solidity >=0.8.20;
 
-import {BranchlessMath} from "./utils/BranchlessMath.sol";
+import {PrimitiveUtils} from "./Primitives.sol";
 
 /**
  * @dev Utilities for compute the GMP gas price, gas cost and gas needed.
  */
 library GasUtils {
-    using BranchlessMath for uint256;
+    using PrimitiveUtils for uint256;
 
     /**
      * @dev Compute the amount of gas used by the `GatewayProxy`.

src/Gateway.sol

@@ -4,7 +4,6 @@
 pragma solidity >=0.8.0;
 
 import {Schnorr} from "../lib/frost-evm/sol/Schnorr.sol";
-import {BranchlessMath} from "./utils/BranchlessMath.sol";
 import {GasUtils} from "./GasUtils.sol";
 import {RouteStore} from "./storage/Routes.sol";
 import {ShardStore} from "./storage/Shards.sol";
@@ -32,7 +31,9 @@ contract Gateway is IGateway, UUPSUpgradeable, OwnableUpgradeable {
     using PrimitiveUtils for GmpMessage;
     using PrimitiveUtils for GmpCallback;
     using PrimitiveUtils for address;
-    using BranchlessMath for uint256;
+    using PrimitiveUtils for uint256;
+    using PrimitiveUtils for uint64;
+    using PrimitiveUtils for bool;
     using ShardStore for ShardStore.MainStorage;
     using RouteStore for RouteStore.MainStorage;
     using RouteStore for RouteStore.NetworkInfo;
@@ -221,9 +222,8 @@ contract Gateway is IGateway, UUPSUpgradeable, OwnableUpgradeable {
             uint64 nextNonce = uint64(nonces[msg.sender]++);
 
             // Create GMP message and update nonce
-            GmpMessage memory message = GmpMessage(
-                source, networkId(), destinationAddress, network, gasLimit, nextNonce, data
-            );
+            GmpMessage memory message =
+                GmpMessage(source, networkId(), destinationAddress, network, gasLimit, nextNonce, data);
 
             bytes32 messageId = message.messageId();
             emit GmpCreated(
@@ -276,7 +276,7 @@ contract Gateway is IGateway, UUPSUpgradeable, OwnableUpgradeable {
         // Cap the GMP gas limit to 50% of the block gas limit
         // OBS: we assume the remaining 50% is enough for the Gateway execution, which is a safe assumption
         // once most EVM blockchains have gas limits above 10M and don't need more than 60k gas for the Gateway execution.
-        uint256 gasLimit = BranchlessMath.min(callback.gasLimit, block.gaslimit >> 1);
+        uint256 gasLimit = callback.gasLimit.min(block.gaslimit >> 1);
         unchecked {
             // Add `all but one 64th` to the gas needed, as the defined by EIP-150
             // https://eips.ethereum.org/EIPS/eip-150
@@ -313,8 +313,7 @@ contract Gateway is IGateway, UUPSUpgradeable, OwnableUpgradeable {
         }
 
         // Update GMP status
-        GmpStatus status =
-            GmpStatus(BranchlessMath.ternary(success, uint256(GmpStatus.SUCCESS), uint256(GmpStatus.REVERT)));
+        GmpStatus status = GmpStatus(success.ternary(uint256(GmpStatus.SUCCESS), uint256(GmpStatus.REVERT)));
 
         // Persist gmp execution status on storage
         messages[msgId] = status;
@@ -470,9 +469,7 @@ contract Gateway is IGateway, UUPSUpgradeable, OwnableUpgradeable {
         // Compute the Batch signing hash
         rootHash = PrimitiveUtils.hash(batch.version, batch.batchId, uint256(rootHash));
         bytes32 signingHash = keccak256(
-            abi.encodePacked(
-                "Analog GMP v2", networkId(), bytes32(uint256(uint160(address(this)))), rootHash
-            )
+            abi.encodePacked("Analog GMP v2", networkId(), bytes32(uint256(uint160(address(this)))), rootHash)
         );
 
         // Load shard from storage
@@ -495,7 +492,7 @@ contract Gateway is IGateway, UUPSUpgradeable, OwnableUpgradeable {
             gasUsed += initialGas - gasleft();
 
             // Compute refund amount
-            uint256 refund = BranchlessMath.min(gasUsed * tx.gasprice, address(this).balance);
+            uint256 refund = (gasUsed * tx.gasprice).min(address(this).balance);
 
             // Refund the gas used
             assembly ("memory-safe") {

src/GmpProxy.sol

@@ -5,12 +5,9 @@ pragma solidity >=0.8.0;
 
 import {IGmpReceiver} from "../../src/interfaces/IGmpReceiver.sol";
 import {IGateway} from "../../src/interfaces/IGateway.sol";
-import {BranchlessMath} from "../../src/utils/BranchlessMath.sol";
 import {console} from "forge-std/console.sol";
 
 contract GmpProxy is IGmpReceiver {
-    using BranchlessMath for uint256;
-
     event MessageReceived(GmpMessage msg);
 
     struct GmpMessage {
@@ -32,8 +29,8 @@ contract GmpProxy is IGmpReceiver {
     }
 
     function sendMessage(GmpMessage calldata message) external payable returns (bytes32) {
-        uint256 value = address(this).balance.min(msg.value);
-        return GATEWAY.submitMessage{value: value}(message.dest, message.destNetwork, message.gasLimit, message.data);
+        return
+            GATEWAY.submitMessage{value: msg.value}(message.dest, message.destNetwork, message.gasLimit, message.data);
     }
 
     function onGmpReceived(bytes32 id, uint128 srcNetwork, bytes32 src, uint64 nonce, bytes calldata payload)

src/Primitives.sol

@@ -3,8 +3,6 @@
 
 pragma solidity >=0.8.0;
 
-import {BranchlessMath} from "./utils/BranchlessMath.sol";
-
 uint8 constant GMP_VERSION = 1;
 
 /**
@@ -209,6 +207,74 @@ library PrimitiveUtils {
         }
     }
 
+    /**
+     * @dev Returns the smallest of two numbers.
+     */
+    function min(uint256 x, uint256 y) internal pure returns (uint256) {
+        return ternary(x < y, x, y);
+    }
+
+    /**
+     * @dev Returns the largest of two numbers.
+     */
+    function max(uint256 x, uint256 y) internal pure returns (uint256) {
+        return ternary(x > y, x, y);
+    }
+
+    /**
+     * @dev If `condition` is true returns `a`, otherwise returns `b`.
+     */
+    function ternary(bool condition, uint256 a, uint256 b) internal pure returns (uint256 r) {
+        // branchless select, works because:
+        // b ^ (a ^ b) == a
+        // b ^ 0 == b
+        //
+        // This is better than doing `condition ? a : b` because:
+        // - Consumes less gas
+        // - Constant gas cost regardless the inputs
+        // - Reduces the final bytecode size
+        assembly {
+            r := xor(b, mul(xor(a, b), condition))
+        }
+    }
+
+    /**
+     * @dev If `condition` is true returns `a`, otherwise returns `b`.
+     * see `BranchlessMath.ternary`
+     */
+    function ternaryU64(bool condition, uint64 a, uint64 b) internal pure returns (uint64 r) {
+        assembly {
+            r := xor(b, mul(xor(a, b), condition))
+        }
+    }
+
+    /**
+     * @dev Cast a boolean (false or true) to a uint256 (0 or 1) with no jump.
+     */
+    function toUint(bool b) internal pure returns (uint256 u) {
+        assembly ("memory-safe") {
+            u := iszero(iszero(b))
+        }
+    }
+
+    /**
+     * @dev Aligns `x` to 32 bytes.
+     */
+    function align32(uint256 x) internal pure returns (uint256 r) {
+        unchecked {
+            r = (x + 31) & 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffe0;
+        }
+    }
+
+    /**
+     * @dev Convert byte count to 256bit word count, rounded up.
+     */
+    function toWordCount(uint256 byteCount) internal pure returns (uint256 words) {
+        assembly {
+            words := add(shr(5, byteCount), gt(and(byteCount, 0x1f), 0))
+        }
+    }
+
     function toAddress(bytes32 sender) internal pure returns (address) {
         return address(uint160(uint256(sender)));
     }

src/storage/Routes.sol

@@ -3,19 +3,14 @@
 pragma solidity ^0.8.20;
 
 import {Signature, Route, MAX_PAYLOAD_SIZE} from "../Primitives.sol";
-import {EnumerableSet, Pointer} from "../utils/EnumerableSet.sol";
-import {BranchlessMath} from "../utils/BranchlessMath.sol";
-import {StoragePtr} from "../utils/Pointer.sol";
+import {EnumerableMap} from "@openzeppelin/contracts/utils/structs/EnumerableMap.sol";
 import {GasUtils} from "../GasUtils.sol";
 
 /**
  * @dev EIP-7201 Route's Storage
  */
 library RouteStore {
-    using Pointer for StoragePtr;
-    using Pointer for uint256;
-    using EnumerableSet for EnumerableSet.Map;
-    using BranchlessMath for uint256;
+    using EnumerableMap for EnumerableMap.UintToUintMap;
 
     /**
      * @dev Namespace of the routes storage `analog.one.gateway.routes`.
@@ -69,32 +64,26 @@ library RouteStore {
      * @custom:storage-location erc7201:analog.one.gateway.routes
      */
     struct MainStorage {
-        EnumerableSet.Map routes;
+        EnumerableMap.UintToUintMap routeIds;
+        mapping(uint16 => NetworkInfo) routes;
     }
 
     error RouteNotExists(uint16 id);
     error IndexOutOfBounds(uint256 index);
+    error ZeroGatewayForNewRoute();
+    error InvalidRouteParameters();
 
     function getMainStorage() internal pure returns (MainStorage storage $) {
         assembly {
             $.slot := _EIP7201_NAMESPACE
         }
     }
 
-    /**
-     * @dev Converts a `StoragePtr` into an `NetworkInfo`.
-     */
-    function pointerToRoute(StoragePtr ptr) private pure returns (NetworkInfo storage route) {
-        assembly {
-            route.slot := ptr
-        }
-    }
-
     /**
      * @dev Returns true if the value is in the set. O(1).
      */
     function has(MainStorage storage store, uint16 networkId) internal view returns (bool) {
-        return store.routes.has(bytes32(uint256(networkId)));
+        return store.routeIds.contains(uint256(networkId));
     }
 
     /**
@@ -104,8 +93,11 @@ library RouteStore {
      * already present.
      */
     function getOrAdd(MainStorage storage store, uint16 networkId) private returns (bool, NetworkInfo storage) {
-        (bool success, StoragePtr ptr) = store.routes.tryAdd(bytes32(uint256(networkId)));
-        return (success, pointerToRoute(ptr));
+        bool exists = store.routeIds.contains(uint256(networkId));
+        if (!exists) {
+            store.routeIds.set(uint256(networkId), 1);
+        }
+        return (!exists, store.routes[networkId]);
     }
 
     /**
@@ -115,18 +107,18 @@ library RouteStore {
      * present.
      */
     function remove(MainStorage storage store, uint16 id) internal returns (bool) {
-        StoragePtr ptr = store.routes.remove(bytes32(uint256(id)));
-        if (ptr.isNull()) {
-            return false;
+        bool existed = store.routeIds.remove(uint256(id));
+        if (existed) {
+            delete store.routes[id];
         }
-        return true;
+        return existed;
     }
 
     /**
      * @dev Returns the number of values on the set. O(1).
      */
     function length(MainStorage storage store) internal view returns (uint256) {
-        return store.routes.length();
+        return store.routeIds.length();
     }
 
     /**
@@ -140,11 +132,12 @@ library RouteStore {
      * - `index` must be strictly less than {length}.
      */
     function at(MainStorage storage store, uint256 index) internal view returns (uint16, NetworkInfo storage) {
-        (bytes32 key, StoragePtr value) = store.routes.at(index);
-        if (value.isNull()) {
+        if (index >= store.routeIds.length()) {
             revert IndexOutOfBounds(index);
         }
-        return (uint16(uint256(key)), pointerToRoute(value));
+        (uint256 key,) = store.routeIds.at(index);
+        uint16 networkId = uint16(key);
+        return (networkId, store.routes[networkId]);
     }
 
     /**
@@ -154,32 +147,27 @@ library RouteStore {
      * - `NetworkInfo` must be in the map.
      */
     function get(MainStorage storage store, uint16 id) internal view returns (NetworkInfo storage) {
-        StoragePtr ptr = store.routes.get(bytes32(uint256(id)));
-        if (ptr.isNull()) {
+        if (!store.routeIds.contains(uint256(id))) {
             revert RouteNotExists(id);
         }
-        return pointerToRoute(ptr);
-    }
-
-    /**
-     * @dev Returns the value associated with `NetworkInfo`. O(1).
-     */
-    function tryGet(MainStorage storage store, uint16 id) internal view returns (bool, NetworkInfo storage) {
-        (bool exists, StoragePtr ptr) = store.routes.tryGet(bytes32(uint256(id)));
-        return (exists, pointerToRoute(ptr));
+        return store.routes[id];
     }
 
     function createOrUpdateRoute(MainStorage storage store, Route calldata route) internal {
-        // Update network info
         (bool created, NetworkInfo storage stored) = getOrAdd(store, route.networkId);
-        require((created && route.gateway != bytes32(0)) || !created, "domain separator cannot be zero");
+        if (created) {
+            if (route.gateway == bytes32(0)) revert ZeroGatewayForNewRoute();
+            stored.gateway = route.gateway;
+        }
 
-        // Update gas limit if it's not zero
-        if (route.gasLimit > 0) {
-            stored.gasLimit = route.gasLimit;
+        if (route.relativeGasPriceDenominator == 0 && route.relativeGasPriceNumerator > 0) {
+            revert InvalidRouteParameters();
         }
 
-        // Update relative gas price and base fee if any of them are greater than zero
+        // Update gas limit if it's not zero
+        if (route.gasLimit > 0) stored.gasLimit = route.gasLimit;
+
+        // Update relative gas price and base fee if denominator is greater than zero
         if (route.relativeGasPriceDenominator > 0) {
             stored.relativeGasPriceNumerator = route.relativeGasPriceNumerator;
             stored.relativeGasPriceDenominator = route.relativeGasPriceDenominator;
@@ -208,12 +196,14 @@ library RouteStore {
      * uncallable if the set grows to a point where copying to memory consumes too much gas to fit in a block.
      */
     function listRoutes(MainStorage storage store) internal view returns (Route[] memory) {
-        bytes32[] memory idx = store.routes.keys;
-        Route[] memory routes = new Route[](idx.length);
-        for (uint256 i = 0; i < idx.length; i++) {
-            uint16 networkId = uint16(uint256(idx[i]));
-            (bool success, NetworkInfo storage route) = tryGet(store, networkId);
-            require(success, "route not found");
+        uint256 len = store.routeIds.length();
+        Route[] memory routes = new Route[](len);
+
+        for (uint256 i = 0; i < len; i++) {
+            (uint256 key,) = store.routeIds.at(i);
+            uint16 networkId = uint16(key);
+            NetworkInfo storage route = store.routes[networkId];
+
             routes[i] = Route({
                 networkId: networkId,
                 gasLimit: route.gasLimit,