Skip to content

Commit e87e029

Browse files
author
Artjom Galaktionov
committed
Merge branch 'master' into feat/Vesting
2 parents bd6cff2 + 9d60288 commit e87e029

File tree

9 files changed

+867
-77
lines changed

9 files changed

+867
-77
lines changed

‎contracts/libs/data-structures/IncrementalMerkleTree.sol‎

Lines changed: 194 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,14 @@ pragma solidity ^0.8.4;
1919
*
2020
* | Statistic | _add | _root |
2121
* | --------- | ------------ | ---------------- |
22-
* | count | 106000 | 106000 |
23-
* | mean | 36619 gas | 71941 gas |
24-
* | std | 3617 gas | 4324 gas |
25-
* | min | 34053 gas | 28670 gas |
26-
* | 25% | 34077 gas | 69715 gas |
27-
* | 50% | 36598 gas | 72641 gas |
28-
* | 75% | 39143 gas | 75557 gas |
29-
* | max | 94661 gas | 75637 gas |
22+
* | count | 49999 | 49999 |
23+
* | mean | 38972 gas | 60213 gas |
24+
* | std | 3871 gas | 4996 gas |
25+
* | min | 36251 gas | 31238 gas |
26+
* | 25% | 36263 gas | 57020 gas |
27+
* | 50% | 38954 gas | 60292 gas |
28+
* | 75% | 41657 gas | 63564 gas |
29+
* | max | 96758 gas | 78071 gas |
3030
*
3131
* ## Usage example:
3232
*
@@ -57,6 +57,17 @@ library IncrementalMerkleTree {
5757
IMT _tree;
5858
}
5959

60+
/**
61+
* @notice The function to set the height of the uint256 tree.
62+
* Complexity is O(1).
63+
*
64+
* @param tree self.
65+
* @param height_ The new height of the Merkle tree. Should be greater than the current one.
66+
*/
67+
function setHeight(UintIMT storage tree, uint256 height_) internal {
68+
_setHeight(tree._tree, height_);
69+
}
70+
6071
/**
6172
* @notice The function to add a new element to the uint256 tree.
6273
* Complexity is O(log(n)), where n is the number of elements in the tree.
@@ -68,6 +79,21 @@ library IncrementalMerkleTree {
6879
_add(tree._tree, bytes32(element_));
6980
}
7081

82+
/**
83+
* @notice The function to set a custom hash functions, that will be used to build the Merkle Tree.
84+
*
85+
* @param tree self.
86+
* @param hash1_ The hash function that accepts one argument.
87+
* @param hash2_ The hash function that accepts two arguments.
88+
*/
89+
function setHashers(
90+
UintIMT storage tree,
91+
function(bytes32) view returns (bytes32) hash1_,
92+
function(bytes32, bytes32) view returns (bytes32) hash2_
93+
) internal {
94+
_setHashers(tree._tree, hash1_, hash2_);
95+
}
96+
7197
/**
7298
* @notice The function to return the root hash of the uint256 tree.
7399
* Complexity is O(log(n) + h), where n is the number of elements in the tree and
@@ -98,6 +124,15 @@ library IncrementalMerkleTree {
98124
return _length(tree._tree);
99125
}
100126

127+
/**
128+
* @notice The function to check whether the custom hash functions are set.
129+
* @param tree self.
130+
* @return True if the custom hash functions are set, false otherwise.
131+
*/
132+
function isCustomHasherSet(UintIMT storage tree) internal view returns (bool) {
133+
return tree._tree.isCustomHasherSet;
134+
}
135+
101136
/**
102137
**********************
103138
* Bytes32IMT *
@@ -108,6 +143,17 @@ library IncrementalMerkleTree {
108143
IMT _tree;
109144
}
110145

146+
/**
147+
* @notice The function to set the height of the bytes32 tree.
148+
* Complexity is O(1).
149+
*
150+
* @param tree self.
151+
* @param height_ The new height of the Merkle tree. Should be greater than the current one.
152+
*/
153+
function setHeight(Bytes32IMT storage tree, uint256 height_) internal {
154+
_setHeight(tree._tree, height_);
155+
}
156+
111157
/**
112158
* @notice The function to add a new element to the bytes32 tree.
113159
* Complexity is O(log(n)), where n is the number of elements in the tree.
@@ -116,6 +162,21 @@ library IncrementalMerkleTree {
116162
_add(tree._tree, element_);
117163
}
118164

165+
/**
166+
* @notice The function to set a custom hash functions, that will be used to build the Merkle Tree.
167+
*
168+
* @param tree self.
169+
* @param hash1_ The hash function that accepts one argument.
170+
* @param hash2_ The hash function that accepts two arguments.
171+
*/
172+
function setHashers(
173+
Bytes32IMT storage tree,
174+
function(bytes32) view returns (bytes32) hash1_,
175+
function(bytes32, bytes32) view returns (bytes32) hash2_
176+
) internal {
177+
_setHashers(tree._tree, hash1_, hash2_);
178+
}
179+
119180
/**
120181
* @notice The function to return the root hash of the bytes32 tree.
121182
* Complexity is O(log(n) + h), where n is the number of elements in the tree and
@@ -139,6 +200,15 @@ library IncrementalMerkleTree {
139200
return _length(tree._tree);
140201
}
141202

203+
/**
204+
* @notice The function to check whether the custom hash functions are set.
205+
* @param tree self.
206+
* @return True if the custom hash functions are set, false otherwise.
207+
*/
208+
function isCustomHasherSet(Bytes32IMT storage tree) internal view returns (bool) {
209+
return tree._tree.isCustomHasherSet;
210+
}
211+
142212
/**
143213
************************
144214
* AddressIMT *
@@ -149,6 +219,17 @@ library IncrementalMerkleTree {
149219
IMT _tree;
150220
}
151221

222+
/**
223+
* @notice The function to set the height of the address tree.
224+
* Complexity is O(1).
225+
*
226+
* @param tree self.
227+
* @param height_ The new height of the Merkle tree. Should be greater than the current one.
228+
*/
229+
function setHeight(AddressIMT storage tree, uint256 height_) internal {
230+
_setHeight(tree._tree, height_);
231+
}
232+
152233
/**
153234
* @notice The function to add a new element to the address tree.
154235
* Complexity is O(log(n)), where n is the number of elements in the tree.
@@ -157,6 +238,21 @@ library IncrementalMerkleTree {
157238
_add(tree._tree, bytes32(uint256(uint160(element_))));
158239
}
159240

241+
/**
242+
* @notice The function to set a custom hash functions, that will be used to build the Merkle Tree.
243+
*
244+
* @param tree self.
245+
* @param hash1_ The hash function that accepts one argument.
246+
* @param hash2_ The hash function that accepts two arguments.
247+
*/
248+
function setHashers(
249+
AddressIMT storage tree,
250+
function(bytes32) view returns (bytes32) hash1_,
251+
function(bytes32, bytes32) view returns (bytes32) hash2_
252+
) internal {
253+
_setHashers(tree._tree, hash1_, hash2_);
254+
}
255+
160256
/**
161257
* @notice The function to return the root hash of the address tree.
162258
* Complexity is O(log(n) + h), where n is the number of elements in the tree and
@@ -180,6 +276,15 @@ library IncrementalMerkleTree {
180276
return _length(tree._tree);
181277
}
182278

279+
/**
280+
* @notice The function to check whether the custom hash functions are set.
281+
* @param tree self.
282+
* @return True if the custom hash functions are set, false otherwise.
283+
*/
284+
function isCustomHasherSet(AddressIMT storage tree) internal view returns (bool) {
285+
return tree._tree.isCustomHasherSet;
286+
}
287+
183288
/**
184289
************************
185290
* InnerIMT *
@@ -189,17 +294,47 @@ library IncrementalMerkleTree {
189294
struct IMT {
190295
bytes32[] branches;
191296
uint256 leavesCount;
297+
bool isStrictHeightSet;
298+
bool isCustomHasherSet;
299+
function(bytes32) view returns (bytes32) hash1;
300+
function(bytes32, bytes32) view returns (bytes32) hash2;
192301
}
193302

194-
bytes32 private constant ZERO_HASH = keccak256(abi.encode(0));
303+
function _setHeight(IMT storage tree, uint256 height_) private {
304+
require(
305+
height_ > _height(tree),
306+
"IncrementalMerkleTree: the height must be greater than the current one"
307+
);
195308

196-
function _add(IMT storage tree, bytes32 element_) private {
197-
bytes32 resultValue_;
309+
tree.isStrictHeightSet = true;
198310

199311
assembly {
200-
mstore(0, element_)
201-
resultValue_ := keccak256(0, 32)
312+
sstore(tree.slot, height_)
202313
}
314+
}
315+
316+
function _setHashers(
317+
IMT storage tree,
318+
function(bytes32) view returns (bytes32) hash1_,
319+
function(bytes32, bytes32) view returns (bytes32) hash2_
320+
) private {
321+
require(_length(tree) == 0, "IncrementalMerkleTree: the tree must be empty");
322+
323+
tree.isCustomHasherSet = true;
324+
325+
tree.hash1 = hash1_;
326+
tree.hash2 = hash2_;
327+
}
328+
329+
function _add(IMT storage tree, bytes32 element_) private {
330+
function(bytes32) view returns (bytes32) hash1_ = tree.isCustomHasherSet
331+
? tree.hash1
332+
: _hash1;
333+
function(bytes32, bytes32) view returns (bytes32) hash2_ = tree.isCustomHasherSet
334+
? tree.hash2
335+
: _hash2;
336+
337+
bytes32 resultValue_ = hash1_(element_);
203338

204339
uint256 index_ = 0;
205340
uint256 size_ = ++tree.leavesCount;
@@ -211,56 +346,51 @@ library IncrementalMerkleTree {
211346
}
212347

213348
bytes32 branch_ = tree.branches[index_];
214-
215-
assembly {
216-
mstore(0, branch_)
217-
mstore(32, resultValue_)
218-
219-
resultValue_ := keccak256(0, 64)
220-
}
349+
resultValue_ = hash2_(branch_, resultValue_);
221350

222351
size_ >>= 1;
223352
++index_;
224353
}
225354

226355
if (index_ == treeHeight_) {
356+
if (tree.isStrictHeightSet) {
357+
revert("IncrementalMerkleTree: the tree is full");
358+
}
359+
227360
tree.branches.push(resultValue_);
228361
} else {
229362
tree.branches[index_] = resultValue_;
230363
}
231364
}
232365

233366
function _root(IMT storage tree) private view returns (bytes32) {
367+
function(bytes32) view returns (bytes32) hash1_ = tree.isCustomHasherSet
368+
? tree.hash1
369+
: _hash1;
370+
function(bytes32, bytes32) view returns (bytes32) hash2_ = tree.isCustomHasherSet
371+
? tree.hash2
372+
: _hash2;
373+
234374
uint256 treeHeight_ = tree.branches.length;
235375

236376
if (treeHeight_ == 0) {
237-
return ZERO_HASH;
377+
return hash1_(bytes32(0));
238378
}
239379

240380
uint256 height_;
241381
uint256 size_ = tree.leavesCount;
242-
bytes32 root_ = ZERO_HASH;
243-
bytes32[] memory zeroHashes_ = _getZeroHashes(treeHeight_);
382+
bytes32 root_ = hash1_(bytes32(0));
383+
bytes32[] memory zeroHashes_ = _getZeroHashes(tree, treeHeight_);
244384

245385
while (height_ < treeHeight_) {
246386
if (size_ & 1 == 1) {
247387
bytes32 branch_ = tree.branches[height_];
248388

249-
assembly {
250-
mstore(0, branch_)
251-
mstore(32, root_)
252-
253-
root_ := keccak256(0, 64)
254-
}
389+
root_ = hash2_(branch_, root_);
255390
} else {
256391
bytes32 zeroHash_ = zeroHashes_[height_];
257392

258-
assembly {
259-
mstore(0, root_)
260-
mstore(32, zeroHash_)
261-
262-
root_ := keccak256(0, 64)
263-
}
393+
root_ = hash2_(root_, zeroHash_);
264394
}
265395

266396
size_ >>= 1;
@@ -278,25 +408,44 @@ library IncrementalMerkleTree {
278408
return tree.leavesCount;
279409
}
280410

281-
function _getZeroHashes(uint256 height_) private pure returns (bytes32[] memory) {
411+
function _getZeroHashes(
412+
IMT storage tree,
413+
uint256 height_
414+
) private view returns (bytes32[] memory) {
415+
function(bytes32) view returns (bytes32) hash1_ = tree.isCustomHasherSet
416+
? tree.hash1
417+
: _hash1;
418+
function(bytes32, bytes32) view returns (bytes32) hash2_ = tree.isCustomHasherSet
419+
? tree.hash2
420+
: _hash2;
421+
282422
bytes32[] memory zeroHashes_ = new bytes32[](height_);
283423

284-
zeroHashes_[0] = ZERO_HASH;
424+
zeroHashes_[0] = hash1_(bytes32(0));
285425

286426
for (uint256 i = 1; i < height_; ++i) {
287-
bytes32 result;
288427
bytes32 prevHash_ = zeroHashes_[i - 1];
289428

290-
assembly {
291-
mstore(0, prevHash_)
292-
mstore(32, prevHash_)
429+
zeroHashes_[i] = hash2_(prevHash_, prevHash_);
430+
}
431+
432+
return zeroHashes_;
433+
}
293434

294-
result := keccak256(0, 64)
295-
}
435+
function _hash1(bytes32 a) private pure returns (bytes32 result) {
436+
assembly {
437+
mstore(0, a)
296438

297-
zeroHashes_[i] = result;
439+
result := keccak256(0, 32)
298440
}
441+
}
299442

300-
return zeroHashes_;
443+
function _hash2(bytes32 a, bytes32 b) private pure returns (bytes32 result) {
444+
assembly {
445+
mstore(0, a)
446+
mstore(32, b)
447+
448+
result := keccak256(0, 64)
449+
}
301450
}
302451
}

0 commit comments

Comments
 (0)