From 511153cf3e77bcc15de219bc848b1ab1146d5b6f Mon Sep 17 00:00:00 2001 From: Nick Barry Date: Mon, 24 Oct 2022 02:00:53 -0500 Subject: [PATCH 1/3] circumvent all Solididity safety features for underlying IMT arrays --- contracts/data/IncrementalMerkleTree.sol | 115 +++++++++++++---------- 1 file changed, 63 insertions(+), 52 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index d8837aba..477a2f70 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -6,7 +6,7 @@ library IncrementalMerkleTree { using IncrementalMerkleTree for Tree; struct Tree { - bytes32[][] nodes; + bytes32[][] __nodes; } /** @@ -15,10 +15,8 @@ library IncrementalMerkleTree { * @return treeSize size of tree */ function size(Tree storage t) internal view returns (uint256 treeSize) { - bytes32[][] storage nodes = t.nodes; - assembly { - mstore(0x00, nodes.slot) + mstore(0x00, t.slot) treeSize := sload(keccak256(0x00, 0x20)) } } @@ -30,10 +28,8 @@ library IncrementalMerkleTree { * @return treeHeight one-indexed height of tree */ function height(Tree storage t) internal view returns (uint256 treeHeight) { - bytes32[][] storage nodes = t.nodes; - assembly { - treeHeight := sload(nodes.slot) + treeHeight := sload(t.slot) } } @@ -43,13 +39,11 @@ 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) { assembly { - mstore(0x00, nodes.slot) + mstore(0x00, t.slot) mstore(0x00, add(keccak256(0x00, 0x20), sub(treeHeight, 1))) hash := sload(keccak256(0x00, 0x20)) } @@ -61,7 +55,15 @@ library IncrementalMerkleTree { view returns (bytes32 hash) { - hash = t.nodes[0][index]; + if (index >= t.size()) { + new bytes32[](0)[1]; + } + + assembly { + mstore(0x00, t.slot) + mstore(0x00, keccak256(0x00, 0x20)) + hash := sload(add(keccak256(0x00, 0x20), index)) + } } /** @@ -70,56 +72,60 @@ library IncrementalMerkleTree { * @param hash to add */ function push(Tree storage t, bytes32 hash) internal { - unchecked { - // index to add to tree - uint256 updateIndex = t.size(); + // index to add to tree + uint256 updateIndex = t.size(); - // add new layer if tree is at capacity + // update stored tree size - if (updateIndex == (1 << t.height()) >> 1) { - t.nodes.push(); - } + 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 = updateIndex; + uint256 treeHeight = t.height(); - while (col == t.nodes[row].length) { - t.nodes[row].push(); - row++; - if (col == 0) break; - 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(updateIndex, hash); - } + t.set(updateIndex, hash); } function pop(Tree storage t) internal { + uint256 treeSize = t.size(); + + if (treeSize == 0) { + new bytes32[](0)[1]; + } + unchecked { // index to remove from tree - uint256 updateIndex = t.size() - 1; + uint256 updateIndex = treeSize - 1; - // remove columns if rows are too long + // update stored tree size - uint256 row; - uint256 col = updateIndex; - - while (col != t.nodes[row].length) { - t.nodes[row].pop(); - row++; - col >>= 1; - if (col == 0) break; + assembly { + mstore(0x00, t.slot) + sstore(keccak256(0x00, 0x20), updateIndex) } // if new tree is full, remove excess layer // if no layer is removed, recalculate hashes - if (updateIndex == (1 << t.height()) >> 2) { - t.nodes.pop(); + uint256 treeHeight = t.height(); + + 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)); } @@ -137,25 +143,36 @@ library IncrementalMerkleTree { uint256 index, bytes32 hash ) internal { - _set(t.nodes, 0, index, t.size(), hash); + uint256 treeSize = t.size(); + + if (index >= treeSize) { + new bytes32[](0)[1]; + } + + _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 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 rowLength, bytes32 hash ) private { - bytes32[] storage row = nodes[rowIndex]; + bytes32[] storage row; + + assembly { + mstore(0x00, t.slot) + row.slot := add(keccak256(0x00, 0x20), rowIndex) + } // store hash in array via assembly to avoid array length sload @@ -191,13 +208,7 @@ library IncrementalMerkleTree { } } - _set( - nodes, - rowIndex + 1, - colIndex >> 1, - (rowLength + 1) >> 1, - hash - ); + _set(t, rowIndex + 1, colIndex >> 1, (rowLength + 1) >> 1, hash); } } } From ad3f47d5aaa09d06dfd3d6014c7d84192924a89e Mon Sep 17 00:00:00 2001 From: NouDaimon Date: Mon, 24 Oct 2022 13:03:44 +0300 Subject: [PATCH 2/3] wrap root functionality in assembly block --- contracts/data/IncrementalMerkleTree.sol | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index 477a2f70..aaae104b 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -39,10 +39,9 @@ library IncrementalMerkleTree { * @return hash root hash */ function root(Tree storage t) internal view returns (bytes32 hash) { - uint256 treeHeight = t.height(); - - if (treeHeight > 0) { - assembly { + assembly { + 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)) From d455bb90f7d1a49ce9adebee91a084718f59f011 Mon Sep 17 00:00:00 2001 From: NouDaimon Date: Mon, 24 Oct 2022 13:52:34 +0300 Subject: [PATCH 3/3] remove unrequired mstore operation --- contracts/data/IncrementalMerkleTree.sol | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/contracts/data/IncrementalMerkleTree.sol b/contracts/data/IncrementalMerkleTree.sol index aaae104b..1fccf35f 100644 --- a/contracts/data/IncrementalMerkleTree.sol +++ b/contracts/data/IncrementalMerkleTree.sol @@ -186,7 +186,6 @@ library IncrementalMerkleTree { if (colIndex & 1 == 1) { // sibling is on the left assembly { - mstore(0x00, row.slot) let sibling := sload( add(keccak256(0x00, 0x20), sub(colIndex, 1)) ) @@ -197,7 +196,6 @@ library IncrementalMerkleTree { } 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)) ) @@ -206,8 +204,8 @@ library IncrementalMerkleTree { hash := keccak256(0x00, 0x40) } } - - _set(t, rowIndex + 1, colIndex >> 1, (rowLength + 1) >> 1, hash); } + + _set(t, rowIndex + 1, colIndex >> 1, (rowLength + 1) >> 1, hash); } }