IMT optimization

contracts/data/IncrementalMerkleTree.sol

@@ -15,19 +15,26 @@ 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;
+ bytes32[][] storage nodes = t.nodes;
+
+ assembly {
+ mstore(0x00, nodes.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) {
+ bytes32[][] storage nodes = t.nodes;
+
+ assembly {
+ treeHeight := sload(nodes.slot)
+ }
  }
 
  /**
@@ -36,11 +43,15 @@ library IncrementalMerkleTree {
  * @return hash root hash
  */
  function root(Tree storage t) internal view returns (bytes32 hash) {
+ bytes32[][] storage nodes = t.nodes;
+
  uint256 treeHeight = t.height();
 
  if (treeHeight > 0) {
- unchecked {
- hash = t.nodes[treeHeight - 1][0];
+ assembly {
+ mstore(0x00, nodes.slot)
+ mstore(0x00, add(keccak256(0x00, 0x20), sub(treeHeight, 1)))
+ hash := sload(keccak256(0x00, 0x20))
  }
  }
  }
@@ -59,59 +70,58 @@ library IncrementalMerkleTree {
  */
  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
 
- if (treeSize == (1 << treeHeight) >> 1) {
+ if (updateIndex == (1 << t.height()) >> 1) {
  t.nodes.push();
- treeHeight++;
  }
 
  // add new columns if rows are full
 
  uint256 row;
- uint256 col = treeSize;
+ uint256 col = updateIndex;
 
- while (row < treeHeight && t.nodes[row].length <= col) {
+ while (col == t.nodes[row].length) {
  t.nodes[row].push();
  row++;
+ if (col == 0) break;
  col >>= 1;
  }
 
  // 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
-
- if (treeSize == (1 << treeHeight) >> 2) {
- treeHeight--;
- t.nodes.pop();
- }
+ unchecked {
+ // index to remove from tree
+ uint256 updateIndex = t.size() - 1;
 
- // remove columns if rows are too long
+  // remove columns if rows are too long
 
- uint256 row;
- uint256 col = treeSize;
+  uint256 row;
+  uint256 col = updateIndex;
 
- while (row < treeHeight && t.nodes[row].length > col) {
- t.nodes[row].pop();
- row++;
- col = (col + 1) >> 1;
- }
+ while (col != t.nodes[row].length) {
+ t.nodes[row].pop();
+ row++;
+ col >>= 1;
+ if (col == 0) break;
+ }
 
- // recalculate hashes
+ // if new tree is full, remove excess layer
+ // if no layer is removed, recalculate hashes
 
- if (treeSize > 0) {
- t.set(treeSize - 1, t.at(treeSize - 1));
+ if (updateIndex == (1 << t.height()) >> 2) {
+ t.nodes.pop();
+ } else {
+ t.set(updateIndex - 1, t.at(updateIndex - 1));
+ }
  }
  }
 
@@ -122,31 +132,34 @@ 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);
- }
+ _set(t.nodes, 0, index, t.size(), hash);
  }
 
  /**
  * @notice update element in tree and recursively recalculate hashes
  * @param nodes internal tree structure 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,
  uint256 rowIndex,
  uint256 colIndex,
- uint256 rootIndex,
+ uint256 rowLength,
  bytes32 hash
  ) private {
  bytes32[] storage row = nodes[rowIndex];
 
- row[colIndex] = hash;
+ // store hash in array via assembly to avoid array length sload
+
+ assembly {
+ mstore(0x00, row.slot)
+ sstore(add(keccak256(0x00, 0x20), colIndex), hash)
+ }
 
- if (rowIndex == rootIndex) return;
+ if (rowLength == 1) return;
 
  unchecked {
  if (colIndex & 1 == 1) {
@@ -160,7 +173,7 @@ library IncrementalMerkleTree {
  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)
@@ -173,7 +186,13 @@ library IncrementalMerkleTree {
  }
  }
 
- _set(nodes, rowIndex + 1, colIndex >> 1, rootIndex, hash);
+ _set(
+ nodes,
+ rowIndex + 1,
+ colIndex >> 1,
+ (rowLength + 1) >> 1,
+ hash
+ );
  }
  }
 }

test/data/IncrementalMerkleTree.ts

@@ -134,11 +134,19 @@ describe('IncrementalMerkleTree', function () {
  });
 
  describe('reverts if', () => {
- it('index is out of bounds', async () => {
+ it('tree is size zero', async () => {
  await expect(instance.callStatic.at(0)).to.be.revertedWithPanic(
  PANIC_CODES.ARRAY_ACCESS_OUT_OF_BOUNDS,
  );
  });
+
+ it('index is out of bounds', async () => {
+ await instance.push(randomHash());
+
+ await expect(instance.callStatic.at(1)).to.be.revertedWithPanic(
+ PANIC_CODES.ARRAY_ACCESS_OUT_OF_BOUNDS,
+ );
+ });
  });
  });
 
@@ -186,7 +194,7 @@ describe('IncrementalMerkleTree', function () {
  describe('reverts if', () => {
  it('tree is size zero', async () => {
  await expect(instance.pop()).to.be.revertedWithPanic(
- PANIC_CODES.ARITHMETIC_UNDER_OR_OVERFLOW,
+ PANIC_CODES.ARRAY_ACCESS_OUT_OF_BOUNDS,
  );
  });
  });