optimize array access & remove depth/length constrains

contracts/mocks/MerkleTreeMock.sol

@@ -9,14 +9,22 @@ contract MerkleTreeMock {
 
     MerkleTree.TreeWithHistory private tree;
 
-    constructor(uint32 _depth, uint32 _length) {
+    constructor(uint256 _depth, uint256 _length) {
         tree.initialize(_depth, _length);
     }
 
-    function insert(bytes32 leaf) public returns (uint32) {
+    function insert(bytes32 leaf) public returns (uint256) {
         return tree.insert(leaf);
     }
 
+    function getDepth() public view returns (uint256) {
+        return tree.getDepth();
+    }
+
+    function getLength() public view returns (uint256) {
+        return tree.getLength();
+    }
+
     function getLastRoot() public view returns (bytes32) {
         return tree.getLastRoot();
     }
@@ -26,19 +34,11 @@ contract MerkleTreeMock {
     }
 
     // internal state
-    function depth() public view returns (uint32) {
-        return tree.depth;
-    }
-
-    function length() public view returns (uint32) {
-        return tree.length;
-    }
-
-    function currentRootIndex() public view returns (uint32) {
+    function currentRootIndex() public view returns (uint256) {
         return tree.currentRootIndex;
     }
 
-    function nextLeafIndex() public view returns (uint32) {
+    function nextLeafIndex() public view returns (uint256) {
         return tree.nextLeafIndex;
     }
 

contracts/utils/structs/MerkleTree.sol

@@ -2,20 +2,21 @@
 
 pragma solidity ^0.8.0;
 
+// TODO: replace this import in favor or "../Array.sol" when #3589 is merged
+import "../StorageSlot.sol";
+
 error Full();
 
 library MerkleTree {
-    uint8 private constant MAX_DEPTH = 32;
+    uint256 private constant MAX_DEPTH = 256;
 
     struct TreeWithHistory {
+        uint256 currentRootIndex;
+        uint256 nextLeafIndex;
+        bytes32[] filledSubtrees;
+        bytes32[] zeros;
+        bytes32[] roots;
         function(bytes32, bytes32) view returns (bytes32) fnHash;
-        uint32 depth;
-        uint32 length;
-        uint32 currentRootIndex;
-        uint32 nextLeafIndex;
-        bytes32[MAX_DEPTH] filledSubtrees;
-        bytes32[MAX_DEPTH] zeros;
-        bytes32[2**MAX_DEPTH] roots;
     }
 
     /**
@@ -27,83 +28,96 @@ library MerkleTree {
      */
     function initialize(
         TreeWithHistory storage self,
-        uint32 depth,
-        uint32 length,
+        uint256 depth,
+        uint256 length,
         bytes32 zero,
         function(bytes32, bytes32) view returns (bytes32) fnHash
     ) internal {
         require(depth <= MAX_DEPTH);
 
-        self.depth = depth;
-        self.length = length;
+        setLength(self.filledSubtrees, depth);
+        setLength(self.zeros, depth);
+        setLength(self.roots, length);
         self.fnHash = fnHash;
 
         bytes32 currentZero = zero;
         for (uint32 i = 0; i < depth; ++i) {
-            self.zeros[i] = self.filledSubtrees[i] = currentZero;
+            unsafeAccess(self.zeros, i).value = currentZero;
+            unsafeAccess(self.filledSubtrees, i).value = currentZero;
             currentZero = fnHash(currentZero, currentZero);
         }
 
         // Insert the first root
-        self.roots[0] = currentZero;
+        unsafeAccess(self.roots, 0).value = currentZero;
     }
 
     /**
      * @dev Insert a new leaf in the tree, compute the new root, and store that new root in the history.
      *
-     * For depth < 32, reverts if the MerkleTree is already full.
-     * For depth = 32, reverts when trying to populate the last leaf (nextLeafIndex increment overflow).
-     *
-     * Said differently:
-     * `2 ** depth` entries can be inserted into trees with depth < 32.
-     * `2 ** depth - 1` entries can be inserted into trees with depth = 32.
+     * Note: because of the `nextLeafIndex` overflow, trees of depth 256 can only store 2**256-1 items.
+     * This should never be an issue in practice.
      */
-    function insert(TreeWithHistory storage self, bytes32 leaf) internal returns (uint32) {
+    function insert(TreeWithHistory storage self, bytes32 leaf) internal returns (uint256) {
         // cache read
-        uint32 depth = self.depth;
+        uint256 depth = self.zeros.length;
 
         // Get leaf index
-        uint32 leafIndex = self.nextLeafIndex++;
+        uint256 leafIndex = self.nextLeafIndex++;
 
-        // Check if tree is full.
+        // Check if tree is full. This check will not work if depth == 256 and the tree is full, but realistically
+        // this will never happen.
         if (leafIndex == 1 << depth) revert Full();
 
         // Rebuild branch from leaf to root
-        uint32 currentIndex = leafIndex;
+        uint256 currentIndex = leafIndex;
         bytes32 currentLevelHash = leaf;
         for (uint32 i = 0; i < depth; i++) {
             // Reaching the parent node, is currentLevelHash the left child?
             bool isLeft = currentIndex % 2 == 0;
 
             // If so, next time we will come from the right, so we need to save it
             if (isLeft) {
-                self.filledSubtrees[i] = currentLevelHash;
+                unsafeAccess(self.filledSubtrees, i).value = currentLevelHash;
             }
 
             // Compute the node hash by hasing the current hash with either:
             // - the last value for this level
             // - the zero for this level
             currentLevelHash = self.fnHash(
-                isLeft ? currentLevelHash : self.filledSubtrees[i],
-                isLeft ? self.zeros[i] : currentLevelHash
+                isLeft ? currentLevelHash : unsafeAccess(self.filledSubtrees, i).value,
+                isLeft ? unsafeAccess(self.zeros, i).value : currentLevelHash
             );
 
             // update node index
             currentIndex >>= 1;
         }
 
         // Record new root
-        self.currentRootIndex = (self.currentRootIndex + 1) % self.length;
-        self.roots[self.currentRootIndex] = currentLevelHash;
+        self.currentRootIndex = (self.currentRootIndex + 1) % self.roots.length;
+        unsafeAccess(self.roots, self.currentRootIndex).value = currentLevelHash;
 
         return leafIndex;
     }
 
+    /**
+     * @dev Tree's depth (set at initialization)
+     */
+    function getDepth(TreeWithHistory storage self) internal view returns (uint256) {
+        return self.zeros.length;
+    }
+
+    /**
+     * @dev History length (set at initialization)
+     */
+    function getLength(TreeWithHistory storage self) internal view returns (uint256) {
+        return self.roots.length;
+    }
+
     /**
      * @dev Return the current root of the tree.
      */
     function getLastRoot(TreeWithHistory storage self) internal view returns (bytes32) {
-        return self.roots[self.currentRootIndex];
+        return unsafeAccess(self.roots, self.currentRootIndex).value;
     }
 
     /**
@@ -116,23 +130,43 @@ library MerkleTree {
 
         // cache as uint256 (avoid overflow)
         uint256 currentRootIndex = self.currentRootIndex;
-        uint256 length = self.length;
+        uint256 length = self.roots.length;
 
         // search
         for (uint256 i = length; i > 0; --i) {
-            if (root == self.roots[(currentRootIndex + i) % length]) {
+            if (root == unsafeAccess(self.roots, (currentRootIndex + i) % length).value) {
                 return true;
             }
         }
 
         return false;
     }
 
+    /**
+     * @dev Helper to set the length of an dynamic array. Directly writting to `.length` is forbiden.
+     */
+    function setLength(bytes32[] storage array, uint256 len) private {
+        assembly {
+            sstore(array.slot, len)
+        }
+    }
+
+    // TODO: this is part of PR#3589 (in the Arrays library)
+    function unsafeAccess(bytes32[] storage arr, uint256 pos) internal pure returns (StorageSlot.Bytes32Slot storage) {
+        bytes32 slot;
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0, arr.slot)
+            slot := add(keccak256(0, 0x20), pos)
+        }
+        return StorageSlot.getBytes32Slot(slot);
+    }
+
     // Default hash
     function initialize(
         TreeWithHistory storage self,
-        uint32 depth,
-        uint32 length
+        uint256 depth,
+        uint256 length
     ) internal {
         return initialize(self, depth, length, bytes32(0), _hashPair);
     }

test/utils/structs/Merkletree.test.js

@@ -17,8 +17,8 @@ describe('Merklee tree', function () {
     const leafs = Array(2 ** DEPTH).fill(constants.ZERO_BYTES32);
     const merkleTree = new MerkleTree(leafs, keccak256, { sortPairs: true });
 
-    expect(await this.contract.depth()).to.be.bignumber.equal(DEPTH);
-    expect(await this.contract.length()).to.be.bignumber.equal(LENGTH);
+    expect(await this.contract.getDepth()).to.be.bignumber.equal(DEPTH);
+    expect(await this.contract.getLength()).to.be.bignumber.equal(LENGTH);
     expect(await this.contract.currentRootIndex()).to.be.bignumber.equal('0');
     expect(await this.contract.nextLeafIndex()).to.be.bignumber.equal('0');