From c2eae4f78f27c5838ed355a07eb47e91df410aee Mon Sep 17 00:00:00 2001 From: NouDaimon Date: Sun, 23 Oct 2022 23:11:38 +0300 Subject: [PATCH 01/15] remove length check to avoid storage read --- contracts/data/IncrementalMerkleTree.sol | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index c9a68ff1..10dc9f74 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -128,7 +128,7 @@ library IncrementalMerkleTree { bytes32 hash ) internal { unchecked { - _set(t.nodes, 0, index, t.height() - 1, hash); + _set(t.nodes, 0, index, t.height() - 1, t.size(), hash); } } @@ -145,6 +145,7 @@ library IncrementalMerkleTree { uint256 rowIndex, uint256 colIndex, uint256 rootIndex, + uint256 rowSize, bytes32 hash ) private { bytes32[] storage row = nodes[rowIndex]; @@ -165,7 +166,7 @@ library IncrementalMerkleTree { mstore(0x20, hash) hash := keccak256(0x00, 0x40) } - } else if (colIndex + 1 < row.length) { + } else if (colIndex + 1 < rowSize) { // sibling is on the right (and sibling exists) assembly { mstore(0x00, row.slot) @@ -178,7 +179,8 @@ library IncrementalMerkleTree { } } - _set(nodes, rowIndex + 1, colIndex >> 1, rootIndex, hash); + rowSize = rowSize % 2 == 0 ? rowSize >> 1 : (rowSize >> 1) + 1; + _set(nodes, rowIndex + 1, colIndex >> 1, rootIndex, rowSize, hash); } } } From 4a359716914eef8d3c819bfdb327608008d132d3 Mon Sep 17 00:00:00 2001 From: NouDaimon Date: Sun, 23 Oct 2022 23:12:39 +0300 Subject: [PATCH 02/15] add natspec --- contracts/data/IncrementalMerkleTree.sol | 1 + 1 file changed, 1 insertion(+) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 10dc9f74..8af51aaa 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -138,6 +138,7 @@ library IncrementalMerkleTree { * @param rowIndex index of current row to update * @param colIndex index of current column to update * @param rootIndex index of root row + * @param rowSize length of row at rowIndex * @param hash hash to store at current position */ function _set( From 5dcf49a21f1d1fec0d74f2c8fbf580095202c1bd Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 15:17:10 -0500 Subject: [PATCH 03/15] store IMT hashes via assembly --- contracts/data/IncrementalMerkleTree.sol | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 8af51aaa..3d3d7cde 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -151,7 +151,12 @@ library IncrementalMerkleTree { ) 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; From 4c5850b949ee5b7412a3ddad2cc51b0e23a049b7 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 15:18:18 -0500 Subject: [PATCH 04/15] rename rowSize to rowLength, modify lt comparison for clarity --- contracts/data/IncrementalMerkleTree.sol | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 3d3d7cde..ea322857 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -138,7 +138,7 @@ library IncrementalMerkleTree { * @param rowIndex index of current row to update * @param colIndex index of current column to update * @param rootIndex index of root row - * @param rowSize length of row at rowIndex + * @param rowLength length of row at rowIndex * @param hash hash to store at current position */ function _set( @@ -146,7 +146,7 @@ library IncrementalMerkleTree { uint256 rowIndex, uint256 colIndex, uint256 rootIndex, - uint256 rowSize, + uint256 rowLength, bytes32 hash ) private { bytes32[] storage row = nodes[rowIndex]; @@ -172,7 +172,7 @@ library IncrementalMerkleTree { mstore(0x20, hash) hash := keccak256(0x00, 0x40) } - } else if (colIndex + 1 < rowSize) { + } else if (colIndex < rowLength - 1) { // sibling is on the right (and sibling exists) assembly { mstore(0x00, row.slot) @@ -185,8 +185,17 @@ library IncrementalMerkleTree { } } - rowSize = rowSize % 2 == 0 ? rowSize >> 1 : (rowSize >> 1) + 1; - _set(nodes, rowIndex + 1, colIndex >> 1, rootIndex, rowSize, hash); + rowLength = rowLength % 2 == 0 + ? rowLength >> 1 + : (rowLength >> 1) + 1; + _set( + nodes, + rowIndex + 1, + colIndex >> 1, + rootIndex, + rowLength, + hash + ); } } } From af581c571da7d45f97f0d233bc0cdaf787f404b4 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 15:38:44 -0500 Subject: [PATCH 05/15] remove rootIndex variable --- contracts/data/IncrementalMerkleTree.sol | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index ea322857..39b4fee2 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -128,7 +128,7 @@ library IncrementalMerkleTree { bytes32 hash ) internal { unchecked { - _set(t.nodes, 0, index, t.height() - 1, t.size(), hash); + _set(t.nodes, 0, index, t.size(), hash); } } @@ -137,7 +137,6 @@ library IncrementalMerkleTree { * @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 */ @@ -145,7 +144,6 @@ library IncrementalMerkleTree { bytes32[][] storage nodes, uint256 rowIndex, uint256 colIndex, - uint256 rootIndex, uint256 rowLength, bytes32 hash ) private { @@ -158,7 +156,7 @@ library IncrementalMerkleTree { sstore(add(keccak256(0x00, 0x20), colIndex), hash) } - if (rowIndex == rootIndex) return; + if (rowLength == 1) return; unchecked { if (colIndex & 1 == 1) { @@ -188,14 +186,7 @@ library IncrementalMerkleTree { rowLength = rowLength % 2 == 0 ? rowLength >> 1 : (rowLength >> 1) + 1; - _set( - nodes, - rowIndex + 1, - colIndex >> 1, - rootIndex, - rowLength, - hash - ); + _set(nodes, rowIndex + 1, colIndex >> 1, rowLength, hash); } } } From 045c652eef7ef42cc52bc8c8cedf49dd9bb1ede6 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 15:41:38 -0500 Subject: [PATCH 06/15] optimize rowLength calculation --- contracts/data/IncrementalMerkleTree.sol | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 39b4fee2..d5cff824 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -183,10 +183,13 @@ library IncrementalMerkleTree { } } - rowLength = rowLength % 2 == 0 - ? rowLength >> 1 - : (rowLength >> 1) + 1; - _set(nodes, rowIndex + 1, colIndex >> 1, rowLength, hash); + _set( + nodes, + rowIndex + 1, + colIndex >> 1, + (rowLength >> 1) + (rowLength & 1), + hash + ); } } } From dfebf46eb40c8cd49ef7229e5ea4dba19b02399f Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 16:02:27 -0500 Subject: [PATCH 07/15] optimize IMT size function --- contracts/data/IncrementalMerkleTree.sol | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index d5cff824..a399c2e7 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -15,8 +15,11 @@ 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)) } } From cc90a1dea655bee8aed9d0942e80f3d3d53f453b Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 19:34:20 -0500 Subject: [PATCH 08/15] wrap contents of IMT pop function in unchecked block --- contracts/data/IncrementalMerkleTree.sol | 38 +++++++++++++----------- test/data/IncrementalMerkleTree.ts | 2 +- 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index a399c2e7..8a6f2550 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -91,31 +91,33 @@ library IncrementalMerkleTree { } function pop(Tree storage t) internal { - uint256 treeHeight = t.height(); - uint256 treeSize = t.size() - 1; + unchecked { + uint256 treeHeight = t.height(); + uint256 treeSize = t.size() - 1; - // remove layer if tree has excess capacity + // remove layer if tree has excess capacity - if (treeSize == (1 << treeHeight) >> 2) { - treeHeight--; - t.nodes.pop(); - } + if (treeSize == (1 << treeHeight) >> 2) { + treeHeight--; + t.nodes.pop(); + } - // remove columns if rows are too long + // remove columns if rows are too long - uint256 row; - uint256 col = treeSize; + uint256 row; + uint256 col = treeSize; - while (row < treeHeight && t.nodes[row].length > col) { - t.nodes[row].pop(); - row++; - col = (col + 1) >> 1; - } + while (row < treeHeight && t.nodes[row].length > col) { + t.nodes[row].pop(); + row++; + col = (col + 1) >> 1; + } - // recalculate hashes + // recalculate hashes - if (treeSize > 0) { - t.set(treeSize - 1, t.at(treeSize - 1)); + if (treeSize > 0) { + t.set(treeSize - 1, t.at(treeSize - 1)); + } } } diff --git a/test/data/IncrementalMerkleTree.ts b/test/data/IncrementalMerkleTree.ts index c4f3c98a..ca05afa7 100644 --- a/test/data/IncrementalMerkleTree.ts +++ b/test/data/IncrementalMerkleTree.ts @@ -186,7 +186,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, ); }); }); From ba217cd2f80662b7e11a9099932348da8745d501 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 21:01:54 -0500 Subject: [PATCH 09/15] remove unused unchecked block --- contracts/data/IncrementalMerkleTree.sol | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 8a6f2550..14f2e2ea 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -132,9 +132,7 @@ library IncrementalMerkleTree { uint256 index, bytes32 hash ) internal { - unchecked { - _set(t.nodes, 0, index, t.size(), hash); - } + _set(t.nodes, 0, index, t.size(), hash); } /** From 3fdae957d6a06ca6eec657b0386f47aa2e55ad0e Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 21:05:45 -0500 Subject: [PATCH 10/15] do not recalculate hashes if layer is removed on pop --- contracts/data/IncrementalMerkleTree.sol | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 14f2e2ea..43480f35 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -95,13 +95,6 @@ library IncrementalMerkleTree { 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(); - } - // remove columns if rows are too long uint256 row; @@ -113,9 +106,12 @@ library IncrementalMerkleTree { col = (col + 1) >> 1; } - // recalculate hashes + // if new tree is full, remove excess layer + // if no layer is removed, recalculate hashes - if (treeSize > 0) { + if (treeSize == ((1 << treeHeight) >> 2)) { + t.nodes.pop(); + } else { t.set(treeSize - 1, t.at(treeSize - 1)); } } From bff831ed910b8265184fc1c5d440a9631ed7d8f4 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 22:35:48 -0500 Subject: [PATCH 11/15] emphasize prallel structure of push and pop functions --- contracts/data/IncrementalMerkleTree.sol | 29 ++++++++++++------------ 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 43480f35..8766ea5c 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -63,56 +63,57 @@ 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 { unchecked { - uint256 treeHeight = t.height(); - uint256 treeSize = t.size() - 1; + // index to remove from tree + uint256 updateIndex = t.size() - 1; // remove columns if rows are too long 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].pop(); row++; - col = (col + 1) >> 1; + col >>= 1; + if (col == 0) break; } // if new tree is full, remove excess layer // if no layer is removed, recalculate hashes - if (treeSize == ((1 << treeHeight) >> 2)) { + if (updateIndex == (1 << t.height()) >> 2) { t.nodes.pop(); } else { - t.set(treeSize - 1, t.at(treeSize - 1)); + t.set(updateIndex - 1, t.at(updateIndex - 1)); } } } From 761ce875ebabaccd0074c4f12f12c3e0e0813d04 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 22:37:00 -0500 Subject: [PATCH 12/15] optimize recursive row length calculation --- contracts/data/IncrementalMerkleTree.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 8766ea5c..ca022072 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -187,7 +187,7 @@ library IncrementalMerkleTree { nodes, rowIndex + 1, colIndex >> 1, - (rowLength >> 1) + (rowLength & 1), + (rowLength + 1) >> 1, hash ); } From 24b483431ec061dad34578620fb115b03ad96140 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Sun, 23 Oct 2022 23:08:55 -0500 Subject: [PATCH 13/15] replace lt with not eq comparison --- contracts/data/IncrementalMerkleTree.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index ca022072..1db21e67 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -100,7 +100,7 @@ library IncrementalMerkleTree { uint256 row; uint256 col = updateIndex; - while (col < t.nodes[row].length) { + while (col != t.nodes[row].length) { t.nodes[row].pop(); row++; col >>= 1; From 886ba4b4c6794842afdc006ecdee6244bd358db2 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Mon, 24 Oct 2022 00:43:51 -0500 Subject: [PATCH 14/15] add additional at function revert test --- test/data/IncrementalMerkleTree.ts | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/test/data/IncrementalMerkleTree.ts b/test/data/IncrementalMerkleTree.ts index ca05afa7..fd894783 100644 --- a/test/data/IncrementalMerkleTree.ts +++ b/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, + ); + }); }); }); From e3116c30fb1e31c6f9af9f3289fee45c94dca265 Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Mon, 24 Oct 2022 01:00:26 -0500 Subject: [PATCH 15/15] use assembly for height and root functions --- contracts/data/IncrementalMerkleTree.sol | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 1db21e67..d8837aba 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -27,10 +27,14 @@ library IncrementalMerkleTree { * @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) + } } /** @@ -39,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)) } } }