Skip to content

Commit 66b4e6a

Browse files
heueristikheueristik
authored
Merge pull request #610 from anoma/heueristik/test-only-tree-root-helpers
refactor(contracts): move the test-only root helpers out of MerkleTree
2 parents 6cc459e + 28d4b10 commit 66b4e6a

5 files changed

Lines changed: 81 additions & 68 deletions

File tree

‎contracts/src/libs/MerkleTree.sol‎

Lines changed: 0 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@
22
pragma solidity ^0.8.30;
33

44
import {Arrays} from "@openzeppelin-contracts-5.7.0/utils/Arrays.sol";
5-
import {Math} from "@openzeppelin-contracts-5.7.0/utils/math/Math.sol";
6-
import {SafeCast} from "@openzeppelin-contracts-5.7.0/utils/math/SafeCast.sol";
75

86
import {SHA256} from "../libs/SHA256.sol";
97

@@ -14,8 +12,6 @@ import {SHA256} from "../libs/SHA256.sol";
1412
/// (https://github.com/OpenZeppelin/openzeppelin-contracts/blob/v5.4.0/contracts/utils/structs/MerkleTree.sol).
1513
/// @custom:security-contact security@anoma.foundation
1614
library MerkleTree {
17-
using SafeCast for uint256;
18-
1915
struct Tree {
2016
uint256 _nextLeafIndex;
2117
bytes32[] _sides;
@@ -117,50 +113,4 @@ library MerkleTree {
117113
function isLeftChild(uint256 index) internal pure returns (bool isLeft) {
118114
isLeft = (index & 1) == 0;
119115
}
120-
121-
/// @notice Computes the root of a Merkle tree.
122-
/// @param leaves The leaves of the tree.
123-
/// @param treeDepth The depth of the tree.
124-
/// @return root The computed root.
125-
/// @dev This method should only be used for trees with low depth.
126-
function computeRoot(bytes32[] memory leaves, uint8 treeDepth) internal pure returns (bytes32 root) {
127-
uint256 treeCapacity = uint256(1) << treeDepth; // 2^treeDepth
128-
129-
// Create array of full leaf set with padding if necessary
130-
bytes32[] memory nodes = new bytes32[](treeCapacity);
131-
for (uint256 i = 0; i < treeCapacity; ++i) {
132-
if (i < leaves.length) {
133-
nodes[i] = leaves[i];
134-
} else {
135-
nodes[i] = SHA256.EMPTY_HASH;
136-
}
137-
}
138-
139-
// Build the tree upward
140-
uint256 currentLevelCapacity = treeCapacity;
141-
while (currentLevelCapacity > 1) {
142-
currentLevelCapacity /= 2;
143-
144-
for (uint256 i = 0; i < currentLevelCapacity; ++i) {
145-
nodes[i] = SHA256.hash(nodes[2 * i], nodes[2 * i + 1]);
146-
}
147-
}
148-
149-
root = nodes[0];
150-
}
151-
152-
/// @notice Computes the root of a Merkle tree using the minimal tree depth to fit all leaves.
153-
/// @param leaves The leaves of the tree.
154-
/// @return root The computed root.
155-
/// @dev This method should only be used for trees with low depth.
156-
function computeRoot(bytes32[] memory leaves) internal pure returns (bytes32 root) {
157-
root = MerkleTree.computeRoot({leaves: leaves, treeDepth: computeMinimalTreeDepth(leaves.length)});
158-
}
159-
160-
/// @notice Computes the minimal required tree depth for a number of leaves.
161-
/// @param leavesCount The number of leaves.
162-
/// @return treeDepth The minimal required tree depth.
163-
function computeMinimalTreeDepth(uint256 leavesCount) internal pure returns (uint8 treeDepth) {
164-
treeDepth = Math.log2({value: leavesCount, rounding: Math.Rounding.Ceil}).toUint8();
165-
}
166116
}

‎contracts/test/examples/MerkleTree.e.sol‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
// SPDX-License-Identifier: MIT
22
pragma solidity ^0.8.30;
33

4-
import {MerkleTree} from "../../src/libs/MerkleTree.sol";
54
import {SHA256} from "../../src/libs/SHA256.sol";
5+
import {MerkleTreeReference} from "../libs/MerkleTreeReference.sol";
66

77
contract MerkleTreeExample {
8-
using MerkleTree for bytes32[];
8+
using MerkleTreeReference for bytes32[];
99

1010
uint256 internal constant _N_LEAVES = 7;
1111
uint256 internal constant _N_ROOTS = 8;
@@ -104,7 +104,7 @@ contract MerkleTreeExample {
104104

105105
_heightTwoNodes[3] = _calculateNextLevel(_heightOneNodes[3]);
106106

107-
_roots[3] = MerkleTree.computeRoot(_leaves[3]);
107+
_roots[3] = MerkleTreeReference.computeRoot(_leaves[3]);
108108

109109
_siblings[3] = new bytes32[][](3);
110110

‎contracts/test/libs/MerkleTreeReference.sol‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity ^0.8.30;
3+
4+
import {Math} from "@openzeppelin-contracts-5.7.0/utils/math/Math.sol";
5+
import {SafeCast} from "@openzeppelin-contracts-5.7.0/utils/math/SafeCast.sol";
6+
7+
import {SHA256} from "../../src/libs/SHA256.sol";
8+
9+
/// @notice Computes a Merkle root from a complete leaf set, the straightforward way. The protocol adapter never does
10+
/// this: it maintains the root incrementally in `MerkleTree`. Tests use this to obtain a root the incremental
11+
/// implementation must reproduce.
12+
library MerkleTreeReference {
13+
using SafeCast for uint256;
14+
15+
/// @notice Computes the root of a Merkle tree.
16+
/// @param leaves The leaves of the tree.
17+
/// @param treeDepth The depth of the tree.
18+
/// @return root The computed root.
19+
/// @dev This method should only be used for trees with low depth.
20+
function computeRoot(bytes32[] memory leaves, uint8 treeDepth) internal pure returns (bytes32 root) {
21+
uint256 treeCapacity = uint256(1) << treeDepth; // 2^treeDepth
22+
23+
// Create array of full leaf set with padding if necessary
24+
bytes32[] memory nodes = new bytes32[](treeCapacity);
25+
for (uint256 i = 0; i < treeCapacity; ++i) {
26+
if (i < leaves.length) {
27+
nodes[i] = leaves[i];
28+
} else {
29+
nodes[i] = SHA256.EMPTY_HASH;
30+
}
31+
}
32+
33+
// Build the tree upward
34+
uint256 currentLevelCapacity = treeCapacity;
35+
while (currentLevelCapacity > 1) {
36+
currentLevelCapacity /= 2;
37+
38+
for (uint256 i = 0; i < currentLevelCapacity; ++i) {
39+
nodes[i] = SHA256.hash(nodes[2 * i], nodes[2 * i + 1]);
40+
}
41+
}
42+
43+
root = nodes[0];
44+
}
45+
46+
/// @notice Computes the root of a Merkle tree using the minimal tree depth to fit all leaves.
47+
/// @param leaves The leaves of the tree.
48+
/// @return root The computed root.
49+
/// @dev This method should only be used for trees with low depth.
50+
function computeRoot(bytes32[] memory leaves) internal pure returns (bytes32 root) {
51+
root = computeRoot({leaves: leaves, treeDepth: computeMinimalTreeDepth(leaves.length)});
52+
}
53+
54+
/// @notice Computes the minimal required tree depth for a number of leaves.
55+
/// @param leavesCount The number of leaves.
56+
/// @return treeDepth The minimal required tree depth.
57+
function computeMinimalTreeDepth(uint256 leavesCount) internal pure returns (uint8 treeDepth) {
58+
treeDepth = Math.log2({value: leavesCount, rounding: Math.Rounding.Ceil}).toUint8();
59+
}
60+
}

‎contracts/test/libs/TxGen.sol‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,13 @@ import {RiscZeroMockVerifier} from "risc0-risc0-ethereum-3.0.1/contracts/src/tes
66

77
import {IProtocolAdapter} from "../../src/interfaces/IProtocolAdapter.sol";
88
import {DeltaProof} from "../../src/libs/DeltaProof.sol";
9-
import {MerkleTree} from "../../src/libs/MerkleTree.sol";
109
import {SHA256} from "../../src/libs/SHA256.sol";
1110
import {VerifyingKeys} from "../../src/libs/VerifyingKeys.sol";
1211
import {DeltaGen} from "./DeltaGen.sol";
12+
import {MerkleTreeReference} from "./MerkleTreeReference.sol";
1313

1414
library TxGen {
15-
using MerkleTree for bytes32[];
15+
using MerkleTreeReference for bytes32[];
1616

1717
/// @notice The resource object constituting the atomic unit of state in the Anoma protocol.
1818
/// @param logicRef The hash of the resource logic function.

‎contracts/test/state/MerkleTree.t.sol‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,11 @@ import {Test} from "forge-std-1.16.2/src/Test.sol";
66
import {MerkleTree} from "./../../src/libs/MerkleTree.sol";
77
import {SHA256} from "./../../src/libs/SHA256.sol";
88
import {MerkleTreeExample} from "./../examples/MerkleTree.e.sol";
9+
import {MerkleTreeReference} from "./../libs/MerkleTreeReference.sol";
910

1011
contract MerkleTreeTest is Test, MerkleTreeExample {
1112
using MerkleTree for MerkleTree.Tree;
12-
using MerkleTree for bytes32[];
13+
using MerkleTreeReference for bytes32[];
1314
using OzMerkleTree for OzMerkleTree.Bytes32PushTree;
1415

1516
MerkleTree.Tree internal _merkleTree;
@@ -114,10 +115,10 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
114115
function test_compare_le() public pure {
115116
assertEq(
116117
_computeMinimalTreeDepthNaive(0),
117-
MerkleTree.computeMinimalTreeDepth(0),
118+
MerkleTreeReference.computeMinimalTreeDepth(0),
118119
"naive and optimized should match for 0 leaves"
119120
);
120-
assertEq(MerkleTree.computeMinimalTreeDepth(0), 0, "minimal tree depth for 0 leaves should be 0");
121+
assertEq(MerkleTreeReference.computeMinimalTreeDepth(0), 0, "minimal tree depth for 0 leaves should be 0");
121122
}
122123

123124
function test_computeMinimalTreeDepth_computes_the_right_tree_depths() public pure {
@@ -126,7 +127,9 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
126127

127128
for (uint256 i = 0; i < depths.length; ++i) {
128129
assertEq(
129-
MerkleTree.computeMinimalTreeDepth({leavesCount: i}), depths[i], "tree depth should match expected"
130+
MerkleTreeReference.computeMinimalTreeDepth({leavesCount: i}),
131+
depths[i],
132+
"tree depth should match expected"
130133
);
131134
}
132135
}
@@ -137,7 +140,7 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
137140
for (uint256 i = 0; i < maxLeafCount; ++i) {
138141
assertEq(
139142
_computeMinimalTreeDepthNaive({leavesCount: i}),
140-
MerkleTree.computeMinimalTreeDepth({leavesCount: i}),
143+
MerkleTreeReference.computeMinimalTreeDepth({leavesCount: i}),
141144
"naive and optimized implementations should match"
142145
);
143146
}
@@ -152,21 +155,21 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
152155

153156
assertEq(
154157
_computeMinimalTreeDepthNaive(powerOfTwo - 1),
155-
MerkleTree.computeMinimalTreeDepth(powerOfTwo - 1),
158+
MerkleTreeReference.computeMinimalTreeDepth(powerOfTwo - 1),
156159
"should match for power of 2 minus 1"
157160
);
158161

159162
// Test power of 2
160163
assertEq(
161164
_computeMinimalTreeDepthNaive(powerOfTwo),
162-
MerkleTree.computeMinimalTreeDepth(powerOfTwo),
165+
MerkleTreeReference.computeMinimalTreeDepth(powerOfTwo),
163166
"should match for power of 2"
164167
);
165168

166169
// Test power of 2 + 1
167170
assertEq(
168171
_computeMinimalTreeDepthNaive(powerOfTwo + 1),
169-
MerkleTree.computeMinimalTreeDepth(powerOfTwo + 1),
172+
MerkleTreeReference.computeMinimalTreeDepth(powerOfTwo + 1),
170173
"should match for power of 2 plus 1"
171174
);
172175
}
@@ -179,7 +182,7 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
179182
for (uint256 i = 0; i < testCases.length; i++) {
180183
assertEq(
181184
_computeMinimalTreeDepthNaive(testCases[i]),
182-
MerkleTree.computeMinimalTreeDepth(testCases[i]),
185+
MerkleTreeReference.computeMinimalTreeDepth(testCases[i]),
183186
"should match for large values"
184187
);
185188
}
@@ -190,7 +193,7 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
190193
leavesCount = bound(leavesCount, 0, type(uint128).max);
191194

192195
uint8 original = _computeMinimalTreeDepthNaive(leavesCount);
193-
uint8 optimized = MerkleTree.computeMinimalTreeDepth(leavesCount);
196+
uint8 optimized = MerkleTreeReference.computeMinimalTreeDepth(leavesCount);
194197

195198
assertEq(original, optimized, "Implementations must match");
196199
}
@@ -200,7 +203,7 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
200203
leavesCount = bound(leavesCount, 0, 1000);
201204

202205
uint8 original = _computeMinimalTreeDepthNaive(leavesCount);
203-
uint8 optimized = MerkleTree.computeMinimalTreeDepth(leavesCount);
206+
uint8 optimized = MerkleTreeReference.computeMinimalTreeDepth(leavesCount);
204207

205208
assertEq(original, optimized, "Implementations must match");
206209
}
@@ -210,7 +213,7 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
210213
leavesCount = bound(leavesCount, 1000, 1000000);
211214

212215
uint8 original = _computeMinimalTreeDepthNaive(leavesCount);
213-
uint8 optimized = MerkleTree.computeMinimalTreeDepth(leavesCount);
216+
uint8 optimized = MerkleTreeReference.computeMinimalTreeDepth(leavesCount);
214217

215218
assertEq(original, optimized, "Implementations must match");
216219
}
@@ -223,7 +226,7 @@ contract MerkleTreeTest is Test, MerkleTreeExample {
223226
/// @notice Fuzz test - Optimized implementation only (to measure gas)
224227
function testFuzz_gas_optimized(uint256 leavesCount) public pure {
225228
leavesCount = bound(leavesCount, 0, 1000000);
226-
MerkleTree.computeMinimalTreeDepth(leavesCount);
229+
MerkleTreeReference.computeMinimalTreeDepth(leavesCount);
227230
}
228231

229232
/// @notice Hashes two `bytes32` values.

0 commit comments

Comments
 (0)