diff --git a/CHANGELOG.md b/CHANGELOG.md index 53531d2d..2244c017 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,8 @@ # Changelog -## [none] +## [patch] -- Fixed deprecation warnings. +- Added `IndexedMerkleTree` library that implements Indexed Merkle Tree data structure. ## [3.3.2] diff --git a/README.md b/README.md index 8f493b74..eb5852dd 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,7 @@ contracts │ │ ├── CartesianMerkleTree — "CMT reference implementation" │ │ ├── DynamicSet — "Set for strings and bytes" │ │ ├── IncrementalMerkleTree — "IMT implementation with flexible tree height" +│ │ ├── IndexedMerkleTree — "IndexedMT implementation" │ │ ├── PriorityQueue — "Max queue heap implementation" │ │ ├── SparseMerkleTree — "SMT optimized implementation" │ │ └── memory diff --git a/contracts/libs/data-structures/IndexedMerkleTree.sol b/contracts/libs/data-structures/IndexedMerkleTree.sol new file mode 100644 index 00000000..cee59183 --- /dev/null +++ b/contracts/libs/data-structures/IndexedMerkleTree.sol @@ -0,0 +1,1178 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.21; + +/** + * @notice Indexed Merkle Tree Module + * + * Gas usage for adding and updating 100 elements to an IndexedMT with the keccak256 and poseidon hash functions is detailed below: + * + * Keccak256: + * - CMT.add - 249k + * - CMT.update - 250k + * + * Poseidon: + * - CMT.add - 1.13m + * - CMT.update - 1.13m + * + * Custom hashing functions can be provided before initialization to change how nodes and + * leaves are hashed (useful for e.g. Poseidon-based hashing in zk environments). Default + * keccak based hashing functions are used when custom hashers are not set. + * + * ## Usage Example + * + * ```solidity + * using IndexedMerkleTree for IndexedMerkleTree.UintIndexedMT; + * + * IndexedMerkleTree.UintIndexedMT internal tree; + * + * tree.setHashers(hashFunctions); + * tree.initialize(); + * + * uint256 leafIndex = tree.add(42, 0); + * + * IndexedMerkleTree.Proof memory proof = tree.getProof(leafIndex, 42); + * + * bool ok = tree.verifyProof(proof); + *``` + */ + +library IndexedMerkleTree { + /** + ************************** + * UintIndexedMT * + ************************** + */ + + struct UintIndexedMT { + IndexedMT _indexedMT; + } + + /** + * @notice Initialize the in-storage Indexed Merkle tree wrapper for uint values. + * + * Requirements: + * - The tree must not already be initialized. + * + * @param tree self. + */ + function initialize(UintIndexedMT storage tree) internal { + _initialize(tree._indexedMT); + } + + /** + * @notice Set custom hashing functions to be used by the tree. + * + * Requirements: + * - Must be called before the tree is initialized. + * + * @param tree self. + * @param hashFunctions_ The hash function container (hash2 and hash4). + */ + function setHashers(UintIndexedMT storage tree, HashFunctions memory hashFunctions_) internal { + _setHashers(tree._indexedMT, hashFunctions_); + } + + /** + * @notice Add a new uint value to the Indexed Merkle tree. + * + * Complexity: O(log(levels)) where levels is the current tree height. + * + * @param tree self. + * @param value_ The value to insert. + * @param lowLeafIndex_ A known low leaf index indicating insertion position. + * @return The new leaf index for the inserted value. + */ + function add( + UintIndexedMT storage tree, + uint256 value_, + uint256 lowLeafIndex_ + ) internal returns (uint256) { + return _add(tree._indexedMT, bytes32(value_), lowLeafIndex_); + } + + /** + * @notice Update an existing leaf in the Indexed Merkle tree. + * + * Requirements: + * - leafIndex_ must be valid and initialized. + * + * @param tree self. + * @param leafIndex_ The index of the leaf to update. + * @param currentLowLeafIndex_ The current low-leaf insertion point that precedes the leaf. + * @param newValue_ New value to set. + * @param newLowLeafIndex_ New low-leaf pointer (may be same as currentLowLeafIndex_). + */ + function update( + UintIndexedMT storage tree, + uint256 leafIndex_, + uint256 currentLowLeafIndex_, + uint256 newValue_, + uint256 newLowLeafIndex_ + ) internal { + _update( + tree._indexedMT, + leafIndex_, + currentLowLeafIndex_, + bytes32(newValue_), + newLowLeafIndex_ + ); + } + + /** + * @notice Generate an inclusion/exclusion proof for the leaf at `index_`. + * + * Complexity: O(levels) to build the proof. + * + * @param tree self. + * @param index_ The leaf index to build the proof for. + * @param value_ The value expected at the index (used to validate existence vs exclusion). + * @return A merkle Proof structure for `index_` and `value_`. + */ + function getProof( + UintIndexedMT storage tree, + uint256 index_, + uint256 value_ + ) internal view returns (Proof memory) { + return _proof(tree._indexedMT, index_, bytes32(value_)); + } + + /** + * @notice Verify a proof produced by `getProof` for this tree instance. + * + * @param tree self. + * @param proof_ The proof to verify. + * @return True if the proof matches the current root, false otherwise. + */ + function verifyProof( + UintIndexedMT storage tree, + Proof memory proof_ + ) internal view returns (bool) { + return _verifyProof(tree._indexedMT, proof_); + } + + /** + * @notice Convenience helper to process a raw Proof using the default hashers (keccak256). + * + * @param proof_ A proof as returned by `getProof`. + * @return The computed root when hashing the proof values using default hash functions. + */ + function processProof(Proof memory proof_) internal view returns (bytes32) { + return _processProof(proof_, HashFunctions({hash2: _hash2, hash4: _hash4})); + } + + /** + * @notice Process a proof using the provided hash functions. + * + * @param proof_ A proof as returned by `getProof`. + * @param hashFunctions_ Custom hashing functions to use when processing the proof. + * @return The computed root when hashing the proof using the provided functions. + */ + function processProof( + Proof memory proof_, + HashFunctions memory hashFunctions_ + ) internal view returns (bytes32) { + return _processProof(proof_, hashFunctions_); + } + + /** + * @notice Get the current Merkle root for the tree. + * + * @param tree self. + * @return The bytes32 root hash. + */ + function getRoot(UintIndexedMT storage tree) internal view returns (bytes32) { + return _getRoot(tree._indexedMT); + } + + /** + * @notice Get the current number of levels in the Indexed Merkle tree. + * + * @param tree self. + * @return The number of levels used by the tree (>= 1 when initialized). + */ + function getTreeLevels(UintIndexedMT storage tree) internal view returns (uint256) { + return _getTreeLevels(tree._indexedMT); + } + + /** + * @notice Read data for a leaf in the tree. + * + * @param tree self. + * @param leafIndex_ The leaf index to query. + * @return LeafData struct with value and nextLeafIndex. + */ + function getLeafData( + UintIndexedMT storage tree, + uint256 leafIndex_ + ) internal view returns (LeafData memory) { + return _getLeafData(tree._indexedMT, leafIndex_); + } + + /** + * @notice Get the hash of a node at a given index and level. + * + * @param tree self. + * @param index_ Index of the node on the provided level. + * @param level_ Level to query (0 == leaves). + * @return The node hash. + */ + function getNodeHash( + UintIndexedMT storage tree, + uint256 index_, + uint256 level_ + ) internal view returns (bytes32) { + return _getNodeHash(tree._indexedMT, index_, level_); + } + + /** + * @notice Get the total number of leaves in the tree. + * + * @param tree self. + * @return The number of leaves stored at level 0. + */ + function getLeavesCount(UintIndexedMT storage tree) internal view returns (uint256) { + return _getLeavesCount(tree._indexedMT); + } + + /** + * @notice Get the number of nodes present at a specific level of the tree. + * + * @param tree self. + * @param level_ The level to query. + * @return The number of nodes in the specified level. + */ + function getLevelNodesCount( + UintIndexedMT storage tree, + uint256 level_ + ) internal view returns (uint256) { + return _getLevelNodesCount(tree._indexedMT, level_); + } + + /** + * @notice Returns true when custom hash functions were provided before initialization. + * + * @param tree self. + * @return True if custom hashers are set, false otherwise. + */ + function isCustomHasherSet(UintIndexedMT storage tree) internal view returns (bool) { + return _isCustomHasherSet(tree._indexedMT); + } + + /** + ************************** + * Bytes32IndexedMT * + ************************** + */ + + struct Bytes32IndexedMT { + IndexedMT _indexedMT; + } + + /** + * @notice Initialize the in-storage Indexed Merkle tree wrapper for bytes32 values. + * + * Requirements: + * - The tree must not already be initialized. + * + * @param tree self. + */ + function initialize(Bytes32IndexedMT storage tree) internal { + _initialize(tree._indexedMT); + } + + /** + * @notice Set custom hashing functions to be used by the tree. + * + * Requirements: + * - Must be called before the tree is initialized. + * + * @param tree self. + * @param hashFunctions_ The hash function container (hash2 and hash4). + */ + function setHashers( + Bytes32IndexedMT storage tree, + HashFunctions memory hashFunctions_ + ) internal { + _setHashers(tree._indexedMT, hashFunctions_); + } + + /** + * @notice Add a new bytes32 value to the Indexed Merkle tree. + * + * Complexity: O(log(levels)) where levels is the current tree height. + * + * @param tree self. + * @param value_ The value to insert. + * @param lowLeafIndex_ A known low leaf index indicating insertion position. + * @return The new leaf index for the inserted value. + */ + function add( + Bytes32IndexedMT storage tree, + bytes32 value_, + uint256 lowLeafIndex_ + ) internal returns (uint256) { + return _add(tree._indexedMT, value_, lowLeafIndex_); + } + + /** + * @notice Update an existing leaf in the Indexed Merkle tree. + * + * Requirements: + * - leafIndex_ must be valid and initialized. + * + * @param tree self. + * @param leafIndex_ The index of the leaf to update. + * @param currentLowLeafIndex_ The current low-leaf insertion point that precedes the leaf. + * @param newValue_ New value to set. + * @param newLowLeafIndex_ New low-leaf pointer (may be same as currentLowLeafIndex_). + */ + function update( + Bytes32IndexedMT storage tree, + uint256 leafIndex_, + uint256 currentLowLeafIndex_, + bytes32 newValue_, + uint256 newLowLeafIndex_ + ) internal { + _update(tree._indexedMT, leafIndex_, currentLowLeafIndex_, newValue_, newLowLeafIndex_); + } + + /** + * @notice Generate an inclusion/exclusion proof for the leaf at `index_`. + * + * Complexity: O(levels) to build the proof. + * + * @param tree self. + * @param index_ The leaf index to build the proof for. + * @param value_ The value expected at the index (used to validate existence vs exclusion). + * @return A merkle Proof structure for `index_` and `value_`. + */ + function getProof( + Bytes32IndexedMT storage tree, + uint256 index_, + bytes32 value_ + ) internal view returns (Proof memory) { + return _proof(tree._indexedMT, index_, value_); + } + + /** + * @notice Verify a proof produced by `getProof` for this tree instance. + * + * @param tree self. + * @param proof_ The proof to verify. + * @return True if the proof matches the current root, false otherwise. + */ + function verifyProof( + Bytes32IndexedMT storage tree, + Proof memory proof_ + ) internal view returns (bool) { + return _verifyProof(tree._indexedMT, proof_); + } + + /** + * @notice Get the current Merkle root for the tree. + * + * @param tree self. + * @return The bytes32 root hash. + */ + function getRoot(Bytes32IndexedMT storage tree) internal view returns (bytes32) { + return _getRoot(tree._indexedMT); + } + + /** + * @notice Get the current number of levels in the Indexed Merkle tree. + * + * @param tree self. + * @return The number of levels used by the tree (>= 1 when initialized). + */ + function getTreeLevels(Bytes32IndexedMT storage tree) internal view returns (uint256) { + return _getTreeLevels(tree._indexedMT); + } + + /** + * @notice Read data for a leaf in the tree. + * + * @param tree self. + * @param leafIndex_ The leaf index to query. + * @return LeafData struct with value and nextLeafIndex. + */ + function getLeafData( + Bytes32IndexedMT storage tree, + uint256 leafIndex_ + ) internal view returns (LeafData memory) { + return _getLeafData(tree._indexedMT, leafIndex_); + } + + /** + * @notice Get the hash of a node at a given index and level. + * + * @param tree self. + * @param index_ Index of the node on the provided level. + * @param level_ Level to query (0 == leaves). + * @return The node hash. + */ + function getNodeHash( + Bytes32IndexedMT storage tree, + uint256 index_, + uint256 level_ + ) internal view returns (bytes32) { + return _getNodeHash(tree._indexedMT, index_, level_); + } + + /** + * @notice Get the total number of leaves in the tree. + * + * @param tree self. + * @return The number of leaves stored at level 0. + */ + function getLeavesCount(Bytes32IndexedMT storage tree) internal view returns (uint256) { + return _getLeavesCount(tree._indexedMT); + } + + /** + * @notice Get the number of nodes present at a specific level of the tree. + * + * @param tree self. + * @param level_ The level to query. + * @return The number of nodes in the specified level. + */ + function getLevelNodesCount( + Bytes32IndexedMT storage tree, + uint256 level_ + ) internal view returns (uint256) { + return _getLevelNodesCount(tree._indexedMT, level_); + } + + /** + * @notice Returns true when custom hash functions were provided before initialization. + * + * @param tree self. + * @return True if custom hashers are set, false otherwise. + */ + function isCustomHasherSet(Bytes32IndexedMT storage tree) internal view returns (bool) { + return _isCustomHasherSet(tree._indexedMT); + } + + /** + ************************** + * AddressIndexedMT * + ************************** + */ + + struct AddressIndexedMT { + IndexedMT _indexedMT; + } + + /** + * @notice Initialize the in-storage Indexed Merkle tree wrapper for address values. + * + * Requirements: + * - The tree must not already be initialized. + * + * @param tree self. + */ + function initialize(AddressIndexedMT storage tree) internal { + _initialize(tree._indexedMT); + } + + /** + * @notice Set custom hashing functions to be used by the tree. + * + * Requirements: + * - Must be called before the tree is initialized. + * + * @param tree self. + * @param hashFunctions_ The hash function container (hash2 and hash4). + */ + function setHashers( + AddressIndexedMT storage tree, + HashFunctions memory hashFunctions_ + ) internal { + _setHashers(tree._indexedMT, hashFunctions_); + } + + /** + * @notice Add a new address value to the Indexed Merkle tree. + * + * Complexity: O(log(levels)) where levels is the current tree height. + * + * @param tree self. + * @param value_ The address value to insert. + * @param lowLeafIndex_ A known low leaf index indicating insertion position. + * @return The new leaf index for the inserted value. + */ + function add( + AddressIndexedMT storage tree, + address value_, + uint256 lowLeafIndex_ + ) internal returns (uint256) { + return _add(tree._indexedMT, bytes32(uint256(uint160(value_))), lowLeafIndex_); + } + + /** + * @notice Update an existing leaf in the Indexed Merkle tree. + * + * Requirements: + * - leafIndex_ must be valid and initialized. + * + * @param tree self. + * @param leafIndex_ The index of the leaf to update. + * @param currentLowLeafIndex_ The current low-leaf insertion point that precedes the leaf. + * @param newValue_ New address value to set. + * @param newLowLeafIndex_ New low-leaf pointer (may be same as currentLowLeafIndex_). + */ + function update( + AddressIndexedMT storage tree, + uint256 leafIndex_, + uint256 currentLowLeafIndex_, + address newValue_, + uint256 newLowLeafIndex_ + ) internal { + _update( + tree._indexedMT, + leafIndex_, + currentLowLeafIndex_, + bytes32(uint256(uint160(newValue_))), + newLowLeafIndex_ + ); + } + + /** + * @notice Generate an inclusion/exclusion proof for the leaf at `index_`. + * + * Complexity: O(levels) to build the proof. + * + * @param tree self. + * @param index_ The leaf index to build the proof for. + * @param value_ The value expected at the index (used to validate existence vs exclusion). + * @return A merkle Proof structure for `index_` and `value_`. + */ + function getProof( + AddressIndexedMT storage tree, + uint256 index_, + address value_ + ) internal view returns (Proof memory) { + return _proof(tree._indexedMT, index_, bytes32(uint256(uint160(value_)))); + } + + /** + * @notice Verify a proof produced by `getProof` for this tree instance. + * + * @param tree self. + * @param proof_ The proof to verify. + * @return True if the proof matches the current root, false otherwise. + */ + function verifyProof( + AddressIndexedMT storage tree, + Proof memory proof_ + ) internal view returns (bool) { + return _verifyProof(tree._indexedMT, proof_); + } + + /** + * @notice Get the current Merkle root for the tree. + * + * @param tree self. + * @return The bytes32 root hash. + */ + function getRoot(AddressIndexedMT storage tree) internal view returns (bytes32) { + return _getRoot(tree._indexedMT); + } + + /** + * @notice Get the current number of levels in the Indexed Merkle tree. + * + * @param tree self. + * @return The number of levels used by the tree (>= 1 when initialized). + */ + function getTreeLevels(AddressIndexedMT storage tree) internal view returns (uint256) { + return _getTreeLevels(tree._indexedMT); + } + + /** + * @notice Read data for a leaf in the tree. + * + * @param tree self. + * @param leafIndex_ The leaf index to query. + * @return LeafData struct with value and nextLeafIndex. + */ + function getLeafData( + AddressIndexedMT storage tree, + uint256 leafIndex_ + ) internal view returns (LeafData memory) { + return _getLeafData(tree._indexedMT, leafIndex_); + } + + /** + * @notice Get the hash of a node at a given index and level. + * + * @param tree self. + * @param index_ Index of the node on the provided level. + * @param level_ Level to query (0 == leaves). + * @return The node hash. + */ + function getNodeHash( + AddressIndexedMT storage tree, + uint256 index_, + uint256 level_ + ) internal view returns (bytes32) { + return _getNodeHash(tree._indexedMT, index_, level_); + } + + /** + * @notice Get the total number of leaves in the tree. + * + * @param tree self. + * @return The number of leaves stored at level 0. + */ + function getLeavesCount(AddressIndexedMT storage tree) internal view returns (uint256) { + return _getLeavesCount(tree._indexedMT); + } + + /** + * @notice Get the number of nodes present at a specific level of the tree. + * + * @param tree self. + * @param level_ The level to query. + * @return The number of nodes in the specified level. + */ + function getLevelNodesCount( + AddressIndexedMT storage tree, + uint256 level_ + ) internal view returns (uint256) { + return _getLevelNodesCount(tree._indexedMT, level_); + } + + /** + * @notice Returns true when custom hash functions were provided before initialization. + * + * @param tree self. + * @return True if custom hashers are set, false otherwise. + */ + function isCustomHasherSet(AddressIndexedMT storage tree) internal view returns (bool) { + return _isCustomHasherSet(tree._indexedMT); + } + + /** + ************************** + * InnerIndexedMT * + ************************** + */ + + /** + * @notice Level index used for leaf nodes. + */ + uint256 internal constant LEAVES_LEVEL = 0; + + /** + * @notice A sentinel zero index used in linked-leaf pointers. + */ + uint64 internal constant ZERO_IDX = 0; + + /** + * @notice Zero hash representation used for empty nodes / default values. + */ + bytes32 internal constant ZERO_HASH = bytes32(0); + + /** + * @notice Core storage structure for the Indexed Merkle tree. + * + * @param leavesData compact storage of leaf metadata (value + pointer to next leaf). + * @param nodes mapping of level => array of node hashes; level 0 is leaves, top index is root. + * @param levelsCount current number of levels present in the tree (>= 1 after init). + * @param isCustomHasherSet true when caller provided custom hash functions before init. + * @param hash2 A two-input hash function used to hash node pairs. + * @param hash4 A four-input hash function used to hash leaf metadata (active flag, idx, value, nextIndex). + */ + struct IndexedMT { + LeafData[] leavesData; + mapping(uint256 level => bytes32[] nodeHashes) nodes; + uint256 levelsCount; + bool isCustomHasherSet; + function(bytes32, bytes32) view returns (bytes32) hash2; + function(bytes32, bytes32, bytes32, bytes32) view returns (bytes32) hash4; + } + + /** + * @notice Container type for custom hashing functions. + */ + struct HashFunctions { + function(bytes32, bytes32) view returns (bytes32) hash2; + function(bytes32, bytes32, bytes32, bytes32) view returns (bytes32) hash4; + } + + /** + * @notice Merkle proof returned by `getProof` and used by `verifyProof`. + * + * @param root The root hash for which this proof should verify. + * @param siblings Array of sibling hashes used to reconstruct the root from the leaf. + * @param existence Whether the supplied index/value exists (true) or this is an exclusion proof (false). + * @param index The leaf index (position) within the leaves level used to compute the proof. + * @param value The stored value for the leaf referenced by `index` (or candidate value for an exclusion check). + * @param nextLeafIndex For the indexed tree the leaf contains a pointer to the next leaf; used when hashing leaves. + */ + struct Proof { + bytes32 root; + bytes32[] siblings; + bool existence; + uint256 index; + bytes32 value; + uint256 nextLeafIndex; + } + + /** + * @notice The main leaf metadata struct + * @param value The stored bytes32 value for the leaf. + * @param nextLeafIndex Index of the next active leaf (ZERO_IDX when none). + */ + struct LeafData { + bytes32 value; + uint256 nextLeafIndex; + } + + error ZeroLeafIndex(); + error IndexOutOfBounds(uint256 index, uint256 level); + error InvalidLowLeaf(uint256 lowLeafIndex, bytes32 newValue); + error InvalidProofIndex(uint256 index, bytes32 value); + error NotANodeLevel(); + error NotALowLeafIndex(uint256 leafIndex, uint256 lowLeafIndex); + error IndexedMerkleTreeNotInitialized(); + error IndexedMerkleTreeAlreadyInitialized(); + + modifier onlyInitialized(IndexedMT storage tree) { + if (!_isInitialized(tree)) revert IndexedMerkleTreeNotInitialized(); + _; + } + + /** + * @dev This will create the empty-leaf sentinel and push the initial leaf and the + * corresponding zero-hash for the leaves level. The function reverts if the + * tree is already initialized. + */ + function _initialize(IndexedMT storage tree) private { + if (_isInitialized(tree)) revert IndexedMerkleTreeAlreadyInitialized(); + + tree.leavesData.push(LeafData({value: ZERO_HASH, nextLeafIndex: ZERO_IDX})); + tree.nodes[LEAVES_LEVEL].push(_hashLeaf(0, 0, 0, true, _getHashFunctions(tree).hash4)); + + tree.levelsCount++; + } + + /** + * @dev Set custom hash functions for the indexed tree. + * Must be invoked before initialization (otherwise the tree is already created). + */ + function _setHashers(IndexedMT storage tree, HashFunctions memory hashFunctions_) private { + if (_isInitialized(tree)) revert IndexedMerkleTreeAlreadyInitialized(); + + tree.isCustomHasherSet = true; + + tree.hash2 = hashFunctions_.hash2; + tree.hash4 = hashFunctions_.hash4; + } + + /** + * @dev Insert a new leaf into the indexed tree at the position following lowLeafIndex_. + * Performs necessary checks that lowLeafIndex_ is a valid low-leaf and then updates + * the data structures and merkle node hashes accordingly. + * + * @return The newly allocated leaf index. + */ + function _add( + IndexedMT storage tree, + bytes32 value_, + uint256 lowLeafIndex_ + ) private onlyInitialized(tree) returns (uint256) { + uint256 nextLeafIndex_ = _checkLowLeaf(tree, value_, lowLeafIndex_); + uint256 newLeafIndex_ = _getLeavesCount(tree); + + _updateNextLeafIndex(tree, lowLeafIndex_, newLeafIndex_); + + LeafData memory newLeafData_ = LeafData({value: value_, nextLeafIndex: nextLeafIndex_}); + + _pushLeaf(tree, newLeafIndex_, newLeafData_); + + return newLeafIndex_; + } + + /** + * @dev Update stored value and optionally reposition the leaf by modifying + * the nextLeafIndex links. The function validates provided low-leaf pointers + * and updates merkle node hashes for affected leaves. + */ + function _update( + IndexedMT storage tree, + uint256 leafIndex_, + uint256 currentLowLeafIndex_, + bytes32 newValue_, + uint256 newLowLeafIndex_ + ) private onlyInitialized(tree) { + require(leafIndex_ != ZERO_IDX, ZeroLeafIndex()); + require( + _getLeafNextIndex(tree, currentLowLeafIndex_) == leafIndex_, + NotALowLeafIndex(leafIndex_, currentLowLeafIndex_) + ); + + tree.leavesData[leafIndex_].value = newValue_; + + if (newLowLeafIndex_ != currentLowLeafIndex_ && newLowLeafIndex_ != leafIndex_) { + uint256 newNextLeafIndex_ = _checkLowLeaf(tree, newValue_, newLowLeafIndex_); + + _updateNextLeafIndex(tree, currentLowLeafIndex_, _getLeafNextIndex(tree, leafIndex_)); + _updateNextLeafIndex(tree, newLowLeafIndex_, leafIndex_); + _updateNextLeafIndex(tree, leafIndex_, newNextLeafIndex_); + } else { + _updateMerkleHashes(tree, leafIndex_); + } + } + + function _pushLeaf( + IndexedMT storage tree, + uint256 leafIndex_, + LeafData memory leafData_ + ) private { + tree.leavesData.push(leafData_); + + uint256 levelsCount_ = tree.levelsCount; + uint256 levelIndex_ = leafIndex_; + + HashFunctions memory hashFunctions_ = _getHashFunctions(tree); + + for (uint256 i = 0; i < levelsCount_; i++) { + bytes32 currentLevelNodeHash_; + + if (i == LEAVES_LEVEL) { + currentLevelNodeHash_ = _hashLeaf( + levelIndex_, + leafData_.value, + leafData_.nextLeafIndex, + true, + hashFunctions_.hash4 + ); + } else { + currentLevelNodeHash_ = _calculateNodeHash(tree, levelIndex_, i, hashFunctions_); + } + + if (levelIndex_ == _getLevelNodesCount(tree, i)) { + tree.nodes[i].push(currentLevelNodeHash_); + } else { + tree.nodes[i][levelIndex_] = currentLevelNodeHash_; + } + + if (i + 1 == levelsCount_ && _getLevelNodesCount(tree, i) > 1) { + levelsCount_++; + } + + levelIndex_ /= 2; + } + + tree.levelsCount = levelsCount_; + } + + function _updateNextLeafIndex( + IndexedMT storage tree, + uint256 leafIndex_, + uint256 newNextLeafIndex_ + ) private { + tree.leavesData[leafIndex_].nextLeafIndex = newNextLeafIndex_; + + _updateMerkleHashes(tree, leafIndex_); + } + + function _updateMerkleHashes(IndexedMT storage tree, uint256 leafIndex_) private { + uint256 levelsCount_ = tree.levelsCount; + uint256 levelIndex_ = leafIndex_; + + HashFunctions memory hashFunctions_ = _getHashFunctions(tree); + + for (uint256 i = 0; i < levelsCount_; i++) { + bytes32 currentLevelNodeHash_; + + if (i == LEAVES_LEVEL) { + LeafData memory leafData_ = _getLeafData(tree, levelIndex_); + + currentLevelNodeHash_ = _hashLeaf( + levelIndex_, + leafData_.value, + leafData_.nextLeafIndex, + true, + hashFunctions_.hash4 + ); + } else { + currentLevelNodeHash_ = _calculateNodeHash(tree, levelIndex_, i, hashFunctions_); + } + + tree.nodes[i][levelIndex_] = currentLevelNodeHash_; + + levelIndex_ /= 2; + } + } + + function _proof( + IndexedMT storage tree, + uint256 index_, + bytes32 value_ + ) private view returns (Proof memory) { + LeafData memory leafData_ = _getLeafData(tree, index_); + + Proof memory proof_ = Proof({ + root: _getRoot(tree), + siblings: new bytes32[](tree.levelsCount - 1), + existence: false, + index: index_, + value: leafData_.value, + nextLeafIndex: leafData_.nextLeafIndex + }); + + if (leafData_.value == value_) { + proof_.existence = true; + } else if (!_isLowLeaf(tree, value_, index_)) { + revert InvalidProofIndex(index_, value_); + } + + HashFunctions memory hashFunctions_ = _getHashFunctions(tree); + + uint256 parentIndex_ = index_; + + for (uint256 i = 0; i < proof_.siblings.length; ++i) { + uint256 currentLevelIndex_ = parentIndex_ % 2 == 0 + ? parentIndex_ + 1 + : parentIndex_ - 1; + + proof_.siblings[i] = currentLevelIndex_ < _getLevelNodesCount(tree, i) + ? tree.nodes[i][currentLevelIndex_] + : _getZeroNodeHash(i, hashFunctions_); + + parentIndex_ /= 2; + } + + return proof_; + } + + function _verifyProof( + IndexedMT storage tree, + Proof memory proof_ + ) private view returns (bool) { + return _processProof(proof_, _getHashFunctions(tree)) == _getRoot(tree); + } + + function _processProof( + Proof memory proof_, + HashFunctions memory hashFunctions_ + ) private view returns (bytes32) { + bytes32 computedHash_ = _hashLeaf( + proof_.index, + proof_.value, + proof_.nextLeafIndex, + true, + hashFunctions_.hash4 + ); + + for (uint256 i = 0; i < proof_.siblings.length; ++i) { + if ((proof_.index >> i) & 1 == 1) { + computedHash_ = _hashNode(proof_.siblings[i], computedHash_, hashFunctions_.hash2); + } else { + computedHash_ = _hashNode(computedHash_, proof_.siblings[i], hashFunctions_.hash2); + } + } + + return computedHash_; + } + + function _getRoot(IndexedMT storage tree) private view returns (bytes32) { + return tree.nodes[tree.levelsCount - 1][0]; + } + + function _getTreeLevels(IndexedMT storage tree) private view returns (uint256) { + return tree.levelsCount; + } + + function _getLeavesCount(IndexedMT storage tree) private view returns (uint256) { + return _getLevelNodesCount(tree, LEAVES_LEVEL); + } + + function _getLevelNodesCount( + IndexedMT storage tree, + uint256 level_ + ) private view returns (uint256) { + return tree.nodes[level_].length; + } + + function _getNodeHash( + IndexedMT storage tree, + uint256 index_, + uint256 level_ + ) private view returns (bytes32) { + return tree.nodes[level_][index_]; + } + + function _getLeafData( + IndexedMT storage tree, + uint256 index_ + ) private view returns (LeafData memory) { + _checkIndexExistence(tree, index_, LEAVES_LEVEL); + + return tree.leavesData[index_]; + } + + function _getLeafNextIndex( + IndexedMT storage tree, + uint256 index_ + ) private view returns (uint256) { + _checkIndexExistence(tree, index_, LEAVES_LEVEL); + + return tree.leavesData[index_].nextLeafIndex; + } + + function _calculateNodeHash( + IndexedMT storage tree, + uint256 index_, + uint256 level_, + HashFunctions memory hashFunctions_ + ) private view returns (bytes32) { + uint256 childrenLevel_ = level_ - 1; + uint256 leftChild_ = index_ * 2; + uint256 rightChild_ = index_ * 2 + 1; + + bytes32 leftChildHash_ = _getNodeHash(tree, leftChild_, childrenLevel_); + bytes32 rightChildHash_ = rightChild_ < _getLevelNodesCount(tree, childrenLevel_) + ? _getNodeHash(tree, rightChild_, childrenLevel_) + : _getZeroNodeHash(childrenLevel_, hashFunctions_); + + return _hashNode(leftChildHash_, rightChildHash_, hashFunctions_.hash2); + } + + function _checkIndexExistence( + IndexedMT storage tree, + uint256 index_, + uint256 level_ + ) private view { + if (index_ >= tree.nodes[level_].length) { + revert IndexOutOfBounds(index_, level_); + } + } + + /** + * @dev Validate whether `lowLeafIndex_` is a valid low leaf for the supplied value_. + * If valid, returns the next leaf index after the low leaf. + */ + function _checkLowLeaf( + IndexedMT storage tree, + bytes32 value_, + uint256 lowLeafIndex_ + ) private view returns (uint256) { + if (!_isLowLeaf(tree, value_, lowLeafIndex_)) { + revert InvalidLowLeaf(lowLeafIndex_, value_); + } + + return _getLeafData(tree, lowLeafIndex_).nextLeafIndex; + } + + /** + * @dev Returns true when `lowLeafIndex_` is an insertion point for `value_`. + * That means low leaf's value is < value_ and the next leaf's value (if present) + * is strictly greater than value_. + */ + function _isLowLeaf( + IndexedMT storage tree, + bytes32 value_, + uint256 lowLeafIndex_ + ) private view returns (bool) { + LeafData memory lowLeafData = _getLeafData(tree, lowLeafIndex_); + + uint256 nextLeafIndex_ = lowLeafData.nextLeafIndex; + + return + lowLeafData.value < value_ && + (nextLeafIndex_ == ZERO_IDX || _getLeafData(tree, nextLeafIndex_).value > value_); + } + + function _isInitialized(IndexedMT storage tree) private view returns (bool) { + return tree.levelsCount > 0; + } + + function _isCustomHasherSet(IndexedMT storage tree) private view returns (bool) { + return tree.isCustomHasherSet; + } + + /** + * @dev Return the canonical zero node hash for a given level. For level 0 this is + * the hash of an inactive leaf; for higher levels it is the hash of two identical + * zero child hashes. + */ + function _getZeroNodeHash( + uint256 level_, + HashFunctions memory hashFunctions_ + ) private view returns (bytes32) { + if (level_ == 0) { + return _hashLeaf(0, 0, 0, false, hashFunctions_.hash4); + } + + bytes32 prevLevelNodeHash_ = _getZeroNodeHash(level_ - 1, hashFunctions_); + + return _hashNode(prevLevelNodeHash_, prevLevelNodeHash_, hashFunctions_.hash2); + } + + function _hashNode( + bytes32 leftChildHash_, + bytes32 rightChildHash_, + function(bytes32, bytes32) view returns (bytes32) hash2_ + ) private view returns (bytes32) { + return hash2_(leftChildHash_, rightChildHash_); + } + + function _hashLeaf( + uint256 leafIndex_, + bytes32 value_, + uint256 nextLeafIndex_, + bool isActive_, + function(bytes32, bytes32, bytes32, bytes32) view returns (bytes32) hash4_ + ) private view returns (bytes32) { + return + hash4_( + bytes32(uint256(isActive_ ? 1 : 0)), + bytes32(leafIndex_), + value_, + bytes32(nextLeafIndex_) + ); + } + + function _getHashFunctions( + IndexedMT storage tree + ) private view returns (HashFunctions memory) { + return + HashFunctions({ + hash2: tree.isCustomHasherSet ? tree.hash2 : _hash2, + hash4: tree.isCustomHasherSet ? tree.hash4 : _hash4 + }); + } + + function _hash2(bytes32 a_, bytes32 b_) private pure returns (bytes32 result_) { + assembly { + mstore(0, a_) + mstore(32, b_) + + result_ := keccak256(0, 64) + } + } + + /** + * @dev The decision not to update the free memory pointer is due to the temporary nature of the hash arguments. + */ + function _hash4( + bytes32 a_, + bytes32 b_, + bytes32 c_, + bytes32 d_ + ) private pure returns (bytes32 result_) { + assembly { + let freePtr_ := mload(64) + + mstore(freePtr_, a_) + mstore(add(freePtr_, 32), b_) + mstore(add(freePtr_, 64), c_) + mstore(add(freePtr_, 96), d_) + + result_ := keccak256(freePtr_, 128) + } + } +} diff --git a/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol new file mode 100644 index 00000000..e17686ff --- /dev/null +++ b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol @@ -0,0 +1,248 @@ +// SPDX-License-Identifier: MIT +// solhint-disable +pragma solidity ^0.8.21; + +import {IndexedMerkleTree} from "../../../libs/data-structures/IndexedMerkleTree.sol"; + +library PoseidonUnit2L { + function poseidon(uint256[2] calldata) public pure returns (uint256) {} +} + +library PoseidonUnit4L { + function poseidon(uint256[4] calldata) public pure returns (uint256) {} +} + +contract IndexedMerkleTreeMock { + using IndexedMerkleTree for *; + + IndexedMerkleTree.UintIndexedMT internal _uintTree; + IndexedMerkleTree.Bytes32IndexedMT internal _bytes32Tree; + IndexedMerkleTree.AddressIndexedMT internal _addressTree; + + function initializeUintTree() external { + _uintTree.initialize(); + } + + function initializeBytes32Tree() external { + _bytes32Tree.initialize(); + } + + function initializeAddressTree() external { + _addressTree.initialize(); + } + + function setUintPoseidonHasher() external { + _uintTree.setHashers(IndexedMerkleTree.HashFunctions({hash2: _hash2, hash4: _hash4})); + } + + function setBytes32PoseidonHasher() external { + _bytes32Tree.setHashers(IndexedMerkleTree.HashFunctions({hash2: _hash2, hash4: _hash4})); + } + + function setAddressPoseidonHasher() external { + _addressTree.setHashers(IndexedMerkleTree.HashFunctions({hash2: _hash2, hash4: _hash4})); + } + + function addUint(uint256 value_, uint256 lowLeafIndex_) external returns (uint256) { + return _uintTree.add(value_, lowLeafIndex_); + } + + function addBytes32(bytes32 value_, uint256 lowLeafIndex_) external returns (uint256) { + return _bytes32Tree.add(value_, lowLeafIndex_); + } + + function addAddress(address value_, uint256 lowLeafIndex_) external returns (uint256) { + return _addressTree.add(value_, lowLeafIndex_); + } + + function updateUint( + uint256 indexToUpdate_, + uint256 currentLowLeafIndex_, + uint256 newValue_, + uint256 newLowLeafIndex_ + ) external { + _uintTree.update(indexToUpdate_, currentLowLeafIndex_, newValue_, newLowLeafIndex_); + } + + function updateBytes32( + uint256 indexToUpdate_, + uint256 currentLowLeafIndex_, + bytes32 newValue_, + uint256 newLowLeafIndex_ + ) external { + _bytes32Tree.update(indexToUpdate_, currentLowLeafIndex_, newValue_, newLowLeafIndex_); + } + + function updateAddress( + uint256 indexToUpdate_, + uint256 currentLowLeafIndex_, + address newValue_, + uint256 newLowLeafIndex_ + ) external { + _addressTree.update(indexToUpdate_, currentLowLeafIndex_, newValue_, newLowLeafIndex_); + } + + function getProofUint( + uint256 index_, + uint256 value_ + ) external view returns (IndexedMerkleTree.Proof memory) { + return _uintTree.getProof(index_, value_); + } + + function getProofBytes32( + uint256 index_, + bytes32 value_ + ) external view returns (IndexedMerkleTree.Proof memory) { + return _bytes32Tree.getProof(index_, value_); + } + + function getProofAddress( + uint256 index_, + address value_ + ) external view returns (IndexedMerkleTree.Proof memory) { + return _addressTree.getProof(index_, value_); + } + + function verifyProofUint(IndexedMerkleTree.Proof memory proof_) external view returns (bool) { + return _uintTree.verifyProof(proof_); + } + + function verifyProofBytes32( + IndexedMerkleTree.Proof memory proof_ + ) external view returns (bool) { + return _bytes32Tree.verifyProof(proof_); + } + + function verifyProofAddress( + IndexedMerkleTree.Proof memory proof_ + ) external view returns (bool) { + return _addressTree.verifyProof(proof_); + } + + function processProof(IndexedMerkleTree.Proof memory proof_) external view returns (bytes32) { + return IndexedMerkleTree.processProof(proof_); + } + + function processProofPoseidon( + IndexedMerkleTree.Proof memory proof_ + ) external view returns (bytes32) { + return + IndexedMerkleTree.processProof( + proof_, + IndexedMerkleTree.HashFunctions({hash2: _hash2, hash4: _hash4}) + ); + } + + function getRootUint() external view returns (bytes32) { + return _uintTree.getRoot(); + } + + function getRootBytes32() external view returns (bytes32) { + return _bytes32Tree.getRoot(); + } + + function getRootAddress() external view returns (bytes32) { + return _addressTree.getRoot(); + } + + function getTreeLevelsUint() external view returns (uint256) { + return _uintTree.getTreeLevels(); + } + + function getTreeLevelsBytes32() external view returns (uint256) { + return _bytes32Tree.getTreeLevels(); + } + + function getTreeLevelsAddress() external view returns (uint256) { + return _addressTree.getTreeLevels(); + } + + function getLeafDataUint( + uint256 leafIndex_ + ) external view returns (IndexedMerkleTree.LeafData memory) { + return _uintTree.getLeafData(leafIndex_); + } + + function getLeafDataBytes32( + uint256 leafIndex_ + ) external view returns (IndexedMerkleTree.LeafData memory) { + return _bytes32Tree.getLeafData(leafIndex_); + } + + function getLeafDataAddress( + uint256 leafIndex_ + ) external view returns (IndexedMerkleTree.LeafData memory) { + return _addressTree.getLeafData(leafIndex_); + } + + function getNodeHashUint(uint256 index_, uint256 level_) external view returns (bytes32) { + return _uintTree.getNodeHash(index_, level_); + } + + function getNodeHashBytes32(uint256 index_, uint256 level_) external view returns (bytes32) { + return _bytes32Tree.getNodeHash(index_, level_); + } + + function getNodeHashAddress(uint256 index_, uint256 level_) external view returns (bytes32) { + return _addressTree.getNodeHash(index_, level_); + } + + function getLeavesCountUint() external view returns (uint256) { + return _uintTree.getLeavesCount(); + } + + function getLeavesCountBytes32() external view returns (uint256) { + return _bytes32Tree.getLeavesCount(); + } + + function getLeavesCountAddress() external view returns (uint256) { + return _addressTree.getLeavesCount(); + } + + function getLevelNodesCountUint(uint256 level_) external view returns (uint256) { + return _uintTree.getLevelNodesCount(level_); + } + + function getLevelNodesCountBytes32(uint256 level_) external view returns (uint256) { + return _bytes32Tree.getLevelNodesCount(level_); + } + + function getLevelNodesCountAddress(uint256 level_) external view returns (uint256) { + return _addressTree.getLevelNodesCount(level_); + } + + function isCustomHasherSetUint() external view returns (bool) { + return _uintTree.isCustomHasherSet(); + } + + function isCustomHasherSetBytes32() external view returns (bool) { + return _bytes32Tree.isCustomHasherSet(); + } + + function isCustomHasherSetAddress() external view returns (bool) { + return _addressTree.isCustomHasherSet(); + } + + function _hash2(bytes32 element1_, bytes32 element2_) internal pure returns (bytes32) { + return bytes32(PoseidonUnit2L.poseidon([uint256(element1_), uint256(element2_)])); + } + + function _hash4( + bytes32 element1_, + bytes32 element2_, + bytes32 element3_, + bytes32 element4_ + ) internal pure returns (bytes32) { + return + bytes32( + PoseidonUnit4L.poseidon( + [ + uint256(element1_), + uint256(element2_), + uint256(element3_), + uint256(element4_) + ] + ) + ); + } +} diff --git a/test/helpers/indexed-merkle-tree.ts b/test/helpers/indexed-merkle-tree.ts new file mode 100644 index 00000000..34ad6ccd --- /dev/null +++ b/test/helpers/indexed-merkle-tree.ts @@ -0,0 +1,479 @@ +import { ethers } from "ethers"; + +import { poseidonHash } from "./poseidon-hash.ts"; + +export interface MerkleTreeLevel { + [index: string]: string; +} + +export interface MerkleTreeLevels { + [level: string]: MerkleTreeLevel; +} + +export interface LeavesData { + [index: string]: LeafData; +} + +export interface IndexedLeafData { + index: bigint; + value: string; + nextIndex: bigint; + isActive: boolean; +} + +export interface LeafData { + value: string; + nextLeafIndex: bigint; +} + +export interface Proof { + root: string; + siblings: string[]; + existence: boolean; + index: bigint; + value: string; + nextLeafIndex: bigint; +} + +export const LEAVES_LEVEL = 0n; +export const ZERO_IDX = 0n; + +export function hashNodePoseidon(leftChild: string, rightChild: string): string { + const encodedData = ethers.AbiCoder.defaultAbiCoder().encode(["bytes32", "uint256"], [leftChild, rightChild]); + + return poseidonHash(encodedData); +} + +export function hashIndexedLeafPoseidon(leafData: IndexedLeafData): string { + const encodedData = ethers.AbiCoder.defaultAbiCoder().encode( + ["bool", "uint256", "bytes32", "uint256"], + [leafData.isActive, leafData.index, leafData.value, leafData.nextIndex], + ); + + return poseidonHash(encodedData); +} + +export function hashNode(leftChild: string, rightChild: string): string { + const encodedData = ethers.AbiCoder.defaultAbiCoder().encode(["bytes32", "uint256"], [leftChild, rightChild]); + + return ethers.keccak256(encodedData); +} + +export function hashIndexedLeaf(leafData: IndexedLeafData): string { + const encodedData = ethers.AbiCoder.defaultAbiCoder().encode( + ["bool", "uint256", "bytes32", "uint256"], + [leafData.isActive, leafData.index, leafData.value, leafData.nextIndex], + ); + + return ethers.keccak256(encodedData); +} + +export function encodeBytes32Value(value: bigint): string { + return ethers.toBeHex(value, 32); +} + +export class IndexedMerkleTree { + public zeroHashesCache: string[] = []; + + private levels: MerkleTreeLevels = {}; + private leavesData: LeavesData = {}; + private levelsCount: number; + private maxLevelsCount: number; + + private hashNodeFn: (leftChild: string, rightChild: string) => string; + private hashLeafFn: (leafData: IndexedLeafData) => string; + + public static buildMerkleTree( + leavesData?: IndexedLeafData[], + hashNodeFn: (leftChild: string, rightChild: string) => string = hashNode, + hashLeafFn: (leafData: IndexedLeafData) => string = hashIndexedLeaf, + maxLevelsCount: number = 256, + ): IndexedMerkleTree { + return new IndexedMerkleTree( + leavesData ?? [ + { + index: 0n, + value: encodeBytes32Value(0n), + nextIndex: 0n, + isActive: true, + }, + ], + maxLevelsCount, + hashNodeFn, + hashLeafFn, + ); + } + + private constructor( + leavesData: IndexedLeafData[], + maxLevelsCount: number, + hashNodeFn: (leftChild: string, rightChild: string) => string, + hashLeafFn: (leafData: IndexedLeafData) => string, + ) { + this.hashNodeFn = hashNodeFn; + this.hashLeafFn = hashLeafFn; + + if (leavesData.length === 0) { + throw new Error("Tree must have leaves."); + } + + this._precalculateZeroHashes(maxLevelsCount); + + this.levelsCount = Math.ceil(Math.log2(leavesData.length)) + 1; + this.maxLevelsCount = maxLevelsCount; + + for (let i = 0n; i < BigInt(maxLevelsCount); i++) { + this.levels[i.toString()] = {}; + } + + if (this.levelsCount > this.maxLevelsCount) { + throw new Error(`Invalid maxLevelsCount ${maxLevelsCount} parameter.`); + } + + this._buildTree(leavesData); + } + + public add(value: string, lowLeafIndex: bigint = this.getLowLeafIndex(value)): bigint { + const currentLeavesCount = this.getLevelNodesCount(LEAVES_LEVEL); + + if (currentLeavesCount >= 1n << BigInt(this.maxLevelsCount)) { + throw new Error("Maximum tree capacity reached."); + } + + if (!this._isLowLeaf(lowLeafIndex, value)) { + throw new Error(`Index ${lowLeafIndex} not a low leaf index for the value ${value}`); + } + + const newLeafIndex = currentLeavesCount; + const nextLeafIndex = this.getLeafData(lowLeafIndex).nextLeafIndex; + + this._updateNextLeafIndex(lowLeafIndex, newLeafIndex); + this._pushLeaf(newLeafIndex, { value: value, nextLeafIndex: nextLeafIndex }); + + return newLeafIndex; + } + + public update( + indexToUpdate: bigint, + currentLowLeafIndex: bigint, + newValue: string, + newLowLeafIndex: bigint = this.getLowLeafIndex(newValue), + ) { + if (indexToUpdate == ZERO_IDX) { + throw new Error("Unable to update zero index."); + } + + if (this.getLeafData(currentLowLeafIndex).nextLeafIndex != indexToUpdate) { + throw new Error(`Index ${currentLowLeafIndex} not a low leaf for the element with index ${indexToUpdate}`); + } + + this.leavesData[indexToUpdate.toString()].value = newValue; + + if (newLowLeafIndex != currentLowLeafIndex && newLowLeafIndex != indexToUpdate) { + if (!this._isLowLeaf(newLowLeafIndex, newValue)) { + throw new Error(`Index ${newLowLeafIndex} not a low leaf index for the value ${newValue}`); + } + + const nextLeafIndex = this.getLeafData(newLowLeafIndex).nextLeafIndex; + + this._updateNextLeafIndex(currentLowLeafIndex, this.getLeafData(indexToUpdate).nextLeafIndex); + this._updateNextLeafIndex(newLowLeafIndex, indexToUpdate); + this._updateNextLeafIndex(indexToUpdate, nextLeafIndex); + } else { + this._updateMerkleHashes(indexToUpdate); + } + } + + public getProof(index: bigint, value: string): Proof { + if (index >= this.getLevelNodesCount(LEAVES_LEVEL)) { + throw new Error(`Leaf with index ${index} does not exist.`); + } + + const siblings: string[] = []; + const leafData: LeafData = this.leavesData[index.toString()]; + + let leafExists: boolean; + + if (leafData.value == value) { + leafExists = true; + } else if (this._isLowLeaf(index, value)) { + leafExists = false; + } else { + throw new Error(`Invalid index ${index} for the value ${value}`); + } + + let currentIndex = index; + + for (let level = 0n; level < this.levelsCount - 1; level++) { + const isRightChild = currentIndex % 2n !== 0n; + const siblingIndex = isRightChild ? currentIndex - 1n : currentIndex + 1n; + + let siblingHash: string; + + if (siblingIndex < this.getLevelNodesCount(level)) { + siblingHash = this.levels[level.toString()][siblingIndex.toString()]; + } else { + siblingHash = this.zeroHashesCache[Number(level)]; + } + + siblings.push(siblingHash); + + currentIndex = currentIndex / 2n; + } + + return { + root: this.getRoot(), + siblings: siblings, + existence: leafExists, + index: index, + value: leafData.value, + nextLeafIndex: leafData.nextLeafIndex, + }; + } + + public verifyProof(proof: Proof): boolean { + return this.processProof(proof) == this.getRoot(); + } + + public processProof(proof: Proof): string { + let computedHash = this.hashLeafFn({ + index: proof.index, + nextIndex: proof.nextLeafIndex, + value: proof.value, + isActive: true, + }); + + for (let i = 0; i < proof.siblings.length; ++i) { + if (((proof.index >> BigInt(i)) & 1n) === 1n) { + computedHash = this.hashNodeFn(proof.siblings[i], computedHash); + } else { + computedHash = this.hashNodeFn(computedHash, proof.siblings[i]); + } + } + + return computedHash; + } + + public getLeafData(index: bigint): LeafData { + if (index >= this.getLevelNodesCount(LEAVES_LEVEL)) { + throw new Error(`Leaf with index ${index} does not exist.`); + } + + return this.leavesData[index.toString()]; + } + + public getPrevLeafIndex(index: bigint): bigint { + const leavesCount = this.getLevelNodesCount(LEAVES_LEVEL); + let currentIndex = ZERO_IDX; + + for (let i = 0n; i < leavesCount; i++) { + const currentLeafData = this.getLeafData(currentIndex); + if (currentLeafData.nextLeafIndex == index) { + return currentIndex; + } + + currentIndex = currentLeafData.nextLeafIndex; + } + + throw new Error(`Can't find a previous leaf for the leaf with index ${index}`); + } + + public getLeafIndex(value: string): bigint { + const leavesCount = this.getLevelNodesCount(LEAVES_LEVEL); + + for (let i = 0n; i < leavesCount; i++) { + if (this._cmpValues(this.leavesData[i.toString()].value, value) == 0) { + return i; + } + } + + throw new Error(`Can't find a leaf with value ${value}`); + } + + public getLowLeafIndex(value: string): bigint { + const leavesCount = this.getLevelNodesCount(LEAVES_LEVEL); + + for (let i = 0n; i < leavesCount; i++) { + if (this._isLowLeaf(i, value)) { + return i; + } + } + + throw new Error("Can't find a low leaf index"); + } + + public getRoot(): string { + return this.levels[this.levelsCount - 1][0]; + } + + public getLevelsCount(): number { + return this.levelsCount; + } + + public getLevelHashes(level: bigint): string[] { + return Object.values(this.levels[level.toString()]) || []; + } + + public getLevelNodesCount(level: bigint): bigint { + return BigInt(Object.values(this.levels[level.toString()]).length); + } + + public getLeavesCount(): bigint { + return this.getLevelNodesCount(LEAVES_LEVEL); + } + + private _updateNextLeafIndex(leafIndex: bigint, newNextLeafIndex: bigint): void { + this.leavesData[leafIndex.toString()].nextLeafIndex = newNextLeafIndex; + + this._updateMerkleHashes(leafIndex); + } + + private _updateMerkleHashes(leafIndex: bigint): void { + let levelIndex: bigint = leafIndex; + + for (let level = 0n; level < this.levelsCount; level++) { + let currentLevelNodeHash: string; + + if (level == LEAVES_LEVEL) { + const leafData = this.getLeafData(levelIndex); + + currentLevelNodeHash = this.hashLeafFn({ + index: levelIndex, + value: leafData.value, + nextIndex: leafData.nextLeafIndex, + isActive: true, + }); + } else { + currentLevelNodeHash = this._calculateNodeHash(levelIndex, level); + } + + this.levels[level.toString()][levelIndex.toString()] = currentLevelNodeHash; + + levelIndex /= 2n; + } + } + + private _pushLeaf(leafIndex: bigint, leafData: LeafData) { + this.leavesData[leafIndex.toString()] = leafData; + + let levelIndex = leafIndex; + + for (let level = 0n; level < this.levelsCount; level++) { + let currentLevelNodeHash: string; + + if (level == LEAVES_LEVEL) { + currentLevelNodeHash = this.hashLeafFn({ + index: levelIndex, + value: leafData.value, + nextIndex: leafData.nextLeafIndex, + isActive: true, + }); + } else { + currentLevelNodeHash = this._calculateNodeHash(levelIndex, level); + } + + this.levels[level.toString()][levelIndex.toString()] = currentLevelNodeHash; + + if (level + 1n == BigInt(this.levelsCount) && this.getLevelNodesCount(level) > 1) { + this.levelsCount++; + } + + levelIndex /= 2n; + } + } + + private _precalculateZeroHashes(maxDepth: number): void { + if (this.zeroHashesCache.length > 0) return; + + let currentHash = this.hashLeafFn({ + index: ZERO_IDX, + value: encodeBytes32Value(0n), + nextIndex: ZERO_IDX, + isActive: false, + }); + this.zeroHashesCache.push(currentHash); + + for (let i = 1; i <= maxDepth; i++) { + currentHash = this.hashNodeFn(currentHash, currentHash); + this.zeroHashesCache.push(currentHash); + } + } + + private _createLeafLevel(leavesData: IndexedLeafData[]): string[] { + return leavesData.map((data, index) => { + this.leavesData[index.toString()] = { + value: data.value, + nextLeafIndex: data.nextIndex, + }; + + return this.hashLeafFn(data); + }); + } + + private _buildTree(leavesData: IndexedLeafData[]): void { + let currentLevelHashes = this._createLeafLevel(leavesData); + + currentLevelHashes.forEach((leafHash: string, index: number) => { + this.levels[LEAVES_LEVEL.toString()][index] = leafHash; + }); + + let level = 0; + + while (currentLevelHashes.length > 1) { + const nextLevelHashes: string[] = []; + + for (let i = 0; i < currentLevelHashes.length; i += 2) { + const left = currentLevelHashes[i]; + const right = i + 1 < currentLevelHashes.length ? currentLevelHashes[i + 1] : this.zeroHashesCache[level]; + + nextLevelHashes.push(this.hashNodeFn(left, right)); + } + + level++; + + nextLevelHashes.forEach((leafHash: string, index: number) => { + this.levels[level.toString()][index] = leafHash; + }); + + currentLevelHashes = nextLevelHashes; + } + } + + private _calculateNodeHash(index: bigint, level: bigint): string { + if (level == LEAVES_LEVEL) { + throw new Error("Not a leaves level"); + } + + const childrenLevel = level - 1n; + const leftChild = index * 2n; + const rightChild = index * 2n + 1n; + + const leftChildHash = this.levels[childrenLevel.toString()][leftChild.toString()]; + const rightChildHash = + rightChild < this.getLevelNodesCount(childrenLevel) + ? this.levels[childrenLevel.toString()][rightChild.toString()] + : this.zeroHashesCache[Number(childrenLevel)]; + + return this.hashNodeFn(leftChildHash, rightChildHash); + } + + private _isLowLeaf(index: bigint, value: string): boolean { + const leafData = this.getLeafData(index); + + return ( + this._cmpValues(leafData.value, value) == -1 && + (leafData.nextLeafIndex == ZERO_IDX || + this._cmpValues(this.getLeafData(leafData.nextLeafIndex).value, value) == 1) + ); + } + + private _cmpValues(value0: string, value1: string): number { + if (BigInt(value0) > BigInt(value1)) { + return 1; + } else if (BigInt(value0) < BigInt(value1)) { + return -1; + } else { + return 0; + } + } +} diff --git a/test/libs/data-structures/IndexedMerkleTree.test.ts b/test/libs/data-structures/IndexedMerkleTree.test.ts new file mode 100644 index 00000000..5af893ca --- /dev/null +++ b/test/libs/data-structures/IndexedMerkleTree.test.ts @@ -0,0 +1,1253 @@ +import { expect } from "chai"; +import hre from "hardhat"; + +import { Reverter, getPoseidon } from "@test-helpers"; + +import { IndexedMerkleTreeMock } from "@ethers-v6"; + +import { IndexedMerkleTree as IndexedMerkleTreeLib } from "../../../generated-types/ethers/mock/libs/data-structures/IndexedMerkleTreeMock.ts"; +import { + IndexedMerkleTree, + Proof, + ZERO_IDX, + encodeBytes32Value, + hashIndexedLeaf, + hashIndexedLeafPoseidon, + hashNodePoseidon, +} from "@/test/helpers/indexed-merkle-tree.ts"; + +const { ethers, networkHelpers } = await hre.network.connect(); + +describe("IndexedMerkleTree", () => { + const reverter: Reverter = new Reverter(networkHelpers); + + const LEAVES_LEVEL = 0n; + + let indexedMT: IndexedMerkleTreeMock; + + function getRandomIntInclusive(min: number, max: number): number { + min = Math.ceil(min); + max = Math.floor(max); + + return Math.floor(Math.random() * (max - min + 1)) + min; + } + + function compareProofs( + contractProof: IndexedMerkleTreeLib.ProofStructOutput, + localProof: Proof, + expectedExistence: boolean, + ) { + expect(contractProof.root).to.be.eq(localProof.root); + expect(contractProof.existence).to.be.eq(expectedExistence); + expect(contractProof.existence).to.be.eq(localProof.existence); + expect(contractProof.index).to.be.eq(localProof.index); + expect(contractProof.value).to.be.eq(localProof.value); + expect(contractProof.nextLeafIndex).to.be.eq(localProof.nextLeafIndex); + expect(contractProof.siblings).to.be.deep.eq(localProof.siblings); + } + + function encodeAddressValue(address: string): string { + return ethers.AbiCoder.defaultAbiCoder().encode(["address"], [address]); + } + + function checkInvariant(localIndexedMerkleTree: IndexedMerkleTree) { + const leavesCount = localIndexedMerkleTree.getLeavesCount(); + const usedIndexes: bigint[] = []; + let currentIndex = 0n; + + for (let i = 0n; i < leavesCount; ++i) { + const currentLeafInfo = localIndexedMerkleTree.getLeafData(currentIndex); + + if (currentLeafInfo.nextLeafIndex == ZERO_IDX && i != leavesCount - 1n) { + throw new Error(`Invariant failed: zero index in the middle`); + } + + if (usedIndexes.includes(currentLeafInfo.nextLeafIndex)) { + throw new Error(`Invariant failed: index ${currentLeafInfo.nextLeafIndex} is already used`); + } + + const nextLeafInfo = localIndexedMerkleTree.getLeafData(currentLeafInfo.nextLeafIndex); + + if (currentLeafInfo.nextLeafIndex != ZERO_IDX && currentLeafInfo.value >= nextLeafInfo.value) { + throw new Error(`Invariant failed: invalid next leaf index for the ${currentIndex} index`); + } + + currentIndex = currentLeafInfo.nextLeafIndex; + + usedIndexes.push(currentIndex); + } + } + + before("setup", async () => { + indexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + + await reverter.snapshot(); + }); + + afterEach("cleanup", async () => { + await reverter.revert(); + }); + + describe("UintIndexedMerkleTree", () => { + beforeEach("setup", async () => { + await indexedMT.initializeUintTree(); + + expect(await indexedMT.isCustomHasherSetUint()).to.be.false; + }); + + describe("initialize", () => { + it("should correctly initialize UintIndexedMerkleTree", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + const zeroLeafHash = hashIndexedLeaf({ + index: 0n, + isActive: true, + nextIndex: 0n, + value: ethers.ZeroHash, + }); + + expect(await indexedMT.getRootUint()).to.be.eq(zeroLeafHash); + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + expect(await indexedMT.getTreeLevelsUint()).to.be.eq(1); + expect(await indexedMT.getLeavesCountUint()).to.be.eq(1); + expect(await indexedMT.getNodeHashUint(0, LEAVES_LEVEL)).to.be.eq(zeroLeafHash); + }); + + it("should get exception if try to initialize twice", async () => { + await expect(indexedMT.initializeUintTree()).to.be.revertedWithCustomError( + indexedMT, + "IndexedMerkleTreeAlreadyInitialized", + ); + }); + }); + + describe("add", () => { + it("should correctly add new elements with the increment values", async () => { + const startIndex = 1n; + let lowLeafIndex = 0n; + let lowLeafValue = 0n; + let value = 10n; + + const count = 10; + + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + + for (let i = 0; i < count; ++i) { + const currentIndex = startIndex + BigInt(i); + + await indexedMT.addUint(value, lowLeafIndex); + localIndexedMerkleTree.add(encodeBytes32Value(value)); + + const leafData = await indexedMT.getLeafDataUint(currentIndex); + + expect(leafData.value).to.be.eq(value); + expect(leafData.nextLeafIndex).to.be.eq(0n); + + const leafHash = hashIndexedLeaf({ + index: currentIndex, + value: encodeBytes32Value(value), + nextIndex: 0n, + isActive: true, + }); + expect(await indexedMT.getNodeHashUint(currentIndex, LEAVES_LEVEL)).to.be.eq(leafHash); + + const lowLeafData = await indexedMT.getLeafDataUint(lowLeafIndex); + + expect(lowLeafData.value).to.be.eq(lowLeafValue); + expect(lowLeafData.nextLeafIndex).to.be.eq(currentIndex); + + const lowLeafNewHash = hashIndexedLeaf({ + index: lowLeafIndex, + value: encodeBytes32Value(lowLeafValue), + nextIndex: currentIndex, + isActive: true, + }); + expect(await indexedMT.getNodeHashUint(lowLeafIndex, LEAVES_LEVEL)).to.be.eq(lowLeafNewHash); + + lowLeafIndex = currentIndex; + lowLeafValue = value; + value *= 2n; + + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + + checkInvariant(localIndexedMerkleTree); + } + + const expectedLevelsCount = Math.ceil(Math.log2(count + 1)) + 1; + + expect(await indexedMT.getTreeLevelsUint()).to.be.eq(expectedLevelsCount); + }); + + it("should correctly add 100 random elements", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + const elementsCount = 100n; + + for (let i = 0; i < elementsCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = localIndexedMerkleTree.getLowLeafIndex(currentValue); + + const expectedNextLeafIndex = localIndexedMerkleTree.getLeafData(lowLeafIndex).nextLeafIndex; + + const index = localIndexedMerkleTree.add(currentValue, lowLeafIndex); + await indexedMT.addUint(BigInt(currentValue), lowLeafIndex); + + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + + const leafData = await indexedMT.getLeafDataUint(index); + + expect(leafData.value).to.be.eq(currentValue); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + expect((await indexedMT.getLeafDataUint(lowLeafIndex)).nextLeafIndex).to.be.eq(index); + + checkInvariant(localIndexedMerkleTree); + } + + expect(await indexedMT.getLevelNodesCountUint(LEAVES_LEVEL)).to.be.eq(elementsCount + 1n); + }); + + it("should get exception if pass invalid low leaf index", async () => { + await indexedMT.addUint(10n, 0n); + await indexedMT.addUint(20n, 1n); + + await expect(indexedMT.addUint(5n, 1n)).to.be.revertedWithCustomError(indexedMT, "InvalidLowLeaf"); + await expect(indexedMT.addUint(25n, 1n)).to.be.revertedWithCustomError(indexedMT, "InvalidLowLeaf"); + }); + + it("should get exception if the tree is not initialized", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + + await expect(newIndexedMT.addUint(10n, 0n)).to.be.revertedWithCustomError( + newIndexedMT, + "IndexedMerkleTreeNotInitialized", + ); + }); + }); + + describe("update", () => { + const values: bigint[] = [0n, 30n, 10n, 5n, 20n]; + let localIndexedMerkleTree: IndexedMerkleTree; + + beforeEach("setup", async () => { + localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + + for (let i = 1; i < values.length; i++) { + const lowLeafIndex = localIndexedMerkleTree.getLowLeafIndex(encodeBytes32Value(values[i])); + localIndexedMerkleTree.add(encodeBytes32Value(values[i])); + + await indexedMT.addUint(values[i], lowLeafIndex); + } + + checkInvariant(localIndexedMerkleTree); + }); + + it("should correctly update value without updating next leaf indexes", async () => { + const index = 1n; + const lowLeafIndex = 4n; + const newValue = 25n; + const newLowIndex = 4n; + + localIndexedMerkleTree.update(index, lowLeafIndex, encodeBytes32Value(newValue)); + await indexedMT.updateUint(index, lowLeafIndex, newValue, newLowIndex); + + const expectedNextLeafIndex = ZERO_IDX; + const leafData = await indexedMT.getLeafDataUint(index); + + expect(leafData.value).to.be.eq(newValue); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + const lowLeafData = await indexedMT.getLeafDataUint(lowLeafIndex); + + expect(lowLeafData.nextLeafIndex).to.be.eq(index); + + checkInvariant(localIndexedMerkleTree); + + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + }); + + it("should correctly update value with the same new low leaf index", async () => { + const index = 2n; + const lowLeafIndex = 3n; + const newValue = 9n; + const newLowIndex = 3n; + + const expectedNextLeafIndex = 4n; + + let leafData = await indexedMT.getLeafDataUint(index); + + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + localIndexedMerkleTree.update(index, lowLeafIndex, encodeBytes32Value(newValue)); + await indexedMT.updateUint(index, lowLeafIndex, newValue, newLowIndex); + + leafData = await indexedMT.getLeafDataUint(index); + + expect(leafData.value).to.be.eq(newValue); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + const lowLeafData = await indexedMT.getLeafDataUint(lowLeafIndex); + + expect(lowLeafData.nextLeafIndex).to.be.eq(index); + + checkInvariant(localIndexedMerkleTree); + + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + }); + + it("should correctly update value if the new low leaf index is equal to index", async () => { + const index = 2n; + const lowLeafIndex = 3n; + const newValue = 12n; + const newLowIndex = 2n; + + const expectedNextLeafIndex = 4n; + + let leafData = await indexedMT.getLeafDataUint(index); + + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + localIndexedMerkleTree.update(index, lowLeafIndex, encodeBytes32Value(newValue)); + await indexedMT.updateUint(index, lowLeafIndex, newValue, newLowIndex); + + leafData = await indexedMT.getLeafDataUint(index); + + expect(leafData.value).to.be.eq(newValue); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + const lowLeafData = await indexedMT.getLeafDataUint(lowLeafIndex); + + expect(lowLeafData.nextLeafIndex).to.be.eq(index); + + checkInvariant(localIndexedMerkleTree); + + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + }); + + it("should correctly update values in the random tree", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + + await newIndexedMT.initializeUintTree(); + + const valuesCount = 100n; + + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + + newLocalIndexedMT.add(currentValue, lowLeafIndex); + await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + + checkInvariant(newLocalIndexedMT); + } + + const updatesCount = 100n; + + for (let i = 0; i < updatesCount; i++) { + const randIndex = BigInt(getRandomIntInclusive(1, 99)); + const newValue = ethers.hexlify(ethers.randomBytes(32)); + const currentLowLeafIndex = newLocalIndexedMT.getPrevLeafIndex(randIndex); + const newLowLeafIndex = newLocalIndexedMT.getLowLeafIndex(newValue); + + newLocalIndexedMT.update(randIndex, currentLowLeafIndex, newValue, newLowLeafIndex); + await newIndexedMT.updateUint(randIndex, currentLowLeafIndex, newValue, newLowLeafIndex); + + expect(await newIndexedMT.getRootUint()).to.be.eq(newLocalIndexedMT.getRoot()); + + checkInvariant(newLocalIndexedMT); + } + }); + + it("should get exception if pass zero index", async () => { + await expect(indexedMT.updateUint(0n, 0n, 123n, 0n)).to.be.revertedWithCustomError(indexedMT, "ZeroLeafIndex"); + }); + + it("should get exception if pass invalid current low leaf index", async () => { + const index = 1n; + const invalidLowLeafIndex = 3n; + const newValue = 12n; + const newLowIndex = 2n; + + await expect(indexedMT.updateUint(index, invalidLowLeafIndex, newValue, newLowIndex)) + .to.be.revertedWithCustomError(indexedMT, "NotALowLeafIndex") + .withArgs(index, invalidLowLeafIndex); + }); + + it("should get exception if pas invalid new low index", async () => { + const index = 1n; + const lowLeafIndex = 4n; + const newValue = 12n; + const newLowIndex = 3n; + + await expect(indexedMT.updateUint(index, lowLeafIndex, newValue, newLowIndex)) + .to.be.revertedWithCustomError(indexedMT, "InvalidLowLeaf") + .withArgs(newLowIndex, newValue); + }); + }); + + describe("getProof", () => { + const values: bigint[] = [0n, 10n, 20n, 30n]; + let localIndexedMerkleTree: IndexedMerkleTree; + + beforeEach("setup", async () => { + localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + + for (let i = 1; i < values.length; i++) { + localIndexedMerkleTree.add(encodeBytes32Value(values[i])); + + await indexedMT.addUint(values[i], i - 1); + } + }); + + it("should return correct inclusion proof", async () => { + let index = 2n; + let value = 20n; + let expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); + let proof = await indexedMT.getProofUint(index, value); + + compareProofs(proof, expectedProof, true); + + index = 1n; + value = 10n; + expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); + proof = await indexedMT.getProofUint(index, value); + + compareProofs(proof, expectedProof, true); + }); + + it("should return correct inclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + + await newIndexedMT.initializeUintTree(); + + const valuesCount = 100n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + + newLocalIndexedMT.add(currentValue, lowLeafIndex); + await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + + values.push(currentValue); + + checkInvariant(newLocalIndexedMT); + } + + const proofsCount = 100n; + + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, 99); + const valueToProve = values[randIndex]; + + const index = newLocalIndexedMT.getLeafIndex(valueToProve); + + const expectedProof = newLocalIndexedMT.getProof(index, valueToProve); + const proof = await newIndexedMT.getProofUint(index, valueToProve); + + compareProofs(proof, expectedProof, true); + } + }); + + it("should return correct exclusion proof", async () => { + const index = 1n; + const value = 15n; + const expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); + const proof = await indexedMT.getProofUint(index, value); + + compareProofs(proof, expectedProof, false); + }); + + it("should return correct exclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + + await newIndexedMT.initializeUintTree(); + + const valuesCount = 100n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + + newLocalIndexedMT.add(currentValue, lowLeafIndex); + await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + + values.push(currentValue); + + checkInvariant(newLocalIndexedMT); + } + + const proofsCount = 100n; + + for (let i = 0; i < proofsCount; i++) { + let valueToProve: string; + + do { + valueToProve = ethers.hexlify(ethers.randomBytes(32)); + } while (values.includes(valueToProve)); + + const index = newLocalIndexedMT.getLowLeafIndex(valueToProve); + + const expectedProof = newLocalIndexedMT.getProof(index, valueToProve); + const proof = await newIndexedMT.getProofUint(index, valueToProve); + + compareProofs(proof, expectedProof, false); + } + }); + + it("should get exception if pass invalid index", async () => { + const index = 1n; + const value = 5n; + + await expect(indexedMT.getProofUint(index, value)) + .to.be.revertedWithCustomError(indexedMT, "InvalidProofIndex") + .withArgs(index, value); + }); + }); + + describe("verifyProof", () => { + const values: bigint[] = [0n, 10n, 20n, 30n, 40n, 50n]; + let localIndexedMerkleTree: IndexedMerkleTree; + + beforeEach("setup", async () => { + localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + + for (let i = 1; i < values.length; i++) { + const lowIndex = localIndexedMerkleTree.getLowLeafIndex(encodeBytes32Value(values[i])); + localIndexedMerkleTree.add(encodeBytes32Value(values[i])); + + await indexedMT.addUint(values[i], lowIndex); + } + }); + + it("should correctly verify inclusion proofs", async () => { + let index = 1n; + let proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + + expect(await indexedMT.verifyProofUint(proof)).to.be.true; + + index = 2n; + proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + + expect(await indexedMT.verifyProofUint(proof)).to.be.true; + + index = 5n; + proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + + expect(await indexedMT.verifyProofUint(proof)).to.be.true; + + expect(await indexedMT.processProof(proof)).to.be.eq(localIndexedMerkleTree.getRoot()); + }); + + it("should correctly verify inclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + + await newIndexedMT.initializeUintTree(); + + const valuesCount = 100n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + + newLocalIndexedMT.add(currentValue, lowLeafIndex); + await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + + values.push(currentValue); + } + + const proofsCount = 100n; + + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, 99); + const valueToProve = values[randIndex]; + + const index = newLocalIndexedMT.getLeafIndex(valueToProve); + const proof = newLocalIndexedMT.getProof(index, valueToProve); + + expect(await newIndexedMT.verifyProofUint(proof)).to.be.true; + } + }); + + it("should correctly verify exclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + + await newIndexedMT.initializeUintTree(); + + const valuesCount = 100n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + + newLocalIndexedMT.add(currentValue, lowLeafIndex); + await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + + values.push(currentValue); + } + + const proofsCount = 100n; + + for (let i = 0; i < proofsCount; i++) { + let valueToProve: string; + + do { + valueToProve = ethers.hexlify(ethers.randomBytes(32)); + } while (values.includes(valueToProve)); + + const index = newLocalIndexedMT.getLowLeafIndex(valueToProve); + const proof = newLocalIndexedMT.getProof(index, valueToProve); + + expect(await newIndexedMT.verifyProofUint(proof)).to.be.true; + } + }); + }); + + describe("getters", () => { + it("should get exception if pass invalid index", async () => { + const invalidIndex = 120n; + + await expect(indexedMT.getLeafDataUint(invalidIndex)) + .to.be.revertedWithCustomError(indexedMT, "IndexOutOfBounds") + .withArgs(invalidIndex, LEAVES_LEVEL); + }); + }); + }); + + describe("UintIndexedMerkleTree Poseidon", () => { + beforeEach("setup", async () => { + await indexedMT.setUintPoseidonHasher(); + await indexedMT.initializeUintTree(); + + expect(await indexedMT.isCustomHasherSetUint()).to.be.true; + }); + + describe("initialize", () => { + it("should correctly initialize UintIndexedMerkleTree with Poseidon hash functions", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree( + undefined, + hashNodePoseidon, + hashIndexedLeafPoseidon, + ); + const zeroLeafHash = hashIndexedLeafPoseidon({ + index: 0n, + isActive: true, + nextIndex: 0n, + value: ethers.ZeroHash, + }); + + expect(await indexedMT.getRootUint()).to.be.eq(zeroLeafHash); + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + expect(await indexedMT.getTreeLevelsUint()).to.be.eq(1); + expect(await indexedMT.getLeavesCountUint()).to.be.eq(1); + expect(await indexedMT.getNodeHashUint(0, LEAVES_LEVEL)).to.be.eq(zeroLeafHash); + }); + + it("should get exception if try to initialize twice", async () => { + await expect(indexedMT.initializeUintTree()).to.be.revertedWithCustomError( + indexedMT, + "IndexedMerkleTreeAlreadyInitialized", + ); + }); + }); + + describe("setHashers", () => { + it("should get exception if the IndexedMerkleTree is already initialized", async () => { + await expect(indexedMT.setUintPoseidonHasher()).to.revertedWithCustomError( + indexedMT, + "IndexedMerkleTreeAlreadyInitialized", + ); + }); + }); + + describe("add", () => { + it("should correctly add 15 random elements", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree( + undefined, + hashNodePoseidon, + hashIndexedLeafPoseidon, + ); + const elementsCount = 15n; + + for (let i = 0; i < elementsCount; ++i) { + const randomValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randomValue); + const lowLeafIndex = localIndexedMerkleTree.getLowLeafIndex(encodedValue); + + const expectedNextLeafIndex = localIndexedMerkleTree.getLeafData(lowLeafIndex).nextLeafIndex; + + const index = localIndexedMerkleTree.add(encodedValue, lowLeafIndex); + await indexedMT.addUint(randomValue, lowLeafIndex); + + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMerkleTree.getRoot()); + + const leafData = await indexedMT.getLeafDataUint(index); + + expect(leafData.value).to.be.eq(encodedValue); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + expect((await indexedMT.getLeafDataUint(lowLeafIndex)).nextLeafIndex).to.be.eq(index); + + checkInvariant(localIndexedMerkleTree); + } + + expect(await indexedMT.getLevelNodesCountUint(LEAVES_LEVEL)).to.be.eq(elementsCount + 1n); + }); + }); + + describe("update", () => { + it("should correctly update values in the random tree", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 15n; + + for (let i = 0; i < valuesCount; ++i) { + const randomValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randomValue); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedValue); + + localIndexedMT.add(encodedValue, lowLeafIndex); + await indexedMT.addUint(randomValue, lowLeafIndex); + + checkInvariant(localIndexedMT); + } + + const updatesCount = 10n; + + for (let i = 0; i < updatesCount; i++) { + const randIndex = BigInt(getRandomIntInclusive(1, Number(valuesCount - 1n))); + const newValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const newEncodedValue = encodeBytes32Value(newValue); + const currentLowLeafIndex = localIndexedMT.getPrevLeafIndex(randIndex); + const newLowLeafIndex = localIndexedMT.getLowLeafIndex(newEncodedValue); + + localIndexedMT.update(randIndex, currentLowLeafIndex, newEncodedValue, newLowLeafIndex); + await indexedMT.updateUint(randIndex, currentLowLeafIndex, newValue, newLowLeafIndex); + + expect(await indexedMT.getRootUint()).to.be.eq(localIndexedMT.getRoot()); + + checkInvariant(localIndexedMT); + } + }); + }); + + describe("verifyProof", () => { + it("should correctly verify inclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 15n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const randomValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randomValue); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedValue); + + localIndexedMT.add(encodedValue, lowLeafIndex); + await indexedMT.addUint(randomValue, lowLeafIndex); + + values.push(randomValue); + } + + const proofsCount = 10n; + + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, Number(valuesCount - 1n)); + const valueToProve = values[randIndex]; + const encodedValue = encodeBytes32Value(valueToProve); + + const index = localIndexedMT.getLeafIndex(encodedValue); + const proof = localIndexedMT.getProof(index, encodedValue); + + expect(await indexedMT.processProofPoseidon(proof)).to.be.eq(localIndexedMT.getRoot()); + expect(await indexedMT.verifyProofUint(proof)).to.be.true; + } + }); + }); + }); + + describe("Bytes32IndexedMerkleTree", () => { + beforeEach("setup", async () => { + await indexedMT.setBytes32PoseidonHasher(); + await indexedMT.initializeBytes32Tree(); + + expect(await indexedMT.isCustomHasherSetBytes32()).to.be.true; + }); + + describe("initialize", () => { + it("should correctly initialize Bytes32IndexedMerkleTree", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree( + undefined, + hashNodePoseidon, + hashIndexedLeafPoseidon, + ); + const zeroLeafHash = hashIndexedLeafPoseidon({ + index: 0n, + isActive: true, + nextIndex: 0n, + value: ethers.ZeroHash, + }); + + expect(await indexedMT.getRootBytes32()).to.be.eq(zeroLeafHash); + expect(await indexedMT.getRootBytes32()).to.be.eq(localIndexedMerkleTree.getRoot()); + expect(await indexedMT.getTreeLevelsBytes32()).to.be.eq(1); + expect(await indexedMT.getLeavesCountBytes32()).to.be.eq(1); + expect(await indexedMT.getNodeHashBytes32(0, LEAVES_LEVEL)).to.be.eq(zeroLeafHash); + }); + + it("should get exception if try to initialize twice", async () => { + await expect(indexedMT.initializeBytes32Tree()).to.be.revertedWithCustomError( + indexedMT, + "IndexedMerkleTreeAlreadyInitialized", + ); + }); + }); + + describe("add", () => { + it("should correctly add 10 random elements", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree( + undefined, + hashNodePoseidon, + hashIndexedLeafPoseidon, + ); + const elementsCount = 10n; + + for (let i = 0; i < elementsCount; ++i) { + const randValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randValue); + const lowLeafIndex = localIndexedMerkleTree.getLowLeafIndex(encodedValue); + + const expectedNextLeafIndex = localIndexedMerkleTree.getLeafData(lowLeafIndex).nextLeafIndex; + + const index = localIndexedMerkleTree.add(encodedValue, lowLeafIndex); + await indexedMT.addBytes32(encodedValue, lowLeafIndex); + + expect(await indexedMT.getRootBytes32()).to.be.eq(localIndexedMerkleTree.getRoot()); + + const leafData = await indexedMT.getLeafDataBytes32(index); + + expect(leafData.value).to.be.eq(encodedValue); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + expect((await indexedMT.getLeafDataBytes32(lowLeafIndex)).nextLeafIndex).to.be.eq(index); + } + + expect(await indexedMT.getLevelNodesCountBytes32(LEAVES_LEVEL)).to.be.eq(elementsCount + 1n); + }); + }); + + describe("update", () => { + it("should correctly update values in the random tree", async () => { + const indexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + await indexedMT.setBytes32PoseidonHasher(); + await indexedMT.initializeBytes32Tree(); + + const valuesCount = 10n; + + for (let i = 0; i < valuesCount; ++i) { + const randValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randValue); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedValue); + + localIndexedMT.add(encodedValue, lowLeafIndex); + await indexedMT.addBytes32(encodedValue, lowLeafIndex); + + checkInvariant(localIndexedMT); + } + + const updatesCount = 20n; + + for (let i = 0; i < updatesCount; i++) { + const randIndex = BigInt(getRandomIntInclusive(1, Number(valuesCount - 1n))); + const randValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const newEncodedValue = encodeBytes32Value(randValue); + const currentLowLeafIndex = localIndexedMT.getPrevLeafIndex(randIndex); + const newLowLeafIndex = localIndexedMT.getLowLeafIndex(newEncodedValue); + + localIndexedMT.update(randIndex, currentLowLeafIndex, newEncodedValue, newLowLeafIndex); + await indexedMT.updateBytes32(randIndex, currentLowLeafIndex, newEncodedValue, newLowLeafIndex); + + expect(await indexedMT.getRootBytes32()).to.be.eq(localIndexedMT.getRoot()); + + checkInvariant(localIndexedMT); + } + }); + }); + + describe("getProof", () => { + it("should return correct exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 10n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const randValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randValue); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedValue); + + localIndexedMT.add(encodedValue, lowLeafIndex); + await indexedMT.addBytes32(encodedValue, lowLeafIndex); + + values.push(encodedValue); + } + + const proofsCount = 10n; + + for (let i = 0; i < proofsCount; i++) { + let valueToProve: string; + + do { + valueToProve = encodeBytes32Value(BigInt(ethers.hexlify(ethers.randomBytes(30)))); + } while (values.includes(valueToProve)); + + const index = localIndexedMT.getLowLeafIndex(valueToProve); + + const expectedProof = localIndexedMT.getProof(index, valueToProve); + const proof = await indexedMT.getProofBytes32(index, valueToProve); + + compareProofs(proof, expectedProof, false); + } + }); + }); + + describe("verifyProof", () => { + it("should correctly verify inclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 10n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const randValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randValue); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedValue); + + localIndexedMT.add(encodedValue, lowLeafIndex); + await indexedMT.addBytes32(encodedValue, lowLeafIndex); + + values.push(encodedValue); + } + + const proofsCount = 10n; + + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, values.length - 1); + const valueToProve = values[randIndex]; + + const index = localIndexedMT.getLeafIndex(valueToProve); + const proof = localIndexedMT.getProof(index, valueToProve); + + expect(await indexedMT.verifyProofBytes32(proof)).to.be.true; + } + }); + + it("should correctly verify exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 10n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const randValue = BigInt(ethers.hexlify(ethers.randomBytes(30))); + const encodedValue = encodeBytes32Value(randValue); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedValue); + + localIndexedMT.add(encodedValue, lowLeafIndex); + await indexedMT.addBytes32(encodedValue, lowLeafIndex); + + values.push(encodedValue); + } + + const proofsCount = 10n; + + for (let i = 0; i < proofsCount; i++) { + let valueToProve: string; + + do { + valueToProve = encodeBytes32Value(BigInt(ethers.hexlify(ethers.randomBytes(30)))); + } while (values.includes(valueToProve)); + + const index = localIndexedMT.getLowLeafIndex(valueToProve); + const proof = localIndexedMT.getProof(index, valueToProve); + + expect(await indexedMT.verifyProofBytes32(proof)).to.be.true; + } + }); + }); + }); + + describe("AddressIndexedMerkleTree", () => { + beforeEach("setup", async () => { + await indexedMT.setAddressPoseidonHasher(); + await indexedMT.initializeAddressTree(); + + expect(await indexedMT.isCustomHasherSetAddress()).to.be.true; + }); + + describe("initialize", () => { + it("should correctly initialize AddressIndexedMerkleTree", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree( + undefined, + hashNodePoseidon, + hashIndexedLeafPoseidon, + ); + const zeroLeafHash = hashIndexedLeafPoseidon({ + index: 0n, + isActive: true, + nextIndex: 0n, + value: ethers.ZeroHash, + }); + + expect(await indexedMT.getRootAddress()).to.be.eq(zeroLeafHash); + expect(await indexedMT.getRootAddress()).to.be.eq(localIndexedMerkleTree.getRoot()); + expect(await indexedMT.getTreeLevelsAddress()).to.be.eq(1); + expect(await indexedMT.getLeavesCountAddress()).to.be.eq(1); + expect(await indexedMT.getNodeHashAddress(0, LEAVES_LEVEL)).to.be.eq(zeroLeafHash); + }); + + it("should get exception if try to initialize twice", async () => { + await expect(indexedMT.initializeAddressTree()).to.be.revertedWithCustomError( + indexedMT, + "IndexedMerkleTreeAlreadyInitialized", + ); + }); + }); + + describe("add", () => { + it("should correctly add 10 random elements", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree( + undefined, + hashNodePoseidon, + hashIndexedLeafPoseidon, + ); + const elementsCount = 10n; + + for (let i = 0; i < elementsCount; ++i) { + const randomAddress = ethers.hexlify(ethers.randomBytes(20)); + const encodedAddress = encodeAddressValue(randomAddress); + const lowLeafIndex = localIndexedMerkleTree.getLowLeafIndex(encodedAddress); + + const expectedNextLeafIndex = localIndexedMerkleTree.getLeafData(lowLeafIndex).nextLeafIndex; + + const index = localIndexedMerkleTree.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); + + expect(await indexedMT.getRootAddress()).to.be.eq(localIndexedMerkleTree.getRoot()); + + const leafData = await indexedMT.getLeafDataAddress(index); + + expect(leafData.value).to.be.eq(encodedAddress); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + expect((await indexedMT.getLeafDataAddress(lowLeafIndex)).nextLeafIndex).to.be.eq(index); + } + + expect(await indexedMT.getLevelNodesCountAddress(LEAVES_LEVEL)).to.be.eq(elementsCount + 1n); + }); + }); + + describe("update", () => { + it("should correctly update values in the random tree", async () => { + const indexedMT = await ethers.deployContract("IndexedMerkleTreeMock", { + libraries: { + PoseidonUnit2L: await (await getPoseidon(ethers, 2)).getAddress(), + PoseidonUnit4L: await (await getPoseidon(ethers, 4)).getAddress(), + }, + }); + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + await indexedMT.setAddressPoseidonHasher(); + await indexedMT.initializeAddressTree(); + + const valuesCount = 10n; + + for (let i = 0; i < valuesCount; ++i) { + const randomAddress = ethers.hexlify(ethers.randomBytes(20)); + const encodedAddress = encodeAddressValue(randomAddress); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedAddress); + + localIndexedMT.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); + + checkInvariant(localIndexedMT); + } + + const updatesCount = 20n; + + for (let i = 0; i < updatesCount; i++) { + const randIndex = BigInt(getRandomIntInclusive(1, Number(valuesCount - 1n))); + + const newValue = ethers.hexlify(ethers.randomBytes(20)); + const newEncodedAddress = encodeAddressValue(newValue); + + const currentLowLeafIndex = localIndexedMT.getPrevLeafIndex(randIndex); + const newLowLeafIndex = localIndexedMT.getLowLeafIndex(newEncodedAddress); + + localIndexedMT.update(randIndex, currentLowLeafIndex, newEncodedAddress, newLowLeafIndex); + await indexedMT.updateAddress(randIndex, currentLowLeafIndex, newValue, newLowLeafIndex); + + expect(await indexedMT.getRootAddress()).to.be.eq(localIndexedMT.getRoot()); + + checkInvariant(localIndexedMT); + } + }); + }); + + describe("getProof", () => { + it("should return correct exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 10n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const randomAddress = ethers.hexlify(ethers.randomBytes(20)); + const encodedAddress = encodeAddressValue(randomAddress); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedAddress); + + localIndexedMT.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); + + values.push(randomAddress); + } + + const proofsCount = 10n; + + for (let i = 0; i < proofsCount; i++) { + let addressToProve: string; + + do { + addressToProve = ethers.hexlify(ethers.randomBytes(20)); + } while (values.includes(addressToProve)); + + const encodedAddressToProve = encodeAddressValue(addressToProve); + + const index = localIndexedMT.getLowLeafIndex(encodedAddressToProve); + + const expectedProof = localIndexedMT.getProof(index, encodedAddressToProve); + const proof = await indexedMT.getProofAddress(index, addressToProve); + + compareProofs(proof, expectedProof, false); + } + }); + }); + + describe("verifyProof", () => { + it("should correctly verify inclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 10n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const randomAddress = ethers.hexlify(ethers.randomBytes(20)); + const encodedAddress = encodeAddressValue(randomAddress); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedAddress); + + localIndexedMT.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); + + values.push(randomAddress); + } + + const proofsCount = 10n; + + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, values.length - 1); + const addressToProve = values[randIndex]; + const encodedAddressToProve = encodeAddressValue(addressToProve); + + const index = localIndexedMT.getLeafIndex(encodedAddressToProve); + const proof = localIndexedMT.getProof(index, encodedAddressToProve); + + expect(await indexedMT.verifyProofAddress(proof)).to.be.true; + } + }); + + it("should correctly verify exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(undefined, hashNodePoseidon, hashIndexedLeafPoseidon); + + const valuesCount = 10n; + const values = []; + + for (let i = 0; i < valuesCount; ++i) { + const randomAddress = ethers.hexlify(ethers.randomBytes(20)); + const encodedAddress = encodeAddressValue(randomAddress); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(encodedAddress); + + localIndexedMT.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); + + values.push(randomAddress); + } + + const proofsCount = 10n; + + for (let i = 0; i < proofsCount; i++) { + let addressToProve: string; + + do { + addressToProve = ethers.hexlify(ethers.randomBytes(20)); + } while (values.includes(addressToProve)); + + const encodedAddressToProve = encodeAddressValue(addressToProve); + + const index = localIndexedMT.getLowLeafIndex(encodedAddressToProve); + const proof = localIndexedMT.getProof(index, encodedAddressToProve); + + expect(await indexedMT.verifyProofAddress(proof)).to.be.true; + } + }); + }); + }); +});