Skip to content

Commit 7be68b0

Browse files
authored
Extend trees by adding processProof function (#165)
* Extend tress by adding processProof function * Adjustments + tests * version bump * version bump * update changelog * patch * revert back version * Adjustments * Adjustments
1 parent 85756c4 commit 7be68b0

File tree

10 files changed

+172
-25
lines changed

10 files changed

+172
-25
lines changed

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
# Changelog
22

3+
## [patch]
4+
5+
- Slightly extended Merkle Trees functionality. Added `processProof` function.
6+
37
## [3.2.2]
48

59
- Migrated to Hardhat 3.

contracts/libs/data-structures/CartesianMerkleTree.sol

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,21 @@ library CartesianMerkleTree {
338338
return _verifyProof(treaple._treaple, proof_);
339339
}
340340

341+
/**
342+
* @notice The function to process the proof for inclusion or exclusion of a node in the CMT.
343+
* Complexity is O(log(n)), where n is the max depth of the treaple.
344+
*
345+
* @param hash3_ The hash function that accepts three arguments.
346+
* @param proof_ The CMT proof struct.
347+
* @return The calculated root hash from the proof.
348+
*/
349+
function processProof(
350+
function(bytes32, bytes32, bytes32) view returns (bytes32) hash3_,
351+
Proof memory proof_
352+
) internal view returns (bytes32) {
353+
return _processProof(hash3_, proof_);
354+
}
355+
341356
/**
342357
* @notice The function to get the root of the Cartesian Merkle Tree.
343358
* Complexity is O(1).
@@ -924,6 +939,18 @@ library CartesianMerkleTree {
924939
return false;
925940
}
926941

942+
function(bytes32, bytes32, bytes32) view returns (bytes32) hash3_ = treaple
943+
.isCustomHasherSet
944+
? treaple.hash3
945+
: _hash3;
946+
947+
return _processProof(hash3_, proof_) == proof_.root;
948+
}
949+
950+
function _processProof(
951+
function(bytes32, bytes32, bytes32) view returns (bytes32) hash3_,
952+
Proof memory proof_
953+
) private view returns (bytes32) {
927954
bool directionBit_ = _extractDirectionBit(proof_.directionBits, 0);
928955

929956
bytes32 leftHash_ = proof_.siblings[proof_.siblingsLength - 2];
@@ -933,8 +960,7 @@ library CartesianMerkleTree {
933960
(leftHash_, rightHash_) = (rightHash_, leftHash_);
934961
}
935962

936-
bytes32 computedHash_ = _getNodesHash(
937-
treaple,
963+
bytes32 computedHash_ = hash3_(
938964
proof_.existence ? proof_.key : proof_.nonExistenceKey,
939965
leftHash_,
940966
rightHash_
@@ -950,15 +976,14 @@ library CartesianMerkleTree {
950976
(leftHash_, rightHash_) = (rightHash_, leftHash_);
951977
}
952978

953-
computedHash_ = _getNodesHash(
954-
treaple,
979+
computedHash_ = hash3_(
955980
proof_.siblings[proof_.siblingsLength - i - 2],
956981
leftHash_,
957982
rightHash_
958983
);
959984
}
960985

961-
return computedHash_ == proof_.root;
986+
return computedHash_;
962987
}
963988

964989
function _newNode(CMT storage treaple, bytes32 key_) private returns (uint256) {

contracts/libs/data-structures/IncrementalMerkleTree.sol

Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -118,15 +118,17 @@ library IncrementalMerkleTree {
118118
* @param siblings_ The siblings of the leaf.
119119
* @param directionBits_ The direction bits of the leaf.
120120
* @param leaf_ The leaf.
121+
* @param root_ The root hash of the tree to verify against.
121122
* @return True if the proof is valid, false otherwise.
122123
*/
123124
function verifyProof(
124125
UintIMT storage tree,
125126
bytes32[] memory siblings_,
126127
uint256 directionBits_,
127-
bytes32 leaf_
128+
bytes32 leaf_,
129+
bytes32 root_
128130
) internal view returns (bool) {
129-
return _verifyProof(tree._tree, siblings_, directionBits_, leaf_);
131+
return _verifyProof(tree._tree, siblings_, directionBits_, leaf_, root_);
130132
}
131133

132134
/**
@@ -217,15 +219,36 @@ library IncrementalMerkleTree {
217219
* @param siblings_ The siblings of the leaf.
218220
* @param directionBits_ The direction bits of the leaf.
219221
* @param leaf_ The leaf.
222+
* @param root_ The root hash of the tree to verify against.
220223
* @return True if the proof is valid, false otherwise.
221224
*/
222225
function verifyProof(
223226
Bytes32IMT storage tree,
224227
bytes32[] memory siblings_,
225228
uint256 directionBits_,
226-
bytes32 leaf_
229+
bytes32 leaf_,
230+
bytes32 root_
227231
) internal view returns (bool) {
228-
return _verifyProof(tree._tree, siblings_, directionBits_, leaf_);
232+
return _verifyProof(tree._tree, siblings_, directionBits_, leaf_, root_);
233+
}
234+
235+
/**
236+
* @notice The function to process the proof for inclusion or exclusion of a leaf in the tree.
237+
* Complexity is O(log(n)), where n is the number of elements in the tree.
238+
*
239+
* @param hash2_ The hash function that accepts two arguments.
240+
* @param siblings_ The siblings of the leaf.
241+
* @param directionBits_ The direction bits of the leaf.
242+
* @param leaf_ The leaf.
243+
* @return The calculated root hash from the proof.
244+
*/
245+
function processProof(
246+
function(bytes32, bytes32) view returns (bytes32) hash2_,
247+
bytes32[] memory siblings_,
248+
uint256 directionBits_,
249+
bytes32 leaf_
250+
) internal view returns (bytes32) {
251+
return _processProof(hash2_, siblings_, directionBits_, leaf_);
229252
}
230253

231254
/**
@@ -312,15 +335,17 @@ library IncrementalMerkleTree {
312335
* @param siblings_ The siblings of the leaf.
313336
* @param directionBits_ The direction bits of the leaf.
314337
* @param leaf_ The leaf.
338+
* @param root_ The root hash of the tree to verify against.
315339
* @return True if the proof is valid, false otherwise.
316340
*/
317341
function verifyProof(
318342
AddressIMT storage tree,
319343
bytes32[] memory siblings_,
320344
uint256 directionBits_,
321-
bytes32 leaf_
345+
bytes32 leaf_,
346+
bytes32 root_
322347
) internal view returns (bool) {
323-
return _verifyProof(tree._tree, siblings_, directionBits_, leaf_);
348+
return _verifyProof(tree._tree, siblings_, directionBits_, leaf_, root_);
324349
}
325350

326351
/**
@@ -462,12 +487,22 @@ library IncrementalMerkleTree {
462487
IMT storage tree,
463488
bytes32[] memory siblings_,
464489
uint256 directionBits,
465-
bytes32 leaf_
490+
bytes32 leaf_,
491+
bytes32 root_
466492
) private view returns (bool) {
467493
function(bytes32, bytes32) view returns (bytes32) hash2_ = tree.isCustomHasherSet
468494
? tree.hash2
469495
: _hash2;
470496

497+
return _processProof(hash2_, siblings_, directionBits, leaf_) == root_;
498+
}
499+
500+
function _processProof(
501+
function(bytes32, bytes32) view returns (bytes32) hash2_,
502+
bytes32[] memory siblings_,
503+
uint256 directionBits,
504+
bytes32 leaf_
505+
) private view returns (bytes32) {
471506
bytes32 computedHash_ = leaf_;
472507

473508
for (uint256 i = 0; i < siblings_.length; ++i) {
@@ -478,7 +513,7 @@ library IncrementalMerkleTree {
478513
}
479514
}
480515

481-
return computedHash_ == _root(tree);
516+
return computedHash_;
482517
}
483518

484519
function _height(IMT storage tree) private view returns (uint256) {

contracts/libs/data-structures/SparseMerkleTree.sol

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,23 @@ library SparseMerkleTree {
376376
return _verifyProof(tree._tree, proof_);
377377
}
378378

379+
/**
380+
* @notice The function to process the proof for inclusion or exclusion of a node in the SMT.
381+
* Complexity is O(log(n)), where n is the max depth of the tree.
382+
*
383+
* @param hash2_ The hash function that accepts two arguments.
384+
* @param hash3_ The hash function that accepts three arguments.
385+
* @param proof_ The SMT proof struct.
386+
* @return The calculated root hash from the proof.
387+
*/
388+
function processProof(
389+
function(bytes32, bytes32) view returns (bytes32) hash2_,
390+
function(bytes32, bytes32, bytes32) view returns (bytes32) hash3_,
391+
Proof memory proof_
392+
) internal view returns (bytes32) {
393+
return _processProof(hash2_, hash3_, proof_);
394+
}
395+
379396
/**
380397
* @notice The function to get the root of the Merkle tree.
381398
* Complexity is O(1).
@@ -1112,6 +1129,14 @@ library SparseMerkleTree {
11121129
function(bytes32, bytes32, bytes32) view returns (bytes32) hash3_
11131130
) = _getHashFunctions(tree);
11141131

1132+
return _processProof(hash2_, hash3_, proof_) == proof_.root;
1133+
}
1134+
1135+
function _processProof(
1136+
function(bytes32, bytes32) view returns (bytes32) hash2_,
1137+
function(bytes32, bytes32, bytes32) view returns (bytes32) hash3_,
1138+
Proof memory proof_
1139+
) private view returns (bytes32) {
11151140
bytes32 computedHash_;
11161141

11171142
if (proof_.existence) {
@@ -1139,7 +1164,7 @@ library SparseMerkleTree {
11391164
}
11401165
}
11411166

1142-
return computedHash_ == proof_.root;
1167+
return computedHash_;
11431168
}
11441169

11451170
function _getHashFunctions(

contracts/mock/libs/data-structures/CartesianMerkleTreeMock.sol

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,12 @@ contract CartesianMerkleTreeMock {
114114
return _addressCMT.verifyProof(proof_);
115115
}
116116

117+
function processCMTProof(
118+
CartesianMerkleTree.Proof memory proof_
119+
) external view returns (bytes32) {
120+
return CartesianMerkleTree.processProof(_hash3, proof_);
121+
}
122+
117123
function getUintRoot() external view returns (bytes32) {
118124
return _uintCMT.getRoot();
119125
}

contracts/mock/libs/data-structures/IncrementalMerkleTreeMock.sol

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -70,25 +70,36 @@ contract IncrementalMerkleTreeMock {
7070
function verifyUintProof(
7171
bytes32[] memory siblings_,
7272
uint256 directionBits_,
73-
bytes32 element_
73+
bytes32 element_,
74+
bytes32 root_
7475
) external view returns (bool) {
75-
return _uintTree.verifyProof(siblings_, directionBits_, element_);
76+
return _uintTree.verifyProof(siblings_, directionBits_, element_, root_);
7677
}
7778

7879
function verifyBytes32Proof(
7980
bytes32[] memory siblings_,
8081
uint256 directionBits_,
81-
bytes32 element_
82+
bytes32 element_,
83+
bytes32 root_
8284
) external view returns (bool) {
83-
return _bytes32Tree.verifyProof(siblings_, directionBits_, element_);
85+
return _bytes32Tree.verifyProof(siblings_, directionBits_, element_, root_);
8486
}
8587

8688
function verifyAddressProof(
8789
bytes32[] memory siblings_,
8890
uint256 directionBits_,
89-
bytes32 element_
91+
bytes32 element_,
92+
bytes32 root_
9093
) external view returns (bool) {
91-
return _addressTree.verifyProof(siblings_, directionBits_, element_);
94+
return _addressTree.verifyProof(siblings_, directionBits_, element_, root_);
95+
}
96+
97+
function processIMTProof(
98+
bytes32[] memory siblings_,
99+
uint256 directionBits_,
100+
bytes32 leaf_
101+
) external view returns (bytes32) {
102+
return IncrementalMerkleTree.processProof(_keccak256, siblings_, directionBits_, leaf_);
92103
}
93104

94105
function getUintTreeHeight() external view returns (uint256) {
@@ -134,4 +145,16 @@ contract IncrementalMerkleTreeMock {
134145
function _hash2(bytes32 element1_, bytes32 element2_) internal pure returns (bytes32) {
135146
return bytes32(PoseidonUnit2L.poseidon([uint256(element1_), uint256(element2_)]));
136147
}
148+
149+
function _keccak256(
150+
bytes32 element1_,
151+
bytes32 element2_
152+
) internal pure returns (bytes32 result_) {
153+
assembly {
154+
mstore(0, element1_)
155+
mstore(32, element2_)
156+
157+
result_ := keccak256(0, 64)
158+
}
159+
}
137160
}

contracts/mock/libs/data-structures/SparseMerkleTreeMock.sol

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,10 @@ contract SparseMerkleTreeMock {
119119
return _addressTree.verifyProof(proof_);
120120
}
121121

122+
function processProof(SparseMerkleTree.Proof memory proof_) external view returns (bytes32) {
123+
return SparseMerkleTree.processProof(_hash2, _hash3, proof_);
124+
}
125+
122126
function getUintRoot() external view returns (bytes32) {
123127
return _uintTree.getRoot();
124128
}

test/libs/data-structures/CartesianMerkleTree.test.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,7 +307,9 @@ describe("CartesianMerkleTree", () => {
307307
key: proof.key,
308308
nonExistenceKey: proof.nonExistenceKey,
309309
};
310+
310311
expect(await treaple.verifyUintProof(proofObj)).to.be.true;
312+
expect(await treaple.processCMTProof(proofObj)).to.be.eq(proof.root);
311313
}
312314
});
313315

@@ -553,7 +555,9 @@ describe("CartesianMerkleTree", () => {
553555
key: proof.key,
554556
nonExistenceKey: proof.nonExistenceKey,
555557
};
558+
556559
expect(await treaple.verifyBytes32Proof(proofObj)).to.be.true;
560+
expect(await treaple.processCMTProof(proofObj)).to.be.eq(proof.root);
557561
}
558562
});
559563
});
@@ -684,7 +688,9 @@ describe("CartesianMerkleTree", () => {
684688
key: proof.key,
685689
nonExistenceKey: proof.nonExistenceKey,
686690
};
691+
687692
expect(await treaple.verifyAddressProof(proofObj)).to.be.true;
693+
expect(await treaple.processCMTProof(proofObj)).to.be.eq(proof.root);
688694
}
689695
});
690696
});

0 commit comments

Comments
 (0)