diamond improvements

Author: mudgen

src/diamond/DiamondInspectFacet.sol

@@ -6,7 +6,7 @@ pragma solidity >=0.8.30;
  */
 
 interface IFacet {
-    function functionSelectors() external view returns (bytes4[] memory);
+    function packedSelectors() external view returns (bytes memory);
 }
 
 contract DiamondInspectFacet {
@@ -56,6 +56,60 @@ contract DiamondInspectFacet {
         facet = s.facetNodes[_functionSelector].facet;
     }
 
+    /**
+     * @notice Decodes a packed byte stream into a standard bytes4[] array.
+     * @param packed The packed bytes (e.g., from `bytes.concat`).
+     * @return unpacked The standard padded bytes4[] array.
+     */
+    function unpackSelectors(bytes memory packed) internal pure returns (bytes4[] memory unpacked) {
+        /*
+         * Allocate the output array
+        */
+        uint256 count = packed.length / 4;
+        unpacked = new bytes4[](count);
+        /*
+         * Copy from packed to unpacked
+        */
+        assembly ("memory-safe") {
+            /*
+             * 'src' points to the start of the data in the packed array (skip 32-byte length)
+            */
+            let src := add(packed, 32)
+            /*
+             * 'dst' points to the start of the data in the new selectors array (skip 32-byte length)
+             */
+            let dst := add(unpacked, 32)
+            /*
+             * 'end' is the stopping point for the destination pointer
+             */
+            let end := add(dst, mul(count, 32))
+            /*
+             * While 'dst' is less than 'end', keep copying
+            */
+            for {} lt(dst, end) {} {
+                /*
+                 * A. Load 32 bytes from the packed source.
+                 *    We read "dirty" data (neighboring bytes), but it doesn't matter
+                 *    because we truncate it when writing.
+                 */
+                let value := mload(src)
+                /*
+                 * B. Clearn up the value to extract only the 4 bytes we want.
+                 */
+                value := and(value, 0xFFFFFFFF00000000000000000000000000000000000000000000000000000000)
+                /*
+                 * C. Store the value into the destination
+                 */
+                mstore(dst, value)
+                /*
+                 * D. Advance pointers
+                 */
+                src := add(src, 4) // Move forward 4 bytes in packed source
+                dst := add(dst, 32) // Move forward 32 bytes in destination target
+            }
+        }
+    }
+
     /**
      * @notice Gets the function selectors that are handled by the given facet.
      * @dev If facet is not found return empty array.
@@ -64,7 +118,7 @@ contract DiamondInspectFacet {
      */
     function facetFunctionSelectors(address _facet) external view returns (bytes4[] memory facetSelectors) {
         DiamondStorage storage s = getStorage();
-        facetSelectors = IFacet(_facet).functionSelectors();
+        facetSelectors = unpackSelectors(IFacet(_facet).packedSelectors());
         if (facetSelectors.length == 0 || s.facetNodes[facetSelectors[0]].facet == address(0)) {
             facetSelectors = new bytes4[](0);
         }
@@ -105,13 +159,30 @@ contract DiamondInspectFacet {
         facetsAndSelectors = new Facet[](facetList.facetCount);
         for (uint256 i; i < facetList.facetCount; i++) {
             address facet = s.facetNodes[currentSelector].facet;
-            bytes4[] memory facetSelectors = IFacet(facet).functionSelectors();
+            bytes4[] memory facetSelectors = unpackSelectors(IFacet(facet).packedSelectors());
             facetsAndSelectors[i].facet = facet;
             facetsAndSelectors[i].functionSelectors = facetSelectors;
             currentSelector = s.facetNodes[currentSelector].nextFacetNodeId;
         }
     }
 
+    function at(bytes memory selectors, uint256 index) internal pure returns (bytes4 selector) {
+        assembly ("memory-safe") {
+            /**
+             * 1. Calculate Pointer
+             * add(selectors, 32) - skips the length field of the bytes array
+             * shl(2, index) is the same as index * 4 but cheaper
+             */
+            let ptr := add(add(selectors, 32), shl(2, index))
+            /**
+             * 2. Load & Return
+             * We load 32 bytes, but Solidity truncates to 4 bytes automatically
+             * upon return assignment, so masking is unnecessary.
+             */
+            selector := mload(ptr)
+        }
+    }
+
     struct FunctionFacetPair {
         bytes4 selector;
         address facet;
@@ -134,9 +205,13 @@ contract DiamondInspectFacet {
         uint256 selectorCount;
         for (uint256 i; i < facetList.facetCount; i++) {
             address facet = s.facetNodes[currentSelector].facet;
-            bytes4[] memory facetSelectors = IFacet(facet).functionSelectors();
-            for (uint256 selectorIndex; selectorIndex < facetSelectors.length; selectorIndex++) {
-                bytes4 selector = facetSelectors[selectorIndex];
+            bytes memory selectors = IFacet(facet).packedSelectors();
+            uint256 selectorLength;
+            unchecked {
+                selectorLength = selectors.length / 4;
+            }
+            for (uint256 selectorIndex; selectorIndex < selectorLength; selectorIndex++) {
+                bytes4 selector = at(selectors, selectorIndex);
                 pairs[selectorCount] = FunctionFacetPair(selector, facet);
                 unchecked {
                     selectorCount++;
@@ -146,12 +221,13 @@ contract DiamondInspectFacet {
         }
     }
 
-    function functionSelectors() external pure returns (bytes4[] memory selectors) {
-        selectors = new bytes4[](5);
-        selectors[0] = DiamondInspectFacet.facetAddress.selector;
-        selectors[1] = DiamondInspectFacet.facetFunctionSelectors.selector;
-        selectors[2] = DiamondInspectFacet.facetAddresses.selector;
-        selectors[3] = DiamondInspectFacet.facets.selector;
-        selectors[4] = DiamondInspectFacet.functionFacetPairs.selector;
+    function packedSelectors() external pure returns (bytes memory) {
+        return bytes.concat(
+            DiamondInspectFacet.facetAddress.selector,
+            DiamondInspectFacet.facetFunctionSelectors.selector,
+            DiamondInspectFacet.facetAddresses.selector,
+            DiamondInspectFacet.facets.selector,
+            DiamondInspectFacet.functionFacetPairs.selector
+        );
     }
 }

src/diamond/DiamondMod.sol

@@ -51,61 +51,78 @@ error FunctionSelectorsCallFailed(address _facet);
 error NoSelectorsForFacet(address _facet);
 error NoBytecodeAtAddress(address _contractAddress);
 
-function functionSelectors(address _facet) view returns (bytes4[] memory) {
+/**
+ * @notice Emitted when a function is added to a diamond.
+ *
+ * @param _selector The function selector being added.
+ * @param _facet    The facet address that will handle calls to `_selector`.
+ */
+event DiamondFunctionAdded(bytes4 indexed _selector, address indexed _facet);
+
+error CannotAddFunctionToDiamondThatAlreadyExists(bytes4 _selector);
+error NoFacetsToAdd();
+
+function packedSelectors(address _facet) view returns (bytes memory) {
     if (_facet.code.length == 0) {
         revert NoBytecodeAtAddress(_facet);
     }
-    (bool success, bytes memory data) =
-        _facet.staticcall(abi.encodeWithSelector(bytes4(keccak256("functionSelectors()"))));
+    (bool success, bytes memory selectors) =
+        _facet.staticcall(abi.encodeWithSelector(bytes4(keccak256("packedSelectors()"))));
 
     if (success == false) {
         revert FunctionSelectorsCallFailed(_facet);
     }
-    bytes4[] memory selectors = abi.decode(data, (bytes4[]));
-    if (selectors.length == 0) {
+    if (selectors.length < 4) {
         revert NoSelectorsForFacet(_facet);
     }
     return selectors;
 }
 
-/**
- * @notice Emitted when a function is added to a diamond.
- *
- * @param _selector The function selector being added.
- * @param _facet    The facet address that will handle calls to `_selector`.
- */
-event DiamondFunctionAdded(bytes4 indexed _selector, address indexed _facet);
-
-error CannotAddFunctionToDiamondThatAlreadyExists(bytes4 _selector);
-error NoFacetsToAdd();
+function at(bytes memory selectors, uint256 index) pure returns (bytes4 selector) {
+    assembly ("memory-safe") {
+        /**
+         * 1. Calculate Pointer
+         * add(selectors, 32) - skips the length field of the bytes array
+         * shl(2, index) is the same as index * 4 but cheaper
+         */
+        let ptr := add(add(selectors, 32), shl(2, index))
+        /**
+         * 2. Load & Return
+         * We load 32 bytes, but Solidity truncates to 4 bytes automatically
+         * upon return assignment, so masking is unnecessary.
+         */
+        selector := mload(ptr)
+    }
+}
 
-/**
- * @notice Adds facets and their function selectors to the diamond.
- */
 function addFacets(address[] memory _facets) {
     DiamondStorage storage s = getStorage();
     if (_facets.length == 0) {
         return;
     }
     FacetList memory facetList = s.facetList;
     bytes4 prevFacetNodeId = facetList.lastFacetNodeId;
-    bytes4[] memory facetSelectors = functionSelectors(_facets[0]);
+    bytes memory selectors = packedSelectors(_facets[0]);
+    bytes4 currentFacetNodeId = at(selectors, 0);
     if (facetList.facetCount == 0) {
-        facetList.firstFacetNodeId = facetSelectors[0];
+        facetList.firstFacetNodeId = currentFacetNodeId;
+    } else {
+        s.facetNodes[prevFacetNodeId].nextFacetNodeId = currentFacetNodeId;
     }
     for (uint256 i; i < _facets.length; i++) {
+        uint256 selectorsLength;
         uint256 nextI;
         unchecked {
             nextI = i + 1;
-            facetList.selectorCount += uint32(facetSelectors.length);
+            selectorsLength = selectors.length / 4;
+            facetList.selectorCount += uint32(selectorsLength);
         }
-        bytes4[] memory nextFacetSelectors;
+        bytes memory nextSelectors;
         bytes4 nextFacetNodeId;
         if (nextI < _facets.length) {
-            nextFacetSelectors = functionSelectors(_facets[nextI]);
-            nextFacetNodeId = nextFacetSelectors[0];
+            nextSelectors = packedSelectors(_facets[nextI]);
+            nextFacetNodeId = at(nextSelectors, 0);
         }
-        bytes4 currentFacetNodeId = facetSelectors[0];
         address oldFacet = s.facetNodes[currentFacetNodeId].facet;
         if (oldFacet != address(0)) {
             revert CannotAddFunctionToDiamondThatAlreadyExists(currentFacetNodeId);
@@ -114,8 +131,8 @@ function addFacets(address[] memory _facets) {
         s.facetNodes[currentFacetNodeId] = FacetNode(facet, prevFacetNodeId, nextFacetNodeId);
         emit DiamondFunctionAdded(currentFacetNodeId, facet);
 
-        for (uint256 selectorIndex = 1; selectorIndex < facetSelectors.length; selectorIndex++) {
-            bytes4 selector = facetSelectors[selectorIndex];
+        for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) {
+            bytes4 selector = at(selectors, selectorIndex);
             oldFacet = s.facetNodes[selector].facet;
             if (oldFacet != address(0)) {
                 revert CannotAddFunctionToDiamondThatAlreadyExists(selector);
@@ -124,7 +141,8 @@ function addFacets(address[] memory _facets) {
             emit DiamondFunctionAdded(selector, facet);
         }
         prevFacetNodeId = currentFacetNodeId;
-        facetSelectors = nextFacetSelectors;
+        selectors = nextSelectors;
+        currentFacetNodeId = nextFacetNodeId;
     }
     unchecked {
         facetList.facetCount += uint32(_facets.length);