IMT optimization

contracts/data/IncrementalMerkleTree.sol

@@ -6,7 +6,7 @@ library IncrementalMerkleTree {
     using IncrementalMerkleTree for Tree;
 
     struct Tree {
-        bytes32[][] nodes;
+        bytes32[][] __nodes;
     }
 
     /**
@@ -15,19 +15,30 @@ library IncrementalMerkleTree {
      * @return treeSize size of tree
      */
     function size(Tree storage t) internal view returns (uint256 treeSize) {
-        if (t.height() > 0) {
-            treeSize = t.nodes[0].length;
+        assembly {
+            // assembly block equivalent to:
+            //
+            // if (t.height() > 0) treeSize = t.__nodes[0].length;
+
+            mstore(0x00, t.slot)
+            treeSize := sload(keccak256(0x00, 0x20))
         }
     }
 
     /**
      * @notice query one-indexed height of tree
      * @dev conventional zero-indexed height would require the use of signed integers, so height is one-indexed instead
      * @param t Tree struct storage reference
-     * @return one-indexed height of tree
+     * @return treeHeight one-indexed height of tree
      */
-    function height(Tree storage t) internal view returns (uint256) {
-        return t.nodes.length;
+    function height(Tree storage t) internal view returns (uint256 treeHeight) {
+        assembly {
+            // assembly block equivalent to:
+            //
+            // treeHeight = t.__nodes.length;
+
+            treeHeight := sload(t.slot)
+        }
     }
 
     /**
@@ -36,11 +47,16 @@ library IncrementalMerkleTree {
      * @return hash root hash
      */
     function root(Tree storage t) internal view returns (bytes32 hash) {
-        uint256 treeHeight = t.height();
-
-        if (treeHeight > 0) {
-            unchecked {
-                hash = t.nodes[treeHeight - 1][0];
+        assembly {
+            // assembly block equivalent to:
+            //
+            // if (t.height() > 0) hash = t.__nodes[t.height() - 1][0];
+
+            let treeHeight := sload(t.slot)
+            if gt(treeHeight, 0) {
+                mstore(0x00, t.slot)
+                mstore(0x00, add(keccak256(0x00, 0x20), sub(treeHeight, 1)))
+                hash := sload(keccak256(0x00, 0x20))
             }
         }
     }
@@ -49,7 +65,23 @@ library IncrementalMerkleTree {
         Tree storage t,
         uint256 index
     ) internal view returns (bytes32 hash) {
-        hash = t.nodes[0][index];
+        if (index >= t.size()) {
+            assembly {
+                mstore(0x00, 0x4e487b71)
+                mstore(0x20, 0x32)
+                revert(0x1c, 0x24)
+            }
+        }
+
+        assembly {
+            // assembly block equivalent to:
+            //
+            // hash = t.__nodes[0][index];
+
+            mstore(0x00, t.slot)
+            mstore(0x00, keccak256(0x00, 0x20))
+            hash := sload(add(keccak256(0x00, 0x20), index))
+        }
     }
 
     /**
@@ -58,60 +90,67 @@ library IncrementalMerkleTree {
      * @param hash to add
      */
     function push(Tree storage t, bytes32 hash) internal {
-        unchecked {
-            uint256 treeHeight = t.height();
-            uint256 treeSize = t.size();
+        // index to add to tree
+        uint256 updateIndex = t.size();
 
-            // add new layer if tree is at capacity
+        // update stored tree size
 
-            if (treeSize == (1 << treeHeight) >> 1) {
-                t.nodes.push();
-                treeHeight++;
-            }
+        assembly {
+            mstore(0x00, t.slot)
+            sstore(keccak256(0x00, 0x20), add(updateIndex, 1))
+        }
 
-            // add new columns if rows are full
+        // add new layer if tree is at capacity
 
-            uint256 row;
-            uint256 col = treeSize;
+        uint256 treeHeight = t.height();
 
-            while (row < treeHeight && t.nodes[row].length <= col) {
-                t.nodes[row].push();
-                row++;
-                col >>= 1;
+        if (updateIndex == (1 << treeHeight) >> 1) {
+            // increment tree height in storage
+            assembly {
+                sstore(t.slot, add(treeHeight, 1))
             }
+        }
 
-            // add hash to tree
+        // add hash to tree
 
-            t.set(treeSize, hash);
-        }
+        t.set(updateIndex, hash);
     }
 
     function pop(Tree storage t) internal {
-        uint256 treeHeight = t.height();
-        uint256 treeSize = t.size() - 1;
-
-        // remove layer if tree has excess capacity
+        uint256 treeSize = t.size();
 
-        if (treeSize == (1 << treeHeight) >> 2) {
-            treeHeight--;
-            t.nodes.pop();
+        if (treeSize == 0) {
+            assembly {
+                mstore(0x00, 0x4e487b71)
+                mstore(0x20, 0x32)
+                revert(0x1c, 0x24)
+            }
         }
 
-        // remove columns if rows are too long
+        unchecked {
+            // index to remove from tree
+            uint256 updateIndex = treeSize - 1;
+
+            // update stored tree size
 
-        uint256 row;
-        uint256 col = treeSize;
+            assembly {
+                mstore(0x00, t.slot)
+                sstore(keccak256(0x00, 0x20), updateIndex)
+            }
 
-        while (row < treeHeight && t.nodes[row].length > col) {
-            t.nodes[row].pop();
-            row++;
-            col = (col + 1) >> 1;
-        }
+            // if new tree is full, remove excess layer
+            // if no layer is removed, recalculate hashes
 
-        // recalculate hashes
+            uint256 treeHeight = t.height();
 
-        if (treeSize > 0) {
-            t.set(treeSize - 1, t.at(treeSize - 1));
+            if (updateIndex == (1 << treeHeight) >> 2) {
+                // decrement tree height in storage
+                assembly {
+                    sstore(t.slot, sub(treeHeight, 1))
+                }
+            } else {
+                t.set(updateIndex - 1, t.at(updateIndex - 1));
+            }
         }
     }
 
@@ -122,48 +161,72 @@ library IncrementalMerkleTree {
      * @param hash new hash to add
      */
     function set(Tree storage t, uint256 index, bytes32 hash) internal {
-        unchecked {
-            _set(t.nodes, 0, index, t.height() - 1, hash);
+        uint256 treeSize = t.size();
+
+        if (index >= treeSize) {
+            assembly {
+                mstore(0x00, 0x4e487b71)
+                mstore(0x20, 0x32)
+                revert(0x1c, 0x24)
+            }
         }
+
+        _set(t, 0, index, treeSize, hash);
     }
 
     /**
      * @notice update element in tree and recursively recalculate hashes
-     * @param nodes internal tree structure storage reference
+     * @param t Tree struct storage reference
      * @param rowIndex index of current row to update
      * @param colIndex index of current column to update
-     * @param rootIndex index of root row
+     * @param rowLength length of row at rowIndex
      * @param hash hash to store at current position
      */
     function _set(
-        bytes32[][] storage nodes,
+        Tree storage t,
         uint256 rowIndex,
         uint256 colIndex,
-        uint256 rootIndex,
+        uint256 rowLength,
         bytes32 hash
     ) private {
-        bytes32[] storage row = nodes[rowIndex];
+        bytes32[] storage row;
+
+        assembly {
+            // assembly block equivalent to:
+            //
+            // row = nodes[rowIndex];
+
+            mstore(0x00, t.slot)
+            row.slot := add(keccak256(0x00, 0x20), rowIndex)
+        }
 
-        row[colIndex] = hash;
+        // store hash in array via assembly to avoid array length sload
 
-        if (rowIndex == rootIndex) return;
+        assembly {
+            // assembly block equivalent to:
+            //
+            // row[colIndex] = hash;
+
+            mstore(0x00, row.slot)
+            sstore(add(keccak256(0x00, 0x20), colIndex), hash)
+        }
+
+        if (rowLength == 1) return;
 
         unchecked {
             if (colIndex & 1 == 1) {
                 // sibling is on the left
                 assembly {
-                    mstore(0x00, row.slot)
                     let sibling := sload(
                         add(keccak256(0x00, 0x20), sub(colIndex, 1))
                     )
                     mstore(0x00, sibling)
                     mstore(0x20, hash)
                     hash := keccak256(0x00, 0x40)
                 }
-            } else if (colIndex + 1 < row.length) {
+            } else if (colIndex < rowLength - 1) {
                 // sibling is on the right (and sibling exists)
                 assembly {
-                    mstore(0x00, row.slot)
                     let sibling := sload(
                         add(keccak256(0x00, 0x20), add(colIndex, 1))
                     )
@@ -172,8 +235,8 @@ library IncrementalMerkleTree {
                     hash := keccak256(0x00, 0x40)
                 }
             }
-
-            _set(nodes, rowIndex + 1, colIndex >> 1, rootIndex, hash);
         }
+
+        _set(t, rowIndex + 1, colIndex >> 1, (rowLength + 1) >> 1, hash);
     }
 }

test/data/IncrementalMerkleTree.ts

@@ -140,11 +140,19 @@ describe('IncrementalMerkleTree', () => {
     });
 
     describe('reverts if', () => {
-      it('index is out of bounds', async () => {
+      it('tree is size zero', async () => {
         await expect(
           instance.$at.staticCall(STORAGE_SLOT, 0),
         ).to.be.revertedWithPanic(PANIC_CODES.ARRAY_ACCESS_OUT_OF_BOUNDS);
       });
+
+      it('index is out of bounds', async () => {
+        await instance.$push(STORAGE_SLOT, randomHash());
+
+        await expect(
+          instance.$at.staticCall(STORAGE_SLOT, 1),
+        ).to.be.revertedWithPanic(PANIC_CODES.ARRAY_ACCESS_OUT_OF_BOUNDS);
+      });
     });
   });
 
@@ -196,7 +204,7 @@ describe('IncrementalMerkleTree', () => {
     describe('reverts if', () => {
       it('tree is size zero', async () => {
         await expect(instance.$pop(STORAGE_SLOT)).to.be.revertedWithPanic(
-          PANIC_CODES.ARITHMETIC_OVERFLOW,
+          PANIC_CODES.ARRAY_ACCESS_OUT_OF_BOUNDS,
         );
       });
     });