diff --git a/packages/imt.sol/contracts/BinaryIMT.sol b/packages/imt.sol/contracts/BinaryIMT.sol index 06c2a6a1f..23c4e2068 100644 --- a/packages/imt.sol/contracts/BinaryIMT.sol +++ b/packages/imt.sol/contracts/BinaryIMT.sol @@ -1,188 +1,27 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.4; -import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; +import {InternalBinaryIMT, BinaryIMTData} from "./internal/InternalBinaryIMT.sol"; -// Each incremental tree has certain properties and data that will -// be used to add new leaves. -struct BinaryIMTData { - uint256 depth; // Depth of the tree (levels - 1). - uint256 root; // Root hash of the tree. - uint256 numberOfLeaves; // Number of leaves of the tree. - mapping(uint256 => uint256) zeroes; // Zero hashes used for empty nodes (level -> zero hash). - // The nodes of the subtrees used in the last addition of a leaf (level -> [left node, right node]). - mapping(uint256 => uint256[2]) lastSubtrees; // Caching these values is essential to efficient appends. - bool useDefaultZeroes; -} - -error ValueGreaterThanSnarkScalarField(); -error DepthNotSupported(); -error WrongDefaultZeroIndex(); -error TreeIsFull(); -error NewLeafCannotEqualOldLeaf(); -error LeafDoesNotExist(); -error LeafIndexOutOfRange(); -error WrongMerkleProofPath(); - -/// @title Incremental binary Merkle tree. -/// @dev The incremental tree allows to calculate the root hash each time a leaf is added, ensuring -/// the integrity of the tree. library BinaryIMT { - uint8 public constant MAX_DEPTH = 32; - uint256 public constant SNARK_SCALAR_FIELD = - 21888242871839275222246405745257275088548364400416034343698204186575808495617; - - uint256 public constant Z_0 = 0; - uint256 public constant Z_1 = 14744269619966411208579211824598458697587494354926760081771325075741142829156; - uint256 public constant Z_2 = 7423237065226347324353380772367382631490014989348495481811164164159255474657; - uint256 public constant Z_3 = 11286972368698509976183087595462810875513684078608517520839298933882497716792; - uint256 public constant Z_4 = 3607627140608796879659380071776844901612302623152076817094415224584923813162; - uint256 public constant Z_5 = 19712377064642672829441595136074946683621277828620209496774504837737984048981; - uint256 public constant Z_6 = 20775607673010627194014556968476266066927294572720319469184847051418138353016; - uint256 public constant Z_7 = 3396914609616007258851405644437304192397291162432396347162513310381425243293; - uint256 public constant Z_8 = 21551820661461729022865262380882070649935529853313286572328683688269863701601; - uint256 public constant Z_9 = 6573136701248752079028194407151022595060682063033565181951145966236778420039; - uint256 public constant Z_10 = 12413880268183407374852357075976609371175688755676981206018884971008854919922; - uint256 public constant Z_11 = 14271763308400718165336499097156975241954733520325982997864342600795471836726; - uint256 public constant Z_12 = 20066985985293572387227381049700832219069292839614107140851619262827735677018; - uint256 public constant Z_13 = 9394776414966240069580838672673694685292165040808226440647796406499139370960; - uint256 public constant Z_14 = 11331146992410411304059858900317123658895005918277453009197229807340014528524; - uint256 public constant Z_15 = 15819538789928229930262697811477882737253464456578333862691129291651619515538; - uint256 public constant Z_16 = 19217088683336594659449020493828377907203207941212636669271704950158751593251; - uint256 public constant Z_17 = 21035245323335827719745544373081896983162834604456827698288649288827293579666; - uint256 public constant Z_18 = 6939770416153240137322503476966641397417391950902474480970945462551409848591; - uint256 public constant Z_19 = 10941962436777715901943463195175331263348098796018438960955633645115732864202; - uint256 public constant Z_20 = 15019797232609675441998260052101280400536945603062888308240081994073687793470; - uint256 public constant Z_21 = 11702828337982203149177882813338547876343922920234831094975924378932809409969; - uint256 public constant Z_22 = 11217067736778784455593535811108456786943573747466706329920902520905755780395; - uint256 public constant Z_23 = 16072238744996205792852194127671441602062027943016727953216607508365787157389; - uint256 public constant Z_24 = 17681057402012993898104192736393849603097507831571622013521167331642182653248; - uint256 public constant Z_25 = 21694045479371014653083846597424257852691458318143380497809004364947786214945; - uint256 public constant Z_26 = 8163447297445169709687354538480474434591144168767135863541048304198280615192; - uint256 public constant Z_27 = 14081762237856300239452543304351251708585712948734528663957353575674639038357; - uint256 public constant Z_28 = 16619959921569409661790279042024627172199214148318086837362003702249041851090; - uint256 public constant Z_29 = 7022159125197495734384997711896547675021391130223237843255817587255104160365; - uint256 public constant Z_30 = 4114686047564160449611603615418567457008101555090703535405891656262658644463; - uint256 public constant Z_31 = 12549363297364877722388257367377629555213421373705596078299904496781819142130; - uint256 public constant Z_32 = 21443572485391568159800782191812935835534334817699172242223315142338162256601; + using InternalBinaryIMT for *; function defaultZero(uint256 index) public pure returns (uint256) { - if (index == 0) return Z_0; - if (index == 1) return Z_1; - if (index == 2) return Z_2; - if (index == 3) return Z_3; - if (index == 4) return Z_4; - if (index == 5) return Z_5; - if (index == 6) return Z_6; - if (index == 7) return Z_7; - if (index == 8) return Z_8; - if (index == 9) return Z_9; - if (index == 10) return Z_10; - if (index == 11) return Z_11; - if (index == 12) return Z_12; - if (index == 13) return Z_13; - if (index == 14) return Z_14; - if (index == 15) return Z_15; - if (index == 16) return Z_16; - if (index == 17) return Z_17; - if (index == 18) return Z_18; - if (index == 19) return Z_19; - if (index == 20) return Z_20; - if (index == 21) return Z_21; - if (index == 22) return Z_22; - if (index == 23) return Z_23; - if (index == 24) return Z_24; - if (index == 25) return Z_25; - if (index == 26) return Z_26; - if (index == 27) return Z_27; - if (index == 28) return Z_28; - if (index == 29) return Z_29; - if (index == 30) return Z_30; - if (index == 31) return Z_31; - if (index == 32) return Z_32; - - revert WrongDefaultZeroIndex(); + return InternalBinaryIMT._defaultZero(index); } - /// @dev Initializes a tree. - /// @param self: Tree data. - /// @param depth: Depth of the tree. - /// @param zero: Zero value to be used. function init(BinaryIMTData storage self, uint256 depth, uint256 zero) public { - if (zero >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (depth <= 0 || depth > MAX_DEPTH) { - revert DepthNotSupported(); - } - - self.depth = depth; - - for (uint8 i = 0; i < depth; ) { - self.zeroes[i] = zero; - zero = PoseidonT3.hash([zero, zero]); - - unchecked { - ++i; - } - } - - self.root = zero; + InternalBinaryIMT._init(self, depth, zero); } function initWithDefaultZeroes(BinaryIMTData storage self, uint256 depth) public { - if (depth <= 0 || depth > MAX_DEPTH) { - revert DepthNotSupported(); - } - - self.depth = depth; - self.useDefaultZeroes = true; - - self.root = defaultZero(depth); + InternalBinaryIMT._initWithDefaultZeroes(self, depth); } - /// @dev Inserts a leaf in the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be inserted. function insert(BinaryIMTData storage self, uint256 leaf) public returns (uint256) { - uint256 depth = self.depth; - - if (leaf >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (self.numberOfLeaves >= 2 ** depth) { - revert TreeIsFull(); - } - - uint256 index = self.numberOfLeaves; - uint256 hash = leaf; - bool useDefaultZeroes = self.useDefaultZeroes; - - for (uint8 i = 0; i < depth; ) { - if (index & 1 == 0) { - self.lastSubtrees[i] = [hash, useDefaultZeroes ? defaultZero(i) : self.zeroes[i]]; - } else { - self.lastSubtrees[i][1] = hash; - } - - hash = PoseidonT3.hash(self.lastSubtrees[i]); - index >>= 1; - - unchecked { - ++i; - } - } - - self.root = hash; - self.numberOfLeaves += 1; - - return hash; + return InternalBinaryIMT._insert(self, leaf); } - /// @dev Updates a leaf in the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be updated. - /// @param newLeaf: New leaf. - /// @param proofSiblings: Array of the sibling nodes of the proof of membership. - /// @param proofPathIndices: Path of the proof of membership. function update( BinaryIMTData storage self, uint256 leaf, @@ -190,101 +29,24 @@ library BinaryIMT { uint256[] calldata proofSiblings, uint8[] calldata proofPathIndices ) public { - if (newLeaf == leaf) { - revert NewLeafCannotEqualOldLeaf(); - } else if (newLeaf >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (!verify(self, leaf, proofSiblings, proofPathIndices)) { - revert LeafDoesNotExist(); - } - - uint256 depth = self.depth; - uint256 hash = newLeaf; - uint256 updateIndex; - - for (uint8 i = 0; i < depth; ) { - updateIndex |= uint256(proofPathIndices[i]) << uint256(i); - - if (proofPathIndices[i] == 0) { - if (proofSiblings[i] == self.lastSubtrees[i][1]) { - self.lastSubtrees[i][0] = hash; - } - - hash = PoseidonT3.hash([hash, proofSiblings[i]]); - } else { - if (proofSiblings[i] == self.lastSubtrees[i][0]) { - self.lastSubtrees[i][1] = hash; - } - - hash = PoseidonT3.hash([proofSiblings[i], hash]); - } - - unchecked { - ++i; - } - } - - if (updateIndex >= self.numberOfLeaves) { - revert LeafIndexOutOfRange(); - } - - self.root = hash; + InternalBinaryIMT._update(self, leaf, newLeaf, proofSiblings, proofPathIndices); } - /// @dev Removes a leaf from the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be removed. - /// @param proofSiblings: Array of the sibling nodes of the proof of membership. - /// @param proofPathIndices: Path of the proof of membership. function remove( BinaryIMTData storage self, uint256 leaf, uint256[] calldata proofSiblings, uint8[] calldata proofPathIndices ) public { - update(self, leaf, self.useDefaultZeroes ? Z_0 : self.zeroes[0], proofSiblings, proofPathIndices); + InternalBinaryIMT._remove(self, leaf, proofSiblings, proofPathIndices); } - /// @dev Verify if the path is correct and the leaf is part of the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be removed. - /// @param proofSiblings: Array of the sibling nodes of the proof of membership. - /// @param proofPathIndices: Path of the proof of membership. - /// @return True or false. function verify( BinaryIMTData storage self, uint256 leaf, uint256[] calldata proofSiblings, uint8[] calldata proofPathIndices ) private view returns (bool) { - uint256 depth = self.depth; - - if (leaf >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (proofPathIndices.length != depth || proofSiblings.length != depth) { - revert WrongMerkleProofPath(); - } - - uint256 hash = leaf; - - for (uint8 i = 0; i < depth; ) { - if (proofSiblings[i] >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (proofPathIndices[i] != 1 && proofPathIndices[i] != 0) { - revert WrongMerkleProofPath(); - } - - if (proofPathIndices[i] == 0) { - hash = PoseidonT3.hash([hash, proofSiblings[i]]); - } else { - hash = PoseidonT3.hash([proofSiblings[i], hash]); - } - - unchecked { - ++i; - } - } - - return hash == self.root; + return InternalBinaryIMT._verify(self, leaf, proofSiblings, proofPathIndices); } } diff --git a/packages/imt.sol/contracts/Constants.sol b/packages/imt.sol/contracts/Constants.sol new file mode 100644 index 000000000..9f3c45578 --- /dev/null +++ b/packages/imt.sol/contracts/Constants.sol @@ -0,0 +1,5 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.4; + +uint256 constant SNARK_SCALAR_FIELD = 21888242871839275222246405745257275088548364400416034343698204186575808495617; +uint8 constant MAX_DEPTH = 32; diff --git a/packages/imt.sol/contracts/LazyIMT.sol b/packages/imt.sol/contracts/LazyIMT.sol index 83235c2ea..71d869dbc 100644 --- a/packages/imt.sol/contracts/LazyIMT.sol +++ b/packages/imt.sol/contracts/LazyIMT.sol @@ -2,202 +2,44 @@ pragma solidity ^0.8.4; import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; - -struct LazyIMTData { - uint32 maxIndex; - uint40 numberOfLeaves; - mapping(uint256 => uint256) elements; -} +import {InternalLazyIMT, LazyIMTData} from "./internal/InternalLazyIMT.sol"; library LazyIMT { - uint256 public constant SNARK_SCALAR_FIELD = - 21888242871839275222246405745257275088548364400416034343698204186575808495617; - uint8 public constant MAX_DEPTH = 32; - uint40 public constant MAX_INDEX = (1 << 32) - 1; + using InternalLazyIMT for *; function init(LazyIMTData storage self, uint8 depth) public { - require(depth <= MAX_DEPTH, "LazyIMT: Tree too large"); - self.maxIndex = uint32((1 << depth) - 1); - self.numberOfLeaves = 0; + InternalLazyIMT._init(self, depth); + } + + function defaultZero(uint8 index) public pure returns (uint256) { + return InternalLazyIMT._defaultZero(index); } function reset(LazyIMTData storage self) public { - self.numberOfLeaves = 0; + InternalLazyIMT._reset(self); } function indexForElement(uint8 level, uint40 index) public pure returns (uint40) { - // store the elements sparsely - return MAX_INDEX * level + index; + return InternalLazyIMT._indexForElement(level, index); } function insert(LazyIMTData storage self, uint256 leaf) public { - uint40 index = self.numberOfLeaves; - require(leaf < SNARK_SCALAR_FIELD, "LazyIMT: leaf must be < SNARK_SCALAR_FIELD"); - require(index < self.maxIndex, "LazyIMT: tree is full"); - - self.numberOfLeaves = index + 1; - - uint256 hash = leaf; - - for (uint8 i = 0; ; ) { - self.elements[indexForElement(i, index)] = hash; - // it's a left element so we don't hash until there's a right element - if (index & 1 == 0) break; - uint40 elementIndex = indexForElement(i, index - 1); - hash = PoseidonT3.hash([self.elements[elementIndex], hash]); - unchecked { - index >>= 1; - i++; - } - } + InternalLazyIMT._insert(self, leaf); } function update(LazyIMTData storage self, uint256 leaf, uint40 index) public { - require(leaf < SNARK_SCALAR_FIELD, "LazyIMT: leaf must be < SNARK_SCALAR_FIELD"); - uint40 numberOfLeaves = self.numberOfLeaves; - require(index < numberOfLeaves, "LazyIMT: leaf must exist"); - - uint256 hash = leaf; - - for (uint8 i = 0; true; ) { - self.elements[indexForElement(i, index)] = hash; - uint256 levelCount = numberOfLeaves >> (i + 1); - if (levelCount <= index >> 1) break; - if (index & 1 == 0) { - uint40 elementIndex = indexForElement(i, index + 1); - hash = PoseidonT3.hash([hash, self.elements[elementIndex]]); - } else { - uint40 elementIndex = indexForElement(i, index - 1); - hash = PoseidonT3.hash([self.elements[elementIndex], hash]); - } - unchecked { - index >>= 1; - i++; - } - } + InternalLazyIMT._update(self, leaf, index); } function root(LazyIMTData storage self) public view returns (uint256) { - // this will always short circuit if self.numberOfLeaves == 0 - uint40 numberOfLeaves = self.numberOfLeaves; - // dynamically determine a depth - uint8 depth = 1; - while (uint8(2) ** depth < numberOfLeaves) { - depth++; - } - return _root(self, numberOfLeaves, depth); + return InternalLazyIMT._root(self); } function root(LazyIMTData storage self, uint8 depth) public view returns (uint256) { - uint40 numberOfLeaves = self.numberOfLeaves; - require(2 ** depth >= numberOfLeaves, "LazyIMT: ambiguous depth"); - return _root(self, self.numberOfLeaves, depth); + return InternalLazyIMT._root(self, depth); } function _root(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth) internal view returns (uint256) { - require(depth > 0, "LazyIMT: depth must be > 0"); - require(depth <= MAX_DEPTH, "LazyIMT: depth must be < MAX_DEPTH"); - // this should always short circuit if self.numberOfLeaves == 0 - if (numberOfLeaves == 0) return defaultZero(depth); - uint40 index = numberOfLeaves - 1; - - uint256[MAX_DEPTH + 1] memory levels; - - if (index & 1 == 0) { - levels[0] = self.elements[indexForElement(0, index)]; - } else { - levels[0] = defaultZero(0); - } - - for (uint8 i = 0; i < depth; ) { - if (index & 1 == 0) { - levels[i + 1] = PoseidonT3.hash([levels[i], defaultZero(i)]); - } else { - uint256 levelCount = (numberOfLeaves) >> (i + 1); - if (levelCount > index >> 1) { - uint256 parent = self.elements[indexForElement(i + 1, index >> 1)]; - levels[i + 1] = parent; - } else { - uint256 sibling = self.elements[indexForElement(i, index - 1)]; - levels[i + 1] = PoseidonT3.hash([sibling, levels[i]]); - } - } - unchecked { - index >>= 1; - i++; - } - } - return levels[depth]; - } - - uint256 public constant Z_0 = 0; - uint256 public constant Z_1 = 14744269619966411208579211824598458697587494354926760081771325075741142829156; - uint256 public constant Z_2 = 7423237065226347324353380772367382631490014989348495481811164164159255474657; - uint256 public constant Z_3 = 11286972368698509976183087595462810875513684078608517520839298933882497716792; - uint256 public constant Z_4 = 3607627140608796879659380071776844901612302623152076817094415224584923813162; - uint256 public constant Z_5 = 19712377064642672829441595136074946683621277828620209496774504837737984048981; - uint256 public constant Z_6 = 20775607673010627194014556968476266066927294572720319469184847051418138353016; - uint256 public constant Z_7 = 3396914609616007258851405644437304192397291162432396347162513310381425243293; - uint256 public constant Z_8 = 21551820661461729022865262380882070649935529853313286572328683688269863701601; - uint256 public constant Z_9 = 6573136701248752079028194407151022595060682063033565181951145966236778420039; - uint256 public constant Z_10 = 12413880268183407374852357075976609371175688755676981206018884971008854919922; - uint256 public constant Z_11 = 14271763308400718165336499097156975241954733520325982997864342600795471836726; - uint256 public constant Z_12 = 20066985985293572387227381049700832219069292839614107140851619262827735677018; - uint256 public constant Z_13 = 9394776414966240069580838672673694685292165040808226440647796406499139370960; - uint256 public constant Z_14 = 11331146992410411304059858900317123658895005918277453009197229807340014528524; - uint256 public constant Z_15 = 15819538789928229930262697811477882737253464456578333862691129291651619515538; - uint256 public constant Z_16 = 19217088683336594659449020493828377907203207941212636669271704950158751593251; - uint256 public constant Z_17 = 21035245323335827719745544373081896983162834604456827698288649288827293579666; - uint256 public constant Z_18 = 6939770416153240137322503476966641397417391950902474480970945462551409848591; - uint256 public constant Z_19 = 10941962436777715901943463195175331263348098796018438960955633645115732864202; - uint256 public constant Z_20 = 15019797232609675441998260052101280400536945603062888308240081994073687793470; - uint256 public constant Z_21 = 11702828337982203149177882813338547876343922920234831094975924378932809409969; - uint256 public constant Z_22 = 11217067736778784455593535811108456786943573747466706329920902520905755780395; - uint256 public constant Z_23 = 16072238744996205792852194127671441602062027943016727953216607508365787157389; - uint256 public constant Z_24 = 17681057402012993898104192736393849603097507831571622013521167331642182653248; - uint256 public constant Z_25 = 21694045479371014653083846597424257852691458318143380497809004364947786214945; - uint256 public constant Z_26 = 8163447297445169709687354538480474434591144168767135863541048304198280615192; - uint256 public constant Z_27 = 14081762237856300239452543304351251708585712948734528663957353575674639038357; - uint256 public constant Z_28 = 16619959921569409661790279042024627172199214148318086837362003702249041851090; - uint256 public constant Z_29 = 7022159125197495734384997711896547675021391130223237843255817587255104160365; - uint256 public constant Z_30 = 4114686047564160449611603615418567457008101555090703535405891656262658644463; - uint256 public constant Z_31 = 12549363297364877722388257367377629555213421373705596078299904496781819142130; - uint256 public constant Z_32 = 21443572485391568159800782191812935835534334817699172242223315142338162256601; - - function defaultZero(uint8 index) public pure returns (uint256) { - if (index == 0) return Z_0; - if (index == 1) return Z_1; - if (index == 2) return Z_2; - if (index == 3) return Z_3; - if (index == 4) return Z_4; - if (index == 5) return Z_5; - if (index == 6) return Z_6; - if (index == 7) return Z_7; - if (index == 8) return Z_8; - if (index == 9) return Z_9; - if (index == 10) return Z_10; - if (index == 11) return Z_11; - if (index == 12) return Z_12; - if (index == 13) return Z_13; - if (index == 14) return Z_14; - if (index == 15) return Z_15; - if (index == 16) return Z_16; - if (index == 17) return Z_17; - if (index == 18) return Z_18; - if (index == 19) return Z_19; - if (index == 20) return Z_20; - if (index == 21) return Z_21; - if (index == 22) return Z_22; - if (index == 23) return Z_23; - if (index == 24) return Z_24; - if (index == 25) return Z_25; - if (index == 26) return Z_26; - if (index == 27) return Z_27; - if (index == 28) return Z_28; - if (index == 29) return Z_29; - if (index == 30) return Z_30; - if (index == 31) return Z_31; - if (index == 32) return Z_32; - revert("LazyIMT: defaultZero bad index"); + return InternalLazyIMT._root(self, numberOfLeaves, depth); } } diff --git a/packages/imt.sol/contracts/LeanIMT.sol b/packages/imt.sol/contracts/LeanIMT.sol index f250d73c0..1382c1151 100644 --- a/packages/imt.sol/contracts/LeanIMT.sol +++ b/packages/imt.sol/contracts/LeanIMT.sol @@ -2,203 +2,41 @@ pragma solidity ^0.8.4; import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; +import {InternalLeanIMT, LeanIMTData} from "./internal/InternalLeanIMT.sol"; -struct LeanIMTData { - // Tracks the current number of leaves in the tree. - uint256 size; - // Represents the current depth of the tree, which can increase as new leaves are inserted. - uint256 depth; - // A mapping from each level of the tree to the rightmost node at that level. - // Used for efficient updates and root calculations. - mapping(uint256 => uint256) rightmostNodes; - // A mapping from leaf values to their respective indices in the tree. - // This facilitates checks for leaf existence and retrieval of leaf positions. - mapping(uint256 => uint256) leaves; -} - -error WrongSiblingNodes(); -error LeafGreaterThanSnarkScalarField(); -error LeafCannotBeZero(); -error LeafAlreadyExists(); -error LeafDoesNotExist(); - -/// @title Lean Incremental binary Merkle tree. -/// @dev The LeanIMT is an optimized version of the BinaryIMT. -/// This implementation eliminates the use of zeroes, and make the tree depth dynamic. -/// When a node doesn't have the right child, instead of using a zero hash as in the BinaryIMT, -/// the node's value becomes that of its left child. Furthermore, rather than utilizing a static tree depth, -/// it is updated based on the number of leaves in the tree. This approach -/// results in the calculation of significantly fewer hashes, making the tree more efficient. library LeanIMT { - uint256 public constant SNARK_SCALAR_FIELD = - 21888242871839275222246405745257275088548364400416034343698204186575808495617; + using InternalLeanIMT for *; - /// @dev Inserts a new leaf into the incremental merkle tree. - /// The function ensures that the leaf is valid according to the - /// constraints of the tree and then updates the tree's structure accordingly. - /// @param self: A storage reference to the 'LeanIMTData' struct. - /// @param leaf: The value of the new leaf to be inserted into the tree. - /// @return The new hash of the node after the leaf has been inserted. function insert(LeanIMTData storage self, uint256 leaf) public returns (uint256) { - if (leaf >= SNARK_SCALAR_FIELD) { - revert LeafGreaterThanSnarkScalarField(); - } else if (leaf == 0) { - revert LeafCannotBeZero(); - } else if (has(self, leaf)) { - revert LeafAlreadyExists(); - } - - while (2 ** self.depth < self.size + 1) { - self.depth += 1; - } - - uint256 index = self.size; - uint256 node = leaf; - - for (uint256 level = 0; level < self.depth; ) { - if ((index >> level) & 1 == 1) { - node = PoseidonT3.hash([self.rightmostNodes[level], node]); - } else { - self.rightmostNodes[level] = node; - } - - unchecked { - ++level; - } - } - - self.size += 1; - - self.rightmostNodes[self.depth] = node; - self.leaves[leaf] = self.size; - - return node; + return InternalLeanIMT._insert(self, leaf); } - /// @dev Updates the value of an existing leaf and recalculates hashes - /// to maintain tree integrity. - /// @param self: A storage reference to the 'LeanIMTData' struct. - /// @param oldLeaf: The value of the leaf that is to be updated. - /// @param newLeaf: The new value that will replace the oldLeaf in the tree. - /// @param siblingNodes: An array of sibling nodes that are necessary to recalculate the path to the root. - /// @return The new hash of the updated node after the leaf has been updated. function update( LeanIMTData storage self, uint256 oldLeaf, uint256 newLeaf, uint256[] calldata siblingNodes ) public returns (uint256) { - if (newLeaf >= SNARK_SCALAR_FIELD) { - revert LeafGreaterThanSnarkScalarField(); - } else if (!has(self, oldLeaf)) { - revert LeafDoesNotExist(); - } else if (newLeaf != 0 && has(self, newLeaf)) { - revert LeafAlreadyExists(); - } - - uint256 index = indexOf(self, oldLeaf); - uint256 node = newLeaf; - uint256 oldRoot = oldLeaf; - - // A counter that adjusts the level at which sibling nodes are - // accessed and updated during the tree's update process. - // It ensures that the update function correctly navigates and - // modifies the tree's nodes at the appropriate levels, accounting - // for situations where not every level of the tree requires an - // update or a hash calculation. - uint256 s = 0; - - // The number of siblings of a proof can be less than - // the depth of the tree, because in some levels it might not - // be necessary to hash any value. - for (uint256 i = 0; i < siblingNodes.length; ) { - if (siblingNodes[i] >= SNARK_SCALAR_FIELD) { - revert LeafGreaterThanSnarkScalarField(); - } - - uint256 level = i + s; - - if (oldRoot == self.rightmostNodes[level]) { - self.rightmostNodes[level] = node; - - if (oldRoot == self.rightmostNodes[level + 1]) { - s += 1; - } - - uint256 j = 0; - - while (oldRoot == self.rightmostNodes[level + j + 1]) { - self.rightmostNodes[level + j + 1] = node; - - unchecked { - ++s; - ++j; - } - } - - level = i + s; - } - - if ((index >> level) & 1 != 0) { - node = PoseidonT3.hash([siblingNodes[i], node]); - oldRoot = PoseidonT3.hash([siblingNodes[i], oldRoot]); - } else { - node = PoseidonT3.hash([node, siblingNodes[i]]); - oldRoot = PoseidonT3.hash([oldRoot, siblingNodes[i]]); - } - - unchecked { - ++i; - } - } - - if (oldRoot != root(self)) { - revert WrongSiblingNodes(); - } - - self.rightmostNodes[self.depth] = node; - self.leaves[newLeaf] = self.leaves[oldLeaf]; - self.leaves[oldLeaf] = 0; - - return node; + return InternalLeanIMT._update(self, oldLeaf, newLeaf, siblingNodes); } - /// @dev Removes a leaf from the tree by setting its value to zero. - /// This function utilizes the update function to set the leaf's value - /// to zero and update the tree's state accordingly. - /// @param self: A storage reference to the 'LeanIMTData' struct. - /// @param oldLeaf: The value of the leaf to be removed. - /// @param siblingNodes: An array of sibling nodes required for updating the path to the root after removal. - /// @return The new root hash of the tree after the leaf has been removed. function remove( LeanIMTData storage self, uint256 oldLeaf, uint256[] calldata siblingNodes ) public returns (uint256) { - return update(self, oldLeaf, 0, siblingNodes); + return InternalLeanIMT._remove(self, oldLeaf, siblingNodes); } - /// @dev Checks if a leaf exists in the tree. - /// @param self: A storage reference to the 'LeanIMTData' struct. - /// @param leaf: The value of the leaf to check for existence. - /// @return A boolean value indicating whether the leaf exists in the tree. function has(LeanIMTData storage self, uint256 leaf) public view returns (bool) { - return self.leaves[leaf] != 0; + return InternalLeanIMT._has(self, leaf); } - /// @dev Retrieves the index of a given leaf in the tree. - /// @param self: A storage reference to the 'LeanIMTData' struct. - /// @param leaf: The value of the leaf whose index is to be found. - /// @return The index of the specified leaf within the tree. If the leaf is not present, the function returns 0. function indexOf(LeanIMTData storage self, uint256 leaf) public view returns (uint256) { - return self.leaves[leaf] - 1; + return InternalLeanIMT._indexOf(self, leaf); } - /// @dev Retrieves the root of the tree from the 'rightmostNodes' mapping using the - /// current tree depth. - /// @param self: A storage reference to the 'LeanIMTData' struct. - /// @return The root hash of the tree. function root(LeanIMTData storage self) public view returns (uint256) { - return self.rightmostNodes[self.depth]; + return InternalLeanIMT._root(self); } } diff --git a/packages/imt.sol/contracts/QuinaryIMT.sol b/packages/imt.sol/contracts/QuinaryIMT.sol index 2a00d2991..ae0849b11 100644 --- a/packages/imt.sol/contracts/QuinaryIMT.sol +++ b/packages/imt.sol/contracts/QuinaryIMT.sol @@ -2,115 +2,19 @@ pragma solidity ^0.8.4; import {PoseidonT6} from "poseidon-solidity/PoseidonT6.sol"; +import {InternalQuinaryIMT, QuinaryIMTData} from "./internal/InternalQuinaryIMT.sol"; -// Each incremental tree has certain properties and data that will -// be used to add new leaves. -struct QuinaryIMTData { - uint256 depth; // Depth of the tree (levels - 1). - uint256 root; // Root hash of the tree. - uint256 numberOfLeaves; // Number of leaves of the tree. - mapping(uint256 => uint256) zeroes; // Zero hashes used for empty nodes (level -> zero hash). - // The nodes of the subtrees used in the last addition of a leaf (level -> [nodes]). - mapping(uint256 => uint256[5]) lastSubtrees; // Caching these values is essential to efficient appends. -} - -error ValueGreaterThanSnarkScalarField(); -error DepthNotSupported(); -error TreeIsFull(); -error NewLeafCannotEqualOldLeaf(); -error LeafDoesNotExist(); -error LeafIndexOutOfRange(); -error WrongMerkleProofPath(); - -/// @title Incremental quinary Merkle tree. -/// @dev The incremental tree allows to calculate the root hash each time a leaf is added, ensuring -/// the integrity of the tree. library QuinaryIMT { - uint8 public constant MAX_DEPTH = 32; - uint256 public constant SNARK_SCALAR_FIELD = - 21888242871839275222246405745257275088548364400416034343698204186575808495617; + using InternalQuinaryIMT for *; - /// @dev Initializes a tree. - /// @param self: Tree data. - /// @param depth: Depth of the tree. - /// @param zero: Zero value to be used. function init(QuinaryIMTData storage self, uint256 depth, uint256 zero) public { - if (zero >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (depth <= 0 || depth > MAX_DEPTH) { - revert DepthNotSupported(); - } - - self.depth = depth; - - for (uint8 i = 0; i < depth; ) { - self.zeroes[i] = zero; - uint256[5] memory zeroChildren; - - for (uint8 j = 0; j < 5; ) { - zeroChildren[j] = zero; - unchecked { - ++j; - } - } - - zero = PoseidonT6.hash(zeroChildren); - - unchecked { - ++i; - } - } - - self.root = zero; + InternalQuinaryIMT._init(self, depth, zero); } - /// @dev Inserts a leaf in the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be inserted. function insert(QuinaryIMTData storage self, uint256 leaf) public { - uint256 depth = self.depth; - - if (leaf >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (self.numberOfLeaves >= 5 ** depth) { - revert TreeIsFull(); - } - - uint256 index = self.numberOfLeaves; - uint256 hash = leaf; - - for (uint8 i = 0; i < depth; ) { - uint8 position = uint8(index % 5); - - self.lastSubtrees[i][position] = hash; - - if (position == 0) { - for (uint8 j = 1; j < 5; ) { - self.lastSubtrees[i][j] = self.zeroes[i]; - unchecked { - ++j; - } - } - } - - hash = PoseidonT6.hash(self.lastSubtrees[i]); - index /= 5; - - unchecked { - ++i; - } - } - - self.root = hash; - self.numberOfLeaves += 1; + InternalQuinaryIMT._insert(self, leaf); } - /// @dev Updates a leaf in the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be updated. - /// @param newLeaf: New leaf. - /// @param proofSiblings: Array of the sibling nodes of the proof of membership. - /// @param proofPathIndices: Path of the proof of membership. function update( QuinaryIMTData storage self, uint256 leaf, @@ -118,125 +22,24 @@ library QuinaryIMT { uint256[4][] calldata proofSiblings, uint8[] calldata proofPathIndices ) public { - if (newLeaf == leaf) { - revert NewLeafCannotEqualOldLeaf(); - } else if (newLeaf >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (!verify(self, leaf, proofSiblings, proofPathIndices)) { - revert LeafDoesNotExist(); - } - - uint256 depth = self.depth; - uint256 hash = newLeaf; - uint256 updateIndex; - - for (uint8 i = 0; i < depth; ) { - uint256[5] memory nodes; - updateIndex += proofPathIndices[i] * 5 ** i; - - for (uint8 j = 0; j < 5; ) { - if (j < proofPathIndices[i]) { - nodes[j] = proofSiblings[i][j]; - } else if (j == proofPathIndices[i]) { - nodes[j] = hash; - } else { - nodes[j] = proofSiblings[i][j - 1]; - } - unchecked { - ++j; - } - } - - if (nodes[0] == self.lastSubtrees[i][0] || nodes[4] == self.lastSubtrees[i][4]) { - self.lastSubtrees[i][proofPathIndices[i]] = hash; - } - - hash = PoseidonT6.hash(nodes); - - unchecked { - ++i; - } - } - - if (updateIndex >= self.numberOfLeaves) { - revert LeafIndexOutOfRange(); - } - - self.root = hash; + InternalQuinaryIMT._update(self, leaf, newLeaf, proofSiblings, proofPathIndices); } - /// @dev Removes a leaf from the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be removed. - /// @param proofSiblings: Array of the sibling nodes of the proof of membership. - /// @param proofPathIndices: Path of the proof of membership. function remove( QuinaryIMTData storage self, uint256 leaf, uint256[4][] calldata proofSiblings, uint8[] calldata proofPathIndices ) public { - update(self, leaf, self.zeroes[0], proofSiblings, proofPathIndices); + InternalQuinaryIMT._remove(self, leaf, proofSiblings, proofPathIndices); } - /// @dev Verify if the path is correct and the leaf is part of the tree. - /// @param self: Tree data. - /// @param leaf: Leaf to be removed. - /// @param proofSiblings: Array of the sibling nodes of the proof of membership. - /// @param proofPathIndices: Path of the proof of membership. - /// @return True or false. function verify( QuinaryIMTData storage self, uint256 leaf, uint256[4][] calldata proofSiblings, uint8[] calldata proofPathIndices ) private view returns (bool) { - uint256 depth = self.depth; - - if (leaf >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } else if (proofPathIndices.length != depth || proofSiblings.length != depth) { - revert WrongMerkleProofPath(); - } - - uint256 hash = leaf; - - for (uint8 i = 0; i < depth; ) { - uint256[5] memory nodes; - - if (proofPathIndices[i] < 0 || proofPathIndices[i] >= 5) { - revert WrongMerkleProofPath(); - } - - for (uint8 j = 0; j < 5; ) { - if (j < proofPathIndices[i]) { - if (proofSiblings[i][j] >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } - - nodes[j] = proofSiblings[i][j]; - } else if (j == proofPathIndices[i]) { - nodes[j] = hash; - } else { - if (proofSiblings[i][j - 1] >= SNARK_SCALAR_FIELD) { - revert ValueGreaterThanSnarkScalarField(); - } - - nodes[j] = proofSiblings[i][j - 1]; - } - - unchecked { - ++j; - } - } - - hash = PoseidonT6.hash(nodes); - - unchecked { - ++i; - } - } - - return hash == self.root; + return InternalQuinaryIMT._verify(self, leaf, proofSiblings, proofPathIndices); } } diff --git a/packages/imt.sol/contracts/internal/InternalBinaryIMT.sol b/packages/imt.sol/contracts/internal/InternalBinaryIMT.sol new file mode 100644 index 000000000..273f89554 --- /dev/null +++ b/packages/imt.sol/contracts/internal/InternalBinaryIMT.sol @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; +import {SNARK_SCALAR_FIELD, MAX_DEPTH} from "../Constants.sol"; + +// Each incremental tree has certain properties and data that will +// be used to add new leaves. +struct BinaryIMTData { + uint256 depth; // Depth of the tree (levels - 1). + uint256 root; // Root hash of the tree. + uint256 numberOfLeaves; // Number of leaves of the tree. + mapping(uint256 => uint256) zeroes; // Zero hashes used for empty nodes (level -> zero hash). + // The nodes of the subtrees used in the last addition of a leaf (level -> [left node, right node]). + mapping(uint256 => uint256[2]) lastSubtrees; // Caching these values is essential to efficient appends. + bool useDefaultZeroes; +} + +error ValueGreaterThanSnarkScalarField(); +error DepthNotSupported(); +error WrongDefaultZeroIndex(); +error TreeIsFull(); +error NewLeafCannotEqualOldLeaf(); +error LeafDoesNotExist(); +error LeafIndexOutOfRange(); +error WrongMerkleProofPath(); + +/// @title Incremental binary Merkle tree. +/// @dev The incremental tree allows to calculate the root hash each time a leaf is added, ensuring +/// the integrity of the tree. +library InternalBinaryIMT { + uint256 internal constant Z_0 = 0; + uint256 internal constant Z_1 = 14744269619966411208579211824598458697587494354926760081771325075741142829156; + uint256 internal constant Z_2 = 7423237065226347324353380772367382631490014989348495481811164164159255474657; + uint256 internal constant Z_3 = 11286972368698509976183087595462810875513684078608517520839298933882497716792; + uint256 internal constant Z_4 = 3607627140608796879659380071776844901612302623152076817094415224584923813162; + uint256 internal constant Z_5 = 19712377064642672829441595136074946683621277828620209496774504837737984048981; + uint256 internal constant Z_6 = 20775607673010627194014556968476266066927294572720319469184847051418138353016; + uint256 internal constant Z_7 = 3396914609616007258851405644437304192397291162432396347162513310381425243293; + uint256 internal constant Z_8 = 21551820661461729022865262380882070649935529853313286572328683688269863701601; + uint256 internal constant Z_9 = 6573136701248752079028194407151022595060682063033565181951145966236778420039; + uint256 internal constant Z_10 = 12413880268183407374852357075976609371175688755676981206018884971008854919922; + uint256 internal constant Z_11 = 14271763308400718165336499097156975241954733520325982997864342600795471836726; + uint256 internal constant Z_12 = 20066985985293572387227381049700832219069292839614107140851619262827735677018; + uint256 internal constant Z_13 = 9394776414966240069580838672673694685292165040808226440647796406499139370960; + uint256 internal constant Z_14 = 11331146992410411304059858900317123658895005918277453009197229807340014528524; + uint256 internal constant Z_15 = 15819538789928229930262697811477882737253464456578333862691129291651619515538; + uint256 internal constant Z_16 = 19217088683336594659449020493828377907203207941212636669271704950158751593251; + uint256 internal constant Z_17 = 21035245323335827719745544373081896983162834604456827698288649288827293579666; + uint256 internal constant Z_18 = 6939770416153240137322503476966641397417391950902474480970945462551409848591; + uint256 internal constant Z_19 = 10941962436777715901943463195175331263348098796018438960955633645115732864202; + uint256 internal constant Z_20 = 15019797232609675441998260052101280400536945603062888308240081994073687793470; + uint256 internal constant Z_21 = 11702828337982203149177882813338547876343922920234831094975924378932809409969; + uint256 internal constant Z_22 = 11217067736778784455593535811108456786943573747466706329920902520905755780395; + uint256 internal constant Z_23 = 16072238744996205792852194127671441602062027943016727953216607508365787157389; + uint256 internal constant Z_24 = 17681057402012993898104192736393849603097507831571622013521167331642182653248; + uint256 internal constant Z_25 = 21694045479371014653083846597424257852691458318143380497809004364947786214945; + uint256 internal constant Z_26 = 8163447297445169709687354538480474434591144168767135863541048304198280615192; + uint256 internal constant Z_27 = 14081762237856300239452543304351251708585712948734528663957353575674639038357; + uint256 internal constant Z_28 = 16619959921569409661790279042024627172199214148318086837362003702249041851090; + uint256 internal constant Z_29 = 7022159125197495734384997711896547675021391130223237843255817587255104160365; + uint256 internal constant Z_30 = 4114686047564160449611603615418567457008101555090703535405891656262658644463; + uint256 internal constant Z_31 = 12549363297364877722388257367377629555213421373705596078299904496781819142130; + uint256 internal constant Z_32 = 21443572485391568159800782191812935835534334817699172242223315142338162256601; + + function _defaultZero(uint256 index) internal pure returns (uint256) { + if (index == 0) return Z_0; + if (index == 1) return Z_1; + if (index == 2) return Z_2; + if (index == 3) return Z_3; + if (index == 4) return Z_4; + if (index == 5) return Z_5; + if (index == 6) return Z_6; + if (index == 7) return Z_7; + if (index == 8) return Z_8; + if (index == 9) return Z_9; + if (index == 10) return Z_10; + if (index == 11) return Z_11; + if (index == 12) return Z_12; + if (index == 13) return Z_13; + if (index == 14) return Z_14; + if (index == 15) return Z_15; + if (index == 16) return Z_16; + if (index == 17) return Z_17; + if (index == 18) return Z_18; + if (index == 19) return Z_19; + if (index == 20) return Z_20; + if (index == 21) return Z_21; + if (index == 22) return Z_22; + if (index == 23) return Z_23; + if (index == 24) return Z_24; + if (index == 25) return Z_25; + if (index == 26) return Z_26; + if (index == 27) return Z_27; + if (index == 28) return Z_28; + if (index == 29) return Z_29; + if (index == 30) return Z_30; + if (index == 31) return Z_31; + if (index == 32) return Z_32; + revert WrongDefaultZeroIndex(); + } + + /// @dev Initializes a tree. + /// @param self: Tree data. + /// @param depth: Depth of the tree. + /// @param zero: Zero value to be used. + function _init(BinaryIMTData storage self, uint256 depth, uint256 zero) internal { + if (zero >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (depth <= 0 || depth > MAX_DEPTH) { + revert DepthNotSupported(); + } + + self.depth = depth; + + for (uint8 i = 0; i < depth; ) { + self.zeroes[i] = zero; + zero = PoseidonT3.hash([zero, zero]); + + unchecked { + ++i; + } + } + + self.root = zero; + } + + function _initWithDefaultZeroes(BinaryIMTData storage self, uint256 depth) internal { + if (depth <= 0 || depth > MAX_DEPTH) { + revert DepthNotSupported(); + } + + self.depth = depth; + self.useDefaultZeroes = true; + + self.root = _defaultZero(depth); + } + + /// @dev Inserts a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be inserted. + function _insert(BinaryIMTData storage self, uint256 leaf) internal returns (uint256) { + uint256 depth = self.depth; + + if (leaf >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (self.numberOfLeaves >= 2 ** depth) { + revert TreeIsFull(); + } + + uint256 index = self.numberOfLeaves; + uint256 hash = leaf; + bool useDefaultZeroes = self.useDefaultZeroes; + + for (uint8 i = 0; i < depth; ) { + if (index & 1 == 0) { + self.lastSubtrees[i] = [hash, useDefaultZeroes ? _defaultZero(i) : self.zeroes[i]]; + } else { + self.lastSubtrees[i][1] = hash; + } + + hash = PoseidonT3.hash(self.lastSubtrees[i]); + index >>= 1; + + unchecked { + ++i; + } + } + + self.root = hash; + self.numberOfLeaves += 1; + return hash; + } + + /// @dev Updates a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be updated. + /// @param newLeaf: New leaf. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + function _update( + BinaryIMTData storage self, + uint256 leaf, + uint256 newLeaf, + uint256[] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) internal { + if (newLeaf == leaf) { + revert NewLeafCannotEqualOldLeaf(); + } else if (newLeaf >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (!_verify(self, leaf, proofSiblings, proofPathIndices)) { + revert LeafDoesNotExist(); + } + + uint256 depth = self.depth; + uint256 hash = newLeaf; + uint256 updateIndex; + + for (uint8 i = 0; i < depth; ) { + updateIndex |= uint256(proofPathIndices[i]) << uint256(i); + + if (proofPathIndices[i] == 0) { + if (proofSiblings[i] == self.lastSubtrees[i][1]) { + self.lastSubtrees[i][0] = hash; + } + + hash = PoseidonT3.hash([hash, proofSiblings[i]]); + } else { + if (proofSiblings[i] == self.lastSubtrees[i][0]) { + self.lastSubtrees[i][1] = hash; + } + + hash = PoseidonT3.hash([proofSiblings[i], hash]); + } + + unchecked { + ++i; + } + } + + if (updateIndex >= self.numberOfLeaves) { + revert LeafIndexOutOfRange(); + } + + self.root = hash; + } + + /// @dev Removes a leaf from the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be removed. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + function _remove( + BinaryIMTData storage self, + uint256 leaf, + uint256[] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) internal { + _update(self, leaf, self.useDefaultZeroes ? Z_0 : self.zeroes[0], proofSiblings, proofPathIndices); + } + + /// @dev Verify if the path is correct and the leaf is part of the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be removed. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + /// @return True or false. + function _verify( + BinaryIMTData storage self, + uint256 leaf, + uint256[] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) internal view returns (bool) { + uint256 depth = self.depth; + + if (leaf >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (proofPathIndices.length != depth || proofSiblings.length != depth) { + revert WrongMerkleProofPath(); + } + + uint256 hash = leaf; + + for (uint8 i = 0; i < depth; ) { + if (proofSiblings[i] >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (proofPathIndices[i] != 1 && proofPathIndices[i] != 0) { + revert WrongMerkleProofPath(); + } + + if (proofPathIndices[i] == 0) { + hash = PoseidonT3.hash([hash, proofSiblings[i]]); + } else { + hash = PoseidonT3.hash([proofSiblings[i], hash]); + } + + unchecked { + ++i; + } + } + + return hash == self.root; + } +} diff --git a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol new file mode 100644 index 000000000..93fe7b8b5 --- /dev/null +++ b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; +import {SNARK_SCALAR_FIELD, MAX_DEPTH} from "../Constants.sol"; + +struct LazyIMTData { + uint32 maxIndex; + uint40 numberOfLeaves; + mapping(uint256 => uint256) elements; +} + +library InternalLazyIMT { + uint40 internal constant MAX_INDEX = (1 << 32) - 1; + + uint256 internal constant Z_0 = 0; + uint256 internal constant Z_1 = 14744269619966411208579211824598458697587494354926760081771325075741142829156; + uint256 internal constant Z_2 = 7423237065226347324353380772367382631490014989348495481811164164159255474657; + uint256 internal constant Z_3 = 11286972368698509976183087595462810875513684078608517520839298933882497716792; + uint256 internal constant Z_4 = 3607627140608796879659380071776844901612302623152076817094415224584923813162; + uint256 internal constant Z_5 = 19712377064642672829441595136074946683621277828620209496774504837737984048981; + uint256 internal constant Z_6 = 20775607673010627194014556968476266066927294572720319469184847051418138353016; + uint256 internal constant Z_7 = 3396914609616007258851405644437304192397291162432396347162513310381425243293; + uint256 internal constant Z_8 = 21551820661461729022865262380882070649935529853313286572328683688269863701601; + uint256 internal constant Z_9 = 6573136701248752079028194407151022595060682063033565181951145966236778420039; + uint256 internal constant Z_10 = 12413880268183407374852357075976609371175688755676981206018884971008854919922; + uint256 internal constant Z_11 = 14271763308400718165336499097156975241954733520325982997864342600795471836726; + uint256 internal constant Z_12 = 20066985985293572387227381049700832219069292839614107140851619262827735677018; + uint256 internal constant Z_13 = 9394776414966240069580838672673694685292165040808226440647796406499139370960; + uint256 internal constant Z_14 = 11331146992410411304059858900317123658895005918277453009197229807340014528524; + uint256 internal constant Z_15 = 15819538789928229930262697811477882737253464456578333862691129291651619515538; + uint256 internal constant Z_16 = 19217088683336594659449020493828377907203207941212636669271704950158751593251; + uint256 internal constant Z_17 = 21035245323335827719745544373081896983162834604456827698288649288827293579666; + uint256 internal constant Z_18 = 6939770416153240137322503476966641397417391950902474480970945462551409848591; + uint256 internal constant Z_19 = 10941962436777715901943463195175331263348098796018438960955633645115732864202; + uint256 internal constant Z_20 = 15019797232609675441998260052101280400536945603062888308240081994073687793470; + uint256 internal constant Z_21 = 11702828337982203149177882813338547876343922920234831094975924378932809409969; + uint256 internal constant Z_22 = 11217067736778784455593535811108456786943573747466706329920902520905755780395; + uint256 internal constant Z_23 = 16072238744996205792852194127671441602062027943016727953216607508365787157389; + uint256 internal constant Z_24 = 17681057402012993898104192736393849603097507831571622013521167331642182653248; + uint256 internal constant Z_25 = 21694045479371014653083846597424257852691458318143380497809004364947786214945; + uint256 internal constant Z_26 = 8163447297445169709687354538480474434591144168767135863541048304198280615192; + uint256 internal constant Z_27 = 14081762237856300239452543304351251708585712948734528663957353575674639038357; + uint256 internal constant Z_28 = 16619959921569409661790279042024627172199214148318086837362003702249041851090; + uint256 internal constant Z_29 = 7022159125197495734384997711896547675021391130223237843255817587255104160365; + uint256 internal constant Z_30 = 4114686047564160449611603615418567457008101555090703535405891656262658644463; + uint256 internal constant Z_31 = 12549363297364877722388257367377629555213421373705596078299904496781819142130; + uint256 internal constant Z_32 = 21443572485391568159800782191812935835534334817699172242223315142338162256601; + + function _defaultZero(uint8 index) internal pure returns (uint256) { + if (index == 0) return Z_0; + if (index == 1) return Z_1; + if (index == 2) return Z_2; + if (index == 3) return Z_3; + if (index == 4) return Z_4; + if (index == 5) return Z_5; + if (index == 6) return Z_6; + if (index == 7) return Z_7; + if (index == 8) return Z_8; + if (index == 9) return Z_9; + if (index == 10) return Z_10; + if (index == 11) return Z_11; + if (index == 12) return Z_12; + if (index == 13) return Z_13; + if (index == 14) return Z_14; + if (index == 15) return Z_15; + if (index == 16) return Z_16; + if (index == 17) return Z_17; + if (index == 18) return Z_18; + if (index == 19) return Z_19; + if (index == 20) return Z_20; + if (index == 21) return Z_21; + if (index == 22) return Z_22; + if (index == 23) return Z_23; + if (index == 24) return Z_24; + if (index == 25) return Z_25; + if (index == 26) return Z_26; + if (index == 27) return Z_27; + if (index == 28) return Z_28; + if (index == 29) return Z_29; + if (index == 30) return Z_30; + if (index == 31) return Z_31; + if (index == 32) return Z_32; + revert("LazyIMT: defaultZero bad index"); + } + + function _init(LazyIMTData storage self, uint8 depth) internal { + require(depth <= MAX_DEPTH, "LazyIMT: Tree too large"); + self.maxIndex = uint32((1 << depth) - 1); + self.numberOfLeaves = 0; + } + + function _reset(LazyIMTData storage self) internal { + self.numberOfLeaves = 0; + } + + function _indexForElement(uint8 level, uint40 index) internal pure returns (uint40) { + // store the elements sparsely + return MAX_INDEX * level + index; + } + + function _insert(LazyIMTData storage self, uint256 leaf) internal { + uint40 index = self.numberOfLeaves; + require(leaf < SNARK_SCALAR_FIELD, "LazyIMT: leaf must be < SNARK_SCALAR_FIELD"); + require(index < self.maxIndex, "LazyIMT: tree is full"); + + self.numberOfLeaves = index + 1; + + uint256 hash = leaf; + + for (uint8 i = 0; ; ) { + self.elements[_indexForElement(i, index)] = hash; + // it's a left element so we don't hash until there's a right element + if (index & 1 == 0) break; + uint40 elementIndex = _indexForElement(i, index - 1); + hash = PoseidonT3.hash([self.elements[elementIndex], hash]); + unchecked { + index >>= 1; + i++; + } + } + } + + function _update(LazyIMTData storage self, uint256 leaf, uint40 index) internal { + require(leaf < SNARK_SCALAR_FIELD, "LazyIMT: leaf must be < SNARK_SCALAR_FIELD"); + uint40 numberOfLeaves = self.numberOfLeaves; + require(index < numberOfLeaves, "LazyIMT: leaf must exist"); + + uint256 hash = leaf; + + for (uint8 i = 0; true; ) { + self.elements[_indexForElement(i, index)] = hash; + uint256 levelCount = numberOfLeaves >> (i + 1); + if (levelCount <= index >> 1) break; + if (index & 1 == 0) { + uint40 elementIndex = _indexForElement(i, index + 1); + hash = PoseidonT3.hash([hash, self.elements[elementIndex]]); + } else { + uint40 elementIndex = _indexForElement(i, index - 1); + hash = PoseidonT3.hash([self.elements[elementIndex], hash]); + } + unchecked { + index >>= 1; + i++; + } + } + } + + function _root(LazyIMTData storage self) internal view returns (uint256) { + // this will always short circuit if self.numberOfLeaves == 0 + uint40 numberOfLeaves = self.numberOfLeaves; + // dynamically determine a depth + uint8 depth = 1; + while (uint8(2) ** depth < numberOfLeaves) { + depth++; + } + return _root(self, numberOfLeaves, depth); + } + + function _root(LazyIMTData storage self, uint8 depth) internal view returns (uint256) { + uint40 numberOfLeaves = self.numberOfLeaves; + require(2 ** depth >= numberOfLeaves, "LazyIMT: ambiguous depth"); + return _root(self, self.numberOfLeaves, depth); + } + + function _root(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth) internal view returns (uint256) { + require(depth > 0, "LazyIMT: depth must be > 0"); + require(depth <= MAX_DEPTH, "LazyIMT: depth must be < MAX_DEPTH"); + // this should always short circuit if self.numberOfLeaves == 0 + if (numberOfLeaves == 0) return _defaultZero(depth); + uint40 index = numberOfLeaves - 1; + + uint256[MAX_DEPTH + 1] memory levels; + + if (index & 1 == 0) { + levels[0] = self.elements[_indexForElement(0, index)]; + } else { + levels[0] = _defaultZero(0); + } + + for (uint8 i = 0; i < depth; ) { + if (index & 1 == 0) { + levels[i + 1] = PoseidonT3.hash([levels[i], _defaultZero(i)]); + } else { + uint256 levelCount = (numberOfLeaves) >> (i + 1); + if (levelCount > index >> 1) { + uint256 parent = self.elements[_indexForElement(i + 1, index >> 1)]; + levels[i + 1] = parent; + } else { + uint256 sibling = self.elements[_indexForElement(i, index - 1)]; + levels[i + 1] = PoseidonT3.hash([sibling, levels[i]]); + } + } + unchecked { + index >>= 1; + i++; + } + } + return levels[depth]; + } +} diff --git a/packages/imt.sol/contracts/internal/InternalLeanIMT.sol b/packages/imt.sol/contracts/internal/InternalLeanIMT.sol new file mode 100644 index 000000000..be0b695d9 --- /dev/null +++ b/packages/imt.sol/contracts/internal/InternalLeanIMT.sol @@ -0,0 +1,202 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; +import {SNARK_SCALAR_FIELD} from "../Constants.sol"; + +struct LeanIMTData { + // Tracks the current number of leaves in the tree. + uint256 size; + // Represents the current depth of the tree, which can increase as new leaves are inserted. + uint256 depth; + // A mapping from each level of the tree to the rightmost node at that level. + // Used for efficient updates and root calculations. + mapping(uint256 => uint256) rightmostNodes; + // A mapping from leaf values to their respective indices in the tree. + // This facilitates checks for leaf existence and retrieval of leaf positions. + mapping(uint256 => uint256) leaves; +} + +error WrongSiblingNodes(); +error LeafGreaterThanSnarkScalarField(); +error LeafCannotBeZero(); +error LeafAlreadyExists(); +error LeafDoesNotExist(); + +/// @title Lean Incremental binary Merkle tree. +/// @dev The LeanIMT is an optimized version of the BinaryIMT. +/// This implementation eliminates the use of zeroes, and make the tree depth dynamic. +/// When a node doesn't have the right child, instead of using a zero hash as in the BinaryIMT, +/// the node's value becomes that of its left child. Furthermore, rather than utilizing a static tree depth, +/// it is updated based on the number of leaves in the tree. This approach +/// results in the calculation of significantly fewer hashes, making the tree more efficient. +library InternalLeanIMT { + /// @dev Inserts a new leaf into the incremental merkle tree. + /// The function ensures that the leaf is valid according to the + /// constraints of the tree and then updates the tree's structure accordingly. + /// @param self: A storage reference to the 'LeanIMTData' struct. + /// @param leaf: The value of the new leaf to be inserted into the tree. + /// @return The new hash of the node after the leaf has been inserted. + function _insert(LeanIMTData storage self, uint256 leaf) internal returns (uint256) { + if (leaf >= SNARK_SCALAR_FIELD) { + revert LeafGreaterThanSnarkScalarField(); + } else if (leaf == 0) { + revert LeafCannotBeZero(); + } else if (_has(self, leaf)) { + revert LeafAlreadyExists(); + } + + while (2 ** self.depth < self.size + 1) { + self.depth += 1; + } + + uint256 index = self.size; + uint256 node = leaf; + + for (uint256 level = 0; level < self.depth; ) { + if ((index >> level) & 1 == 1) { + node = PoseidonT3.hash([self.rightmostNodes[level], node]); + } else { + self.rightmostNodes[level] = node; + } + + unchecked { + ++level; + } + } + + self.size += 1; + + self.rightmostNodes[self.depth] = node; + self.leaves[leaf] = self.size; + + return node; + } + + /// @dev Updates the value of an existing leaf and recalculates hashes + /// to maintain tree integrity. + /// @param self: A storage reference to the 'LeanIMTData' struct. + /// @param oldLeaf: The value of the leaf that is to be updated. + /// @param newLeaf: The new value that will replace the oldLeaf in the tree. + /// @param siblingNodes: An array of sibling nodes that are necessary to recalculate the path to the root. + /// @return The new hash of the updated node after the leaf has been updated. + function _update( + LeanIMTData storage self, + uint256 oldLeaf, + uint256 newLeaf, + uint256[] calldata siblingNodes + ) internal returns (uint256) { + if (newLeaf >= SNARK_SCALAR_FIELD) { + revert LeafGreaterThanSnarkScalarField(); + } else if (!_has(self, oldLeaf)) { + revert LeafDoesNotExist(); + } else if (newLeaf != 0 && _has(self, newLeaf)) { + revert LeafAlreadyExists(); + } + + uint256 index = _indexOf(self, oldLeaf); + uint256 node = newLeaf; + uint256 oldRoot = oldLeaf; + + // A counter that adjusts the level at which sibling nodes are + // accessed and updated during the tree's update process. + // It ensures that the update function correctly navigates and + // modifies the tree's nodes at the appropriate levels, accounting + // for situations where not every level of the tree requires an + // update or a hash calculation. + uint256 s = 0; + + // The number of siblings of a proof can be less than + // the depth of the tree, because in some levels it might not + // be necessary to hash any value. + for (uint256 i = 0; i < siblingNodes.length; ) { + if (siblingNodes[i] >= SNARK_SCALAR_FIELD) { + revert LeafGreaterThanSnarkScalarField(); + } + + uint256 level = i + s; + + if (oldRoot == self.rightmostNodes[level]) { + self.rightmostNodes[level] = node; + + if (oldRoot == self.rightmostNodes[level + 1]) { + s += 1; + } + + uint256 j = 0; + + while (oldRoot == self.rightmostNodes[level + j + 1]) { + self.rightmostNodes[level + j + 1] = node; + + unchecked { + ++s; + ++j; + } + } + + level = i + s; + } + + if ((index >> level) & 1 != 0) { + node = PoseidonT3.hash([siblingNodes[i], node]); + oldRoot = PoseidonT3.hash([siblingNodes[i], oldRoot]); + } else { + node = PoseidonT3.hash([node, siblingNodes[i]]); + oldRoot = PoseidonT3.hash([oldRoot, siblingNodes[i]]); + } + + unchecked { + ++i; + } + } + + if (oldRoot != _root(self)) { + revert WrongSiblingNodes(); + } + + self.rightmostNodes[self.depth] = node; + self.leaves[newLeaf] = self.leaves[oldLeaf]; + self.leaves[oldLeaf] = 0; + + return node; + } + + /// @dev Removes a leaf from the tree by setting its value to zero. + /// This function utilizes the update function to set the leaf's value + /// to zero and update the tree's state accordingly. + /// @param self: A storage reference to the 'LeanIMTData' struct. + /// @param oldLeaf: The value of the leaf to be removed. + /// @param siblingNodes: An array of sibling nodes required for updating the path to the root after removal. + /// @return The new root hash of the tree after the leaf has been removed. + function _remove( + LeanIMTData storage self, + uint256 oldLeaf, + uint256[] calldata siblingNodes + ) internal returns (uint256) { + return _update(self, oldLeaf, 0, siblingNodes); + } + + /// @dev Checks if a leaf exists in the tree. + /// @param self: A storage reference to the 'LeanIMTData' struct. + /// @param leaf: The value of the leaf to check for existence. + /// @return A boolean value indicating whether the leaf exists in the tree. + function _has(LeanIMTData storage self, uint256 leaf) internal view returns (bool) { + return self.leaves[leaf] != 0; + } + + /// @dev Retrieves the index of a given leaf in the tree. + /// @param self: A storage reference to the 'LeanIMTData' struct. + /// @param leaf: The value of the leaf whose index is to be found. + /// @return The index of the specified leaf within the tree. If the leaf is not present, the function returns 0. + function _indexOf(LeanIMTData storage self, uint256 leaf) internal view returns (uint256) { + return self.leaves[leaf] - 1; + } + + /// @dev Retrieves the root of the tree from the 'rightmostNodes' mapping using the + /// current tree depth. + /// @param self: A storage reference to the 'LeanIMTData' struct. + /// @return The root hash of the tree. + function _root(LeanIMTData storage self) internal view returns (uint256) { + return self.rightmostNodes[self.depth]; + } +} diff --git a/packages/imt.sol/contracts/internal/InternalQuinaryIMT.sol b/packages/imt.sol/contracts/internal/InternalQuinaryIMT.sol new file mode 100644 index 000000000..93f60732c --- /dev/null +++ b/packages/imt.sol/contracts/internal/InternalQuinaryIMT.sol @@ -0,0 +1,241 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {PoseidonT6} from "poseidon-solidity/PoseidonT6.sol"; +import {SNARK_SCALAR_FIELD, MAX_DEPTH} from "../Constants.sol"; + +// Each incremental tree has certain properties and data that will +// be used to add new leaves. +struct QuinaryIMTData { + uint256 depth; // Depth of the tree (levels - 1). + uint256 root; // Root hash of the tree. + uint256 numberOfLeaves; // Number of leaves of the tree. + mapping(uint256 => uint256) zeroes; // Zero hashes used for empty nodes (level -> zero hash). + // The nodes of the subtrees used in the last addition of a leaf (level -> [nodes]). + mapping(uint256 => uint256[5]) lastSubtrees; // Caching these values is essential to efficient appends. +} + +error ValueGreaterThanSnarkScalarField(); +error DepthNotSupported(); +error TreeIsFull(); +error NewLeafCannotEqualOldLeaf(); +error LeafDoesNotExist(); +error LeafIndexOutOfRange(); +error WrongMerkleProofPath(); + +/// @title Incremental quinary Merkle tree. +/// @dev The incremental tree allows to calculate the root hash each time a leaf is added, ensuring +/// the integrity of the tree. +library InternalQuinaryIMT { + /// @dev Initializes a tree. + /// @param self: Tree data. + /// @param depth: Depth of the tree. + /// @param zero: Zero value to be used. + function _init(QuinaryIMTData storage self, uint256 depth, uint256 zero) internal { + if (zero >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (depth <= 0 || depth > MAX_DEPTH) { + revert DepthNotSupported(); + } + + self.depth = depth; + + for (uint8 i = 0; i < depth; ) { + self.zeroes[i] = zero; + uint256[5] memory zeroChildren; + + for (uint8 j = 0; j < 5; ) { + zeroChildren[j] = zero; + unchecked { + ++j; + } + } + + zero = PoseidonT6.hash(zeroChildren); + + unchecked { + ++i; + } + } + + self.root = zero; + } + + /// @dev Inserts a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be inserted. + function _insert(QuinaryIMTData storage self, uint256 leaf) internal { + uint256 depth = self.depth; + + if (leaf >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (self.numberOfLeaves >= 5 ** depth) { + revert TreeIsFull(); + } + + uint256 index = self.numberOfLeaves; + uint256 hash = leaf; + + for (uint8 i = 0; i < depth; ) { + uint8 position = uint8(index % 5); + + self.lastSubtrees[i][position] = hash; + + if (position == 0) { + for (uint8 j = 1; j < 5; ) { + self.lastSubtrees[i][j] = self.zeroes[i]; + unchecked { + ++j; + } + } + } + + hash = PoseidonT6.hash(self.lastSubtrees[i]); + index /= 5; + + unchecked { + ++i; + } + } + + self.root = hash; + self.numberOfLeaves += 1; + } + + /// @dev Updates a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be updated. + /// @param newLeaf: New leaf. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + function _update( + QuinaryIMTData storage self, + uint256 leaf, + uint256 newLeaf, + uint256[4][] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) internal { + if (newLeaf == leaf) { + revert NewLeafCannotEqualOldLeaf(); + } else if (newLeaf >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (!_verify(self, leaf, proofSiblings, proofPathIndices)) { + revert LeafDoesNotExist(); + } + + uint256 depth = self.depth; + uint256 hash = newLeaf; + uint256 updateIndex; + + for (uint8 i = 0; i < depth; ) { + uint256[5] memory nodes; + updateIndex += proofPathIndices[i] * 5 ** i; + + for (uint8 j = 0; j < 5; ) { + if (j < proofPathIndices[i]) { + nodes[j] = proofSiblings[i][j]; + } else if (j == proofPathIndices[i]) { + nodes[j] = hash; + } else { + nodes[j] = proofSiblings[i][j - 1]; + } + unchecked { + ++j; + } + } + + if (nodes[0] == self.lastSubtrees[i][0] || nodes[4] == self.lastSubtrees[i][4]) { + self.lastSubtrees[i][proofPathIndices[i]] = hash; + } + + hash = PoseidonT6.hash(nodes); + + unchecked { + ++i; + } + } + + if (updateIndex >= self.numberOfLeaves) { + revert LeafIndexOutOfRange(); + } + + self.root = hash; + } + + /// @dev Removes a leaf from the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be removed. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + function _remove( + QuinaryIMTData storage self, + uint256 leaf, + uint256[4][] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) internal { + _update(self, leaf, self.zeroes[0], proofSiblings, proofPathIndices); + } + + /// @dev Verify if the path is correct and the leaf is part of the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be removed. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + /// @return True or false. + function _verify( + QuinaryIMTData storage self, + uint256 leaf, + uint256[4][] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) internal view returns (bool) { + uint256 depth = self.depth; + + if (leaf >= SNARK_SCALAR_FIELD) { + revert ValueGreaterThanSnarkScalarField(); + } else if (proofPathIndices.length != depth || proofSiblings.length != depth) { + revert WrongMerkleProofPath(); + } + + uint256 hash = leaf; + + for (uint8 i = 0; i < depth; ) { + uint256[5] memory nodes; + + if (proofPathIndices[i] < 0 || proofPathIndices[i] >= 5) { + revert WrongMerkleProofPath(); + } + + for (uint8 j = 0; j < 5; ) { + if (j < proofPathIndices[i]) { + require( + proofSiblings[i][j] < SNARK_SCALAR_FIELD, + "QuinaryIMT: sibling node must be < SNARK_SCALAR_FIELD" + ); + + nodes[j] = proofSiblings[i][j]; + } else if (j == proofPathIndices[i]) { + nodes[j] = hash; + } else { + require( + proofSiblings[i][j - 1] < SNARK_SCALAR_FIELD, + "QuinaryIMT: sibling node must be < SNARK_SCALAR_FIELD" + ); + + nodes[j] = proofSiblings[i][j - 1]; + } + + unchecked { + ++j; + } + } + + hash = PoseidonT6.hash(nodes); + + unchecked { + ++i; + } + } + + return hash == self.root; + } +} diff --git a/packages/imt.sol/test/BinaryIMT.ts b/packages/imt.sol/test/BinaryIMT.ts index 77ae8a673..ef11ac4c0 100644 --- a/packages/imt.sol/test/BinaryIMT.ts +++ b/packages/imt.sol/test/BinaryIMT.ts @@ -6,15 +6,16 @@ import { poseidon2 } from "poseidon-lite" import { BinaryIMT, BinaryIMTTest } from "../typechain-types" describe("BinaryIMT", () => { - let binaryIMTTest: BinaryIMTTest + const SNARK_SCALAR_FIELD = BigInt("21888242871839275222246405745257275088548364400416034343698204186575808495617") let binaryIMT: BinaryIMT + let binaryIMTTest: BinaryIMTTest let jsBinaryIMT: JSBinaryIMT beforeEach(async () => { const { library, contract } = await run("deploy:imt-test", { library: "BinaryIMT", logs: false }) - binaryIMTTest = contract binaryIMT = library + binaryIMTTest = contract jsBinaryIMT = new JSBinaryIMT(poseidon2, 16, 0, 2) }) @@ -45,9 +46,7 @@ describe("BinaryIMT", () => { describe("# insert", () => { it("Should not insert a leaf if its value is > SNARK_SCALAR_FIELD", async () => { - const leaf = await binaryIMT.SNARK_SCALAR_FIELD() - - const transaction = binaryIMTTest.insert(leaf) + const transaction = binaryIMTTest.insert(SNARK_SCALAR_FIELD) await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField") }) @@ -133,9 +132,7 @@ describe("BinaryIMT", () => { await binaryIMTTest.init(jsBinaryIMT.depth) await binaryIMTTest.insert(1) - const newLeaf = await binaryIMT.SNARK_SCALAR_FIELD() - - const transaction = binaryIMTTest.update(1, newLeaf, [0, 1], [0, 1]) + const transaction = binaryIMTTest.update(1, SNARK_SCALAR_FIELD, [0, 1], [0, 1]) await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField") }) @@ -144,9 +141,7 @@ describe("BinaryIMT", () => { await binaryIMTTest.init(jsBinaryIMT.depth) await binaryIMTTest.insert(1) - const oldLeaf = await binaryIMT.SNARK_SCALAR_FIELD() - - const transaction = binaryIMTTest.update(oldLeaf, 2, [0, 1], [0, 1]) + const transaction = binaryIMTTest.update(SNARK_SCALAR_FIELD, 2, [0, 1], [0, 1]) await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField") }) @@ -257,9 +252,7 @@ describe("BinaryIMT", () => { describe("# remove", () => { it("Should not remove a leaf if its value is > SNARK_SCALAR_FIELD", async () => { - const leaf = await binaryIMT.SNARK_SCALAR_FIELD() - - const transaction = binaryIMTTest.remove(leaf, [0, 1], [0, 1]) + const transaction = binaryIMTTest.remove(SNARK_SCALAR_FIELD, [0, 1], [0, 1]) await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField") }) diff --git a/packages/imt.sol/test/LazyIMT.ts b/packages/imt.sol/test/LazyIMT.ts index 4f62c8aff..4087b72a1 100644 --- a/packages/imt.sol/test/LazyIMT.ts +++ b/packages/imt.sol/test/LazyIMT.ts @@ -7,6 +7,7 @@ import { LazyIMT, LazyIMTTest } from "../typechain-types" const random = () => poseidon2([Math.floor(Math.random() * 2 ** 40), 0]) describe("LazyIMT", () => { + const SNARK_SCALAR_FIELD = BigInt("21888242871839275222246405745257275088548364400416034343698204186575808495617") let lazyIMTTest: LazyIMTTest let lazyIMT: LazyIMT @@ -129,9 +130,9 @@ describe("LazyIMT", () => { await lazyIMTTest.init(depth) - const F = await lazyIMT.SNARK_SCALAR_FIELD() - - await expect(lazyIMTTest.insert(F)).to.be.revertedWith("LazyIMT: leaf must be < SNARK_SCALAR_FIELD") + await expect(lazyIMTTest.insert(SNARK_SCALAR_FIELD)).to.be.revertedWith( + "LazyIMT: leaf must be < SNARK_SCALAR_FIELD" + ) }) }) @@ -217,9 +218,9 @@ describe("LazyIMT", () => { await lazyIMTTest.insert(random()) } - const F = await lazyIMT.SNARK_SCALAR_FIELD() - - await expect(lazyIMTTest.update(F, 0)).to.be.revertedWith("LazyIMT: leaf must be < SNARK_SCALAR_FIELD") + await expect(lazyIMTTest.update(SNARK_SCALAR_FIELD, 0)).to.be.revertedWith( + "LazyIMT: leaf must be < SNARK_SCALAR_FIELD" + ) }) }) diff --git a/packages/imt.sol/test/LeanIMT.ts b/packages/imt.sol/test/LeanIMT.ts index 2a0494daf..80c49370d 100644 --- a/packages/imt.sol/test/LeanIMT.ts +++ b/packages/imt.sol/test/LeanIMT.ts @@ -5,6 +5,7 @@ import { poseidon2 } from "poseidon-lite" import { LeanIMT, LeanIMTTest } from "../typechain-types" describe("LeanIMT", () => { + const SNARK_SCALAR_FIELD = BigInt("21888242871839275222246405745257275088548364400416034343698204186575808495617") let leanIMTTest: LeanIMTTest let leanIMT: LeanIMT let jsLeanIMT: JSLeanIMT @@ -19,9 +20,7 @@ describe("LeanIMT", () => { describe("# insert", () => { it("Should not insert a leaf if its value is > SNARK_SCALAR_FIELD", async () => { - const leaf = await leanIMT.SNARK_SCALAR_FIELD() - - const transaction = leanIMTTest.insert(leaf) + const transaction = leanIMTTest.insert(SNARK_SCALAR_FIELD) await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField") }) @@ -73,9 +72,7 @@ describe("LeanIMT", () => { }) it("Should not update a leaf if its value is > SNARK_SCALAR_FIELD", async () => { - const leaf = await leanIMT.SNARK_SCALAR_FIELD() - - const transaction = leanIMTTest.update(2, leaf, [1, 2, 3, 4]) + const transaction = leanIMTTest.update(2, SNARK_SCALAR_FIELD, [1, 2, 3, 4]) await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField") }) @@ -103,7 +100,7 @@ describe("LeanIMT", () => { const { siblings } = jsLeanIMT.generateProof(0) - siblings[0] = (await leanIMT.SNARK_SCALAR_FIELD()).toBigInt() + siblings[0] = SNARK_SCALAR_FIELD const transaction = leanIMTTest.update(1, 2, siblings) diff --git a/packages/imt.sol/test/QuinaryIMT.ts b/packages/imt.sol/test/QuinaryIMT.ts index 8d570f389..674a6ba08 100644 --- a/packages/imt.sol/test/QuinaryIMT.ts +++ b/packages/imt.sol/test/QuinaryIMT.ts @@ -5,15 +5,16 @@ import { poseidon5 } from "poseidon-lite" import { QuinaryIMT, QuinaryIMTTest } from "../typechain-types" describe("QuinaryIMT", () => { - let quinaryIMTTest: QuinaryIMTTest + const SNARK_SCALAR_FIELD = BigInt("21888242871839275222246405745257275088548364400416034343698204186575808495617") let quinaryIMT: QuinaryIMT + let quinaryIMTTest: QuinaryIMTTest let jsQuinaryIMT: JSQuinaryIMT beforeEach(async () => { const { library, contract } = await run("deploy:imt-test", { library: "QuinaryIMT", arity: 5, logs: false }) - quinaryIMTTest = contract quinaryIMT = library + quinaryIMTTest = contract jsQuinaryIMT = new JSQuinaryIMT(poseidon5, 16, 0, 5) }) @@ -35,9 +36,7 @@ describe("QuinaryIMT", () => { describe("# insert", () => { it("Should not insert a leaf if its value is > SNARK_SCALAR_FIELD", async () => { - const leaf = await quinaryIMT.SNARK_SCALAR_FIELD() - - const transaction = quinaryIMTTest.insert(leaf) + const transaction = quinaryIMTTest.insert(SNARK_SCALAR_FIELD) await expect(transaction).to.be.revertedWithCustomError(quinaryIMT, "ValueGreaterThanSnarkScalarField") }) @@ -96,9 +95,7 @@ describe("QuinaryIMT", () => { await quinaryIMTTest.init(jsQuinaryIMT.depth) await quinaryIMTTest.insert(1) - const newLeaf = await quinaryIMT.SNARK_SCALAR_FIELD() - - const transaction = quinaryIMTTest.update(1, newLeaf, [[0, 1, 2, 3]], [0]) + const transaction = quinaryIMTTest.update(1, SNARK_SCALAR_FIELD, [[0, 1, 2, 3]], [0]) await expect(transaction).to.be.revertedWithCustomError(quinaryIMT, "ValueGreaterThanSnarkScalarField") }) @@ -107,9 +104,7 @@ describe("QuinaryIMT", () => { await quinaryIMTTest.init(jsQuinaryIMT.depth) await quinaryIMTTest.insert(1) - const oldLeaf = await quinaryIMT.SNARK_SCALAR_FIELD() - - const transaction = quinaryIMTTest.update(oldLeaf, 2, [[0, 1, 2, 3]], [0]) + const transaction = quinaryIMTTest.update(SNARK_SCALAR_FIELD, 2, [[0, 1, 2, 3]], [0]) await expect(transaction).to.be.revertedWithCustomError(quinaryIMT, "ValueGreaterThanSnarkScalarField") }) @@ -200,9 +195,7 @@ describe("QuinaryIMT", () => { describe("# remove", () => { it("Should not remove a leaf if its value is > SNARK_SCALAR_FIELD", async () => { - const leaf = await quinaryIMT.SNARK_SCALAR_FIELD() - - const transaction = quinaryIMTTest.remove(leaf, [[0, 1, 2, 3]], [0]) + const transaction = quinaryIMTTest.remove(SNARK_SCALAR_FIELD, [[0, 1, 2, 3]], [0]) await expect(transaction).to.be.revertedWithCustomError(quinaryIMT, "ValueGreaterThanSnarkScalarField") })