From 89f64b7385ea25d0f6de979210d3e4d0e091f9fc Mon Sep 17 00:00:00 2001 From: Oleh Komendant Date: Fri, 10 Oct 2025 17:34:19 +0300 Subject: [PATCH 1/6] Draft IndexedMerkleTree --- .../data-structures/IndexedMerkleTree.sol | 163 ++++++++++++++++++ 1 file changed, 163 insertions(+) create mode 100644 contracts/libs/data-structures/IndexedMerkleTree.sol diff --git a/contracts/libs/data-structures/IndexedMerkleTree.sol b/contracts/libs/data-structures/IndexedMerkleTree.sol new file mode 100644 index 00000000..7baa4995 --- /dev/null +++ b/contracts/libs/data-structures/IndexedMerkleTree.sol @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.21; + +library IndexedMerkleTree { + uint256 internal constant LEAVES_LEVEL = 0; + uint64 internal constant ZERO_IDX = 0; + + bytes32 internal constant ZERO_HASH = bytes32(0); + + struct IndexedMT { + LeafNode[] leaves; + mapping(uint256 => bytes32[]) middleNodes; + uint256 levelsCount; + } + + struct LeafNode { + bytes32 nodeHash; + bytes32 value; + uint256 nextLeafIndex; + } + + error IndexOutOfBounds(uint256 index, uint256 level); + error InvalidLowLeaf(uint256 lowLeafIndex, bytes32 newValue); + + function _init(IndexedMT storage tree) private { + tree.leaves.push( + LeafNode({ + nodeHash: _getZeroNodeHash(LEAVES_LEVEL), + value: ZERO_HASH, + nextLeafIndex: ZERO_IDX + }) + ); + tree.levelsCount++; + } + + function _add(IndexedMT storage tree, bytes32 value_, uint256 lowLeafIndex_) private { + uint256 nextLeafIndex_ = _checkLowLeaf(tree, value_, lowLeafIndex_); + uint256 newLeafIndex_ = tree.leaves.length; + + LeafNode memory newLeafNode_ = LeafNode({ + nodeHash: _hashLeaf(newLeafIndex_, value_, nextLeafIndex_, true), + value: value_, + nextLeafIndex: nextLeafIndex_ + }); + + tree.leaves[lowLeafIndex_].nextLeafIndex = newLeafIndex_; + _updateMerkleHashes(tree, lowLeafIndex_); + + tree.leaves.push(newLeafNode_); + _updateLevels(tree); + _updateMerkleHashes(tree, newLeafIndex_); + } + + function _updateLevels(IndexedMT storage tree) private { + uint256 currentLevel_ = LEAVES_LEVEL; + uint256 currentLevelNodesCount_ = tree.leaves.length; + + while (currentLevelNodesCount_ > 1) { + if (++currentLevel_ == tree.levelsCount) { + tree.middleNodes[++tree.levelsCount].push(0); + } + + currentLevelNodesCount_ = tree.middleNodes[currentLevel_].length; + } + } + + function _updateMerkleHashes(IndexedMT storage tree, uint256 leafIndex_) private { + uint256 levelsCount_ = tree.levelsCount; + + bytes32 leftChildHash_; + bytes32 rightChildHash_; + uint256 currentParentIndex_ = leafIndex_; + + for (uint256 i = 0; i < levelsCount_; ++i) { + if (i == LEAVES_LEVEL) { + LeafNode storage leaf = tree.leaves[leafIndex_]; + + bytes32 newLeafHash_ = _hashLeaf(leafIndex_, leaf.value, leaf.nextLeafIndex, true); + leaf.nodeHash = newLeafHash_; + + bool isLeft_ = leafIndex_ % 2 == 0; + bytes32 secondLeafHash_; + + if (isLeft_ && leafIndex_ == tree.leaves.length - 1) { + secondLeafHash_ = _getZeroNodeHash(LEAVES_LEVEL); + } else if (isLeft_) { + secondLeafHash_ = tree.leaves[leafIndex_ + 1].nodeHash; + } else { + secondLeafHash_ = tree.leaves[leafIndex_ - 1].nodeHash; + } + + leftChildHash_ = isLeft_ ? newLeafHash_ : secondLeafHash_; + rightChildHash_ = isLeft_ ? secondLeafHash_ : newLeafHash_; + } else { + currentParentIndex_ >>= 1; + + bytes32[] storage currentLevel = tree.middleNodes[i]; + + bytes32 newParentHash_ = _hashNode(leftChildHash_, rightChildHash_); + + currentLevel[currentParentIndex_] = newParentHash_; + + bool isLeft_ = currentParentIndex_ % 2 == 0; + bytes32 secondNodeHash_; + + if (isLeft_ && currentParentIndex_ == currentLevel.length - 1) { + secondNodeHash_ = _getZeroNodeHash(i); + } else if (isLeft_) { + secondNodeHash_ = currentLevel[currentParentIndex_ + 1]; + } else { + secondNodeHash_ = currentLevel[currentParentIndex_ - 1]; + } + } + } + } + + function _checkLowLeaf( + IndexedMT storage tree, + bytes32 newValue_, + uint256 lowLeafIndex_ + ) private view returns (uint256 nextLeafIndex_) { + LeafNode storage lowLeaf = tree.leaves[lowLeafIndex_]; + + require( + lowLeafIndex_ > 0 && lowLeafIndex_ < tree.leaves.length, + IndexOutOfBounds(lowLeafIndex_, LEAVES_LEVEL) + ); + + nextLeafIndex_ = lowLeaf.nextLeafIndex; + + require( + lowLeaf.value < newValue_ && + (nextLeafIndex_ == ZERO_IDX || tree.leaves[nextLeafIndex_].value > newValue_), + InvalidLowLeaf(lowLeafIndex_, newValue_) + ); + } + + function _getZeroNodeHash(uint256 level_) private pure returns (bytes32) { + if (level_ == 0) { + return _hashLeaf(0, 0, 0, false); + } + + bytes32 prevLevelNodeHash_ = _getZeroNodeHash(level_ - 1); + + return _hashNode(prevLevelNodeHash_, prevLevelNodeHash_); + } + + function _hashNode( + bytes32 leftChildHash_, + bytes32 rightChildHash_ + ) private pure returns (bytes32) { + return keccak256(abi.encodePacked(leftChildHash_, rightChildHash_)); + } + + function _hashLeaf( + uint256 leafIndex_, + bytes32 value_, + uint256 nextLeafIndex_, + bool isActive_ + ) private pure returns (bytes32) { + return keccak256(abi.encodePacked(isActive_, leafIndex_, value_, nextLeafIndex_)); + } +} From 04451e1359ee93fbc61f1add6a5b9dc8df7c35eb Mon Sep 17 00:00:00 2001 From: Oleh Komendant Date: Fri, 31 Oct 2025 23:52:36 +0200 Subject: [PATCH 2/6] First IndexedMerkleTree implementation --- .../data-structures/IndexedMerkleTree.sol | 276 +++++++++++++----- .../data-structures/IndexedMerkleTreeMock.sol | 45 +++ .../data-structures/IndexedMerkleTree.test.ts | 106 +++++++ 3 files changed, 352 insertions(+), 75 deletions(-) create mode 100644 contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol create mode 100644 test/libs/data-structures/IndexedMerkleTree.test.ts diff --git a/contracts/libs/data-structures/IndexedMerkleTree.sol b/contracts/libs/data-structures/IndexedMerkleTree.sol index 7baa4995..56bcee46 100644 --- a/contracts/libs/data-structures/IndexedMerkleTree.sol +++ b/contracts/libs/data-structures/IndexedMerkleTree.sol @@ -2,115 +2,241 @@ pragma solidity ^0.8.21; library IndexedMerkleTree { + struct UintIndexedMT { + IndexedMT _indexedMT; + } + + function initialize(UintIndexedMT storage tree) internal { + _initialize(tree._indexedMT); + } + + function add( + UintIndexedMT storage tree, + uint256 value_, + uint256 lowLeafIndex_ + ) internal returns (uint256) { + return _add(tree._indexedMT, bytes32(value_), lowLeafIndex_); + } + + function getRoot(UintIndexedMT storage tree) internal view returns (bytes32) { + return _getRoot(tree._indexedMT); + } + + function getTreeLevels(UintIndexedMT storage tree) internal view returns (uint256) { + return _getTreeLevels(tree._indexedMT); + } + + function getLeafData( + UintIndexedMT storage tree, + uint256 leafIndex_ + ) internal view returns (LeafData memory) { + return _getLeafData(tree._indexedMT, leafIndex_); + } + + function getNodeHash( + UintIndexedMT storage tree, + uint256 index_, + uint256 level_ + ) internal view returns (bytes32) { + return _getNodeHash(tree._indexedMT, index_, level_); + } + + function getLeavesCount(UintIndexedMT storage tree) internal view returns (uint256) { + return _getLeavesCount(tree._indexedMT); + } + + function getLevelNodesCount( + UintIndexedMT storage tree, + uint256 level_ + ) internal view returns (uint256) { + return _getLevelNodesCount(tree._indexedMT, level_); + } + uint256 internal constant LEAVES_LEVEL = 0; uint64 internal constant ZERO_IDX = 0; bytes32 internal constant ZERO_HASH = bytes32(0); struct IndexedMT { - LeafNode[] leaves; - mapping(uint256 => bytes32[]) middleNodes; + LeafData[] leavesData; + mapping(uint256 level => bytes32[] nodeHashes) nodes; uint256 levelsCount; } - struct LeafNode { - bytes32 nodeHash; + struct LeafData { bytes32 value; uint256 nextLeafIndex; } error IndexOutOfBounds(uint256 index, uint256 level); error InvalidLowLeaf(uint256 lowLeafIndex, bytes32 newValue); + error NotANodeLevel(); + error IndexedMerkleTreeNotInitialized(); + + modifier onlyInitialized(IndexedMT storage tree) { + if (!_isInitialized(tree)) revert IndexedMerkleTreeNotInitialized(); + _; + } + + function _initialize(IndexedMT storage tree) private { + if (_isInitialized(tree)) revert IndexedMerkleTreeNotInitialized(); + + bytes32 zeroNodeHash_ = _getZeroNodeHash(LEAVES_LEVEL); + + tree.leavesData.push(LeafData({value: ZERO_HASH, nextLeafIndex: ZERO_IDX})); + tree.nodes[LEAVES_LEVEL].push(zeroNodeHash_); - function _init(IndexedMT storage tree) private { - tree.leaves.push( - LeafNode({ - nodeHash: _getZeroNodeHash(LEAVES_LEVEL), - value: ZERO_HASH, - nextLeafIndex: ZERO_IDX - }) - ); tree.levelsCount++; } - function _add(IndexedMT storage tree, bytes32 value_, uint256 lowLeafIndex_) private { + function _add( + IndexedMT storage tree, + bytes32 value_, + uint256 lowLeafIndex_ + ) private onlyInitialized(tree) returns (uint256) { uint256 nextLeafIndex_ = _checkLowLeaf(tree, value_, lowLeafIndex_); - uint256 newLeafIndex_ = tree.leaves.length; + uint256 newLeafIndex_ = _getLeavesCount(tree); - LeafNode memory newLeafNode_ = LeafNode({ - nodeHash: _hashLeaf(newLeafIndex_, value_, nextLeafIndex_, true), - value: value_, - nextLeafIndex: nextLeafIndex_ - }); - - tree.leaves[lowLeafIndex_].nextLeafIndex = newLeafIndex_; + tree.leavesData[lowLeafIndex_].nextLeafIndex = newLeafIndex_; _updateMerkleHashes(tree, lowLeafIndex_); - tree.leaves.push(newLeafNode_); - _updateLevels(tree); - _updateMerkleHashes(tree, newLeafIndex_); + LeafData memory newLeafData_ = LeafData({value: value_, nextLeafIndex: nextLeafIndex_}); + + _pushLeaf(tree, newLeafIndex_, newLeafData_); + + return newLeafIndex_; } - function _updateLevels(IndexedMT storage tree) private { - uint256 currentLevel_ = LEAVES_LEVEL; - uint256 currentLevelNodesCount_ = tree.leaves.length; + function _pushLeaf( + IndexedMT storage tree, + uint256 leafIndex_, + LeafData memory leafData_ + ) private { + tree.leavesData.push(leafData_); + + uint256 levelsCount_ = tree.levelsCount; + uint256 levelIndex_ = leafIndex_; + + for (uint256 i = 0; i < levelsCount_; i++) { + bytes32 currentLevelNodeHash_; - while (currentLevelNodesCount_ > 1) { - if (++currentLevel_ == tree.levelsCount) { - tree.middleNodes[++tree.levelsCount].push(0); + if (i == LEAVES_LEVEL) { + currentLevelNodeHash_ = _hashLeaf( + levelIndex_, + leafData_.value, + leafData_.nextLeafIndex, + true + ); + } else { + currentLevelNodeHash_ = _calculateNodeHash(tree, levelIndex_, i); + } + + if (levelIndex_ == _getLevelNodesCount(tree, i)) { + tree.nodes[i].push(currentLevelNodeHash_); + } else { + tree.nodes[i][levelIndex_] = currentLevelNodeHash_; } - currentLevelNodesCount_ = tree.middleNodes[currentLevel_].length; + if (i + 1 == levelsCount_ && _getLevelNodesCount(tree, i) > 1) { + levelsCount_++; + } + + levelIndex_ /= 2; } + + tree.levelsCount = levelsCount_; } function _updateMerkleHashes(IndexedMT storage tree, uint256 leafIndex_) private { uint256 levelsCount_ = tree.levelsCount; + uint256 levelIndex_ = leafIndex_; - bytes32 leftChildHash_; - bytes32 rightChildHash_; - uint256 currentParentIndex_ = leafIndex_; + for (uint256 i = 0; i < levelsCount_; i++) { + bytes32 currentLevelNodeHash_; - for (uint256 i = 0; i < levelsCount_; ++i) { if (i == LEAVES_LEVEL) { - LeafNode storage leaf = tree.leaves[leafIndex_]; + LeafData memory leafData_ = _getLeafData(tree, levelIndex_); + + currentLevelNodeHash_ = _hashLeaf( + levelIndex_, + leafData_.value, + leafData_.nextLeafIndex, + true + ); + } else { + currentLevelNodeHash_ = _calculateNodeHash(tree, levelIndex_, i); + } + + tree.nodes[i][levelIndex_] = currentLevelNodeHash_; - bytes32 newLeafHash_ = _hashLeaf(leafIndex_, leaf.value, leaf.nextLeafIndex, true); - leaf.nodeHash = newLeafHash_; + levelIndex_ /= 2; + } + } - bool isLeft_ = leafIndex_ % 2 == 0; - bytes32 secondLeafHash_; + function _getRoot(IndexedMT storage tree) private view returns (bytes32) { + return tree.nodes[tree.levelsCount - 1][0]; + } - if (isLeft_ && leafIndex_ == tree.leaves.length - 1) { - secondLeafHash_ = _getZeroNodeHash(LEAVES_LEVEL); - } else if (isLeft_) { - secondLeafHash_ = tree.leaves[leafIndex_ + 1].nodeHash; - } else { - secondLeafHash_ = tree.leaves[leafIndex_ - 1].nodeHash; - } + function _getTreeLevels(IndexedMT storage tree) private view returns (uint256) { + return tree.levelsCount; + } - leftChildHash_ = isLeft_ ? newLeafHash_ : secondLeafHash_; - rightChildHash_ = isLeft_ ? secondLeafHash_ : newLeafHash_; - } else { - currentParentIndex_ >>= 1; + function _getLeavesCount(IndexedMT storage tree) private view returns (uint256) { + return _getLevelNodesCount(tree, LEAVES_LEVEL); + } - bytes32[] storage currentLevel = tree.middleNodes[i]; + function _getLevelNodesCount( + IndexedMT storage tree, + uint256 level_ + ) private view returns (uint256) { + return tree.nodes[level_].length; + } - bytes32 newParentHash_ = _hashNode(leftChildHash_, rightChildHash_); + function _getNodeHash( + IndexedMT storage tree, + uint256 index_, + uint256 level_ + ) private view returns (bytes32) { + return tree.nodes[level_][index_]; + } - currentLevel[currentParentIndex_] = newParentHash_; + function _getLeafData( + IndexedMT storage tree, + uint256 index_ + ) private view returns (LeafData memory) { + _checkIndexExistence(tree, index_, LEAVES_LEVEL); - bool isLeft_ = currentParentIndex_ % 2 == 0; - bytes32 secondNodeHash_; + return tree.leavesData[index_]; + } - if (isLeft_ && currentParentIndex_ == currentLevel.length - 1) { - secondNodeHash_ = _getZeroNodeHash(i); - } else if (isLeft_) { - secondNodeHash_ = currentLevel[currentParentIndex_ + 1]; - } else { - secondNodeHash_ = currentLevel[currentParentIndex_ - 1]; - } - } + function _calculateNodeHash( + IndexedMT storage tree, + uint256 index_, + uint256 level_ + ) private view returns (bytes32) { + if (level_ == LEAVES_LEVEL) { + revert NotANodeLevel(); + } + + 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_); + + return _hashNode(leftChildHash_, rightChildHash_); + } + + function _checkIndexExistence( + IndexedMT storage tree, + uint256 index_, + uint256 level_ + ) private view { + if (index_ >= tree.nodes[level_].length) { + revert IndexOutOfBounds(index_, level_); } } @@ -119,20 +245,20 @@ library IndexedMerkleTree { bytes32 newValue_, uint256 lowLeafIndex_ ) private view returns (uint256 nextLeafIndex_) { - LeafNode storage lowLeaf = tree.leaves[lowLeafIndex_]; + LeafData memory lowLeafData = _getLeafData(tree, lowLeafIndex_); - require( - lowLeafIndex_ > 0 && lowLeafIndex_ < tree.leaves.length, - IndexOutOfBounds(lowLeafIndex_, LEAVES_LEVEL) - ); + nextLeafIndex_ = lowLeafData.nextLeafIndex; - nextLeafIndex_ = lowLeaf.nextLeafIndex; + if ( + lowLeafData.value > newValue_ || + (nextLeafIndex_ != ZERO_IDX && _getLeafData(tree, nextLeafIndex_).value < newValue_) + ) { + revert InvalidLowLeaf(lowLeafIndex_, newValue_); + } + } - require( - lowLeaf.value < newValue_ && - (nextLeafIndex_ == ZERO_IDX || tree.leaves[nextLeafIndex_].value > newValue_), - InvalidLowLeaf(lowLeafIndex_, newValue_) - ); + function _isInitialized(IndexedMT storage tree) private view returns (bool) { + return tree.levelsCount > 0; } function _getZeroNodeHash(uint256 level_) private pure returns (bytes32) { diff --git a/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol new file mode 100644 index 00000000..e7344154 --- /dev/null +++ b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: MIT +// solhint-disable +pragma solidity ^0.8.21; + +import {IndexedMerkleTree} from "../../../libs/data-structures/IndexedMerkleTree.sol"; + +contract IndexedMerkleTreeMock { + using IndexedMerkleTree for *; + + IndexedMerkleTree.UintIndexedMT internal _uintTree; + + function initializeUintTree() external { + _uintTree.initialize(); + } + + function addUint(uint256 value_, uint256 lowLeafIndex_) external returns (uint256) { + return _uintTree.add(value_, lowLeafIndex_); + } + + function getRoot() external view returns (bytes32) { + return _uintTree.getRoot(); + } + + function getTreeLevels() external view returns (uint256) { + return _uintTree.getTreeLevels(); + } + + function getLeafData( + uint256 leafIndex_ + ) external view returns (IndexedMerkleTree.LeafData memory) { + return _uintTree.getLeafData(leafIndex_); + } + + function getNodeHash(uint256 index_, uint256 level_) external view returns (bytes32) { + return _uintTree.getNodeHash(index_, level_); + } + + function getLeavesCount() external view returns (uint256) { + return _uintTree.getLeavesCount(); + } + + function getLevelNodesCount(uint256 level_) external view returns (uint256) { + return _uintTree.getLevelNodesCount(level_); + } +} diff --git a/test/libs/data-structures/IndexedMerkleTree.test.ts b/test/libs/data-structures/IndexedMerkleTree.test.ts new file mode 100644 index 00000000..f09e5921 --- /dev/null +++ b/test/libs/data-structures/IndexedMerkleTree.test.ts @@ -0,0 +1,106 @@ +import { expect } from "chai"; +import hre from "hardhat"; + +import { Reverter } from "@test-helpers"; + +import { IndexedMerkleTreeMock } from "@ethers-v6"; + +const { ethers, networkHelpers } = await hre.network.connect(); + +describe("IndexedMerkleTree", () => { + const reverter: Reverter = new Reverter(networkHelpers); + + const LEAVES_LEVEL = 0n; + + let indexedMT: IndexedMerkleTreeMock; + + function hashLeaf(leafIndex: bigint, value: string, nextLeafIndex: bigint, isActive: boolean): string { + return ethers.solidityPackedKeccak256( + ["bool", "uint256", "bytes32", "uint256"], + [isActive, leafIndex, value, nextLeafIndex], + ); + } + + function encodeBytes32Value(value: bigint): string { + return ethers.toBeHex(value, 32); + } + + before("setup", async () => { + indexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); + + await indexedMT.initializeUintTree(); + + await reverter.snapshot(); + }); + + afterEach("cleanup", async () => { + await reverter.revert(); + }); + + describe("initialize", () => { + it("should correctly initialize IndexedMerkleTree", async () => { + const zeroLeafHash = hashLeaf(0n, ethers.ZeroHash, 0n, false); + + expect(await indexedMT.getRoot()).to.be.eq(zeroLeafHash); + expect(await indexedMT.getTreeLevels()).to.be.eq(1); + expect(await indexedMT.getLeavesCount()).to.be.eq(1); + expect(await indexedMT.getNodeHash(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, + "IndexedMerkleTreeNotInitialized", + ); + }); + }); + + 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; + + for (let i = 0; i < count; ++i) { + const currentIndex = startIndex + BigInt(i); + + await indexedMT.addUint(value, lowLeafIndex); + + const leafData = await indexedMT.getLeafData(currentIndex); + + expect(leafData.value).to.be.eq(value); + expect(leafData.nextLeafIndex).to.be.eq(0n); + + const leafHash = hashLeaf(currentIndex, encodeBytes32Value(value), 0n, true); + expect(await indexedMT.getNodeHash(currentIndex, LEAVES_LEVEL)).to.be.eq(leafHash); + + const lowLeafData = await indexedMT.getLeafData(lowLeafIndex); + + expect(lowLeafData.value).to.be.eq(lowLeafValue); + expect(lowLeafData.nextLeafIndex).to.be.eq(currentIndex); + + const lowLeafNewHash = hashLeaf(lowLeafIndex, encodeBytes32Value(lowLeafValue), currentIndex, true); + expect(await indexedMT.getNodeHash(lowLeafIndex, LEAVES_LEVEL)).to.be.eq(lowLeafNewHash); + + lowLeafIndex = currentIndex; + lowLeafValue = value; + value *= 2n; + } + + const expectedLevelsCount = Math.ceil(Math.log2(count + 1)) + 1; + + expect(await indexedMT.getTreeLevels()).to.be.eq(expectedLevelsCount); + }); + + 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"); + }); + }); +}); From a098b40020ae64244f3a5c9ca88324acc91a95d5 Mon Sep 17 00:00:00 2001 From: Oleh Komendant Date: Mon, 3 Nov 2025 23:02:34 +0200 Subject: [PATCH 3/6] Add getProof and verifyProof logic --- .../data-structures/IndexedMerkleTree.sol | 118 ++++++++++++-- .../data-structures/IndexedMerkleTreeMock.sol | 15 ++ .../data-structures/IndexedMerkleTree.test.ts | 144 +++++++++++++++++- 3 files changed, 264 insertions(+), 13 deletions(-) diff --git a/contracts/libs/data-structures/IndexedMerkleTree.sol b/contracts/libs/data-structures/IndexedMerkleTree.sol index 56bcee46..9f04773f 100644 --- a/contracts/libs/data-structures/IndexedMerkleTree.sol +++ b/contracts/libs/data-structures/IndexedMerkleTree.sol @@ -18,6 +18,25 @@ library IndexedMerkleTree { return _add(tree._indexedMT, bytes32(value_), lowLeafIndex_); } + function getProof( + UintIndexedMT storage tree, + uint256 index_, + uint256 value_ + ) internal view returns (Proof memory) { + return _proof(tree._indexedMT, index_, bytes32(value_)); + } + + function verifyProof( + UintIndexedMT storage tree, + Proof memory proof_ + ) internal view returns (bool) { + return _verifyProof(tree._indexedMT, proof_); + } + + function processProof(Proof memory proof_) internal pure returns (bytes32) { + return _processProof(proof_); + } + function getRoot(UintIndexedMT storage tree) internal view returns (bytes32) { return _getRoot(tree._indexedMT); } @@ -63,6 +82,15 @@ library IndexedMerkleTree { uint256 levelsCount; } + struct Proof { + bytes32 root; + bytes32[] siblings; + bool existence; + uint256 index; + bytes32 value; + uint256 nextLeafIndex; + } + struct LeafData { bytes32 value; uint256 nextLeafIndex; @@ -70,6 +98,7 @@ library IndexedMerkleTree { error IndexOutOfBounds(uint256 index, uint256 level); error InvalidLowLeaf(uint256 lowLeafIndex, bytes32 newValue); + error InvalidProofIndex(uint256 index, bytes32 value); error NotANodeLevel(); error IndexedMerkleTreeNotInitialized(); @@ -81,10 +110,8 @@ library IndexedMerkleTree { function _initialize(IndexedMT storage tree) private { if (_isInitialized(tree)) revert IndexedMerkleTreeNotInitialized(); - bytes32 zeroNodeHash_ = _getZeroNodeHash(LEAVES_LEVEL); - tree.leavesData.push(LeafData({value: ZERO_HASH, nextLeafIndex: ZERO_IDX})); - tree.nodes[LEAVES_LEVEL].push(zeroNodeHash_); + tree.nodes[LEAVES_LEVEL].push(_hashLeaf(0, 0, 0, true)); tree.levelsCount++; } @@ -173,6 +200,64 @@ library IndexedMerkleTree { } } + 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_); + } + + uint256 parentIndex_ = index_; + + for (uint256 i = 0; i < proof_.siblings.length; ++i) { + uint256 currentLevelIndex_ = parentIndex_ % 2 == 0 + ? parentIndex_ + 1 + : parentIndex_ - 1; + + proof_.siblings[i] = tree.nodes[i][currentLevelIndex_]; + + parentIndex_ /= 2; + } + + return proof_; + } + + function _verifyProof( + IndexedMT storage tree, + Proof memory proof_ + ) private view returns (bool) { + return _processProof(proof_) == _getRoot(tree); + } + + function _processProof(Proof memory proof_) private pure returns (bytes32) { + bytes32 computedHash_ = _hashLeaf(proof_.index, proof_.value, proof_.nextLeafIndex, true); + + for (uint256 i = 0; i < proof_.siblings.length; ++i) { + if ((proof_.index >> i) & 1 == 1) { + computedHash_ = _hashNode(proof_.siblings[i], computedHash_); + } else { + computedHash_ = _hashNode(computedHash_, proof_.siblings[i]); + } + } + + return computedHash_; + } + function _getRoot(IndexedMT storage tree) private view returns (bytes32) { return tree.nodes[tree.levelsCount - 1][0]; } @@ -242,19 +327,28 @@ library IndexedMerkleTree { function _checkLowLeaf( IndexedMT storage tree, - bytes32 newValue_, + bytes32 value_, uint256 lowLeafIndex_ - ) private view returns (uint256 nextLeafIndex_) { + ) private view returns (uint256) { + if (!_isLowLeaf(tree, value_, lowLeafIndex_)) { + revert InvalidLowLeaf(lowLeafIndex_, value_); + } + + return _getLeafData(tree, lowLeafIndex_).nextLeafIndex; + } + + function _isLowLeaf( + IndexedMT storage tree, + bytes32 value_, + uint256 lowLeafIndex_ + ) private view returns (bool) { LeafData memory lowLeafData = _getLeafData(tree, lowLeafIndex_); - nextLeafIndex_ = lowLeafData.nextLeafIndex; + uint256 nextLeafIndex_ = lowLeafData.nextLeafIndex; - if ( - lowLeafData.value > newValue_ || - (nextLeafIndex_ != ZERO_IDX && _getLeafData(tree, nextLeafIndex_).value < newValue_) - ) { - revert InvalidLowLeaf(lowLeafIndex_, newValue_); - } + return + lowLeafData.value < value_ && + (nextLeafIndex_ == ZERO_IDX || _getLeafData(tree, nextLeafIndex_).value > value_); } function _isInitialized(IndexedMT storage tree) private view returns (bool) { diff --git a/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol index e7344154..4f9e2c2b 100644 --- a/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol +++ b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol @@ -17,6 +17,21 @@ contract IndexedMerkleTreeMock { return _uintTree.add(value_, lowLeafIndex_); } + function getProof( + uint256 index_, + uint256 value_ + ) external view returns (IndexedMerkleTree.Proof memory) { + return _uintTree.getProof(index_, value_); + } + + function verifyProof(IndexedMerkleTree.Proof memory proof_) external view returns (bool) { + return _uintTree.verifyProof(proof_); + } + + function processProof(IndexedMerkleTree.Proof memory proof_) external pure returns (bytes32) { + return IndexedMerkleTree.processProof(proof_); + } + function getRoot() external view returns (bytes32) { return _uintTree.getRoot(); } diff --git a/test/libs/data-structures/IndexedMerkleTree.test.ts b/test/libs/data-structures/IndexedMerkleTree.test.ts index f09e5921..f899fcde 100644 --- a/test/libs/data-structures/IndexedMerkleTree.test.ts +++ b/test/libs/data-structures/IndexedMerkleTree.test.ts @@ -14,6 +14,10 @@ describe("IndexedMerkleTree", () => { let indexedMT: IndexedMerkleTreeMock; + function hashNode(leftChild: string, rightChild: string): string { + return ethers.solidityPackedKeccak256(["bytes32", "bytes32"], [leftChild, rightChild]); + } + function hashLeaf(leafIndex: bigint, value: string, nextLeafIndex: bigint, isActive: boolean): string { return ethers.solidityPackedKeccak256( ["bool", "uint256", "bytes32", "uint256"], @@ -39,7 +43,7 @@ describe("IndexedMerkleTree", () => { describe("initialize", () => { it("should correctly initialize IndexedMerkleTree", async () => { - const zeroLeafHash = hashLeaf(0n, ethers.ZeroHash, 0n, false); + const zeroLeafHash = hashLeaf(0n, ethers.ZeroHash, 0n, true); expect(await indexedMT.getRoot()).to.be.eq(zeroLeafHash); expect(await indexedMT.getTreeLevels()).to.be.eq(1); @@ -103,4 +107,142 @@ describe("IndexedMerkleTree", () => { await expect(indexedMT.addUint(25n, 1n)).to.be.revertedWithCustomError(indexedMT, "InvalidLowLeaf"); }); }); + + describe("getProof", () => { + const values: bigint[] = [0n, 10n, 20n, 30n]; + let leaves: string[] = []; + + beforeEach("setup", async () => { + for (let i = 0; i < values.length; i++) { + leaves.push( + hashLeaf(BigInt(i), encodeBytes32Value(values[i]), i == values.length - 1 ? 0n : BigInt(i + 1), true), + ); + + if (i > 0) { + await indexedMT.addUint(values[i], i - 1); + } + } + }); + + afterEach("clean", async () => { + leaves = []; + }); + + it("should return correct inclusion proof", async () => { + let index = 2n; + let value = 20n; + let nextLeafIndex = 3n; + let expectedSiblings = [leaves[3], hashNode(leaves[0], leaves[1])]; + let proof = await indexedMT.getProof(index, value); + + expect(proof.root).to.be.eq(await indexedMT.getRoot()); + expect(proof.existence).to.be.true; + expect(proof.index).to.be.eq(index); + expect(proof.value).to.be.eq(value); + expect(proof.nextLeafIndex).to.be.eq(nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedSiblings); + + index = 1n; + value = 10n; + nextLeafIndex = 2n; + expectedSiblings = [leaves[0], hashNode(leaves[2], leaves[3])]; + proof = await indexedMT.getProof(index, value); + + expect(proof.root).to.be.eq(await indexedMT.getRoot()); + expect(proof.existence).to.be.true; + expect(proof.index).to.be.eq(index); + expect(proof.value).to.be.eq(value); + expect(proof.nextLeafIndex).to.be.eq(nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedSiblings); + }); + + it("should return correct exclusion proof", async () => { + const index = 1n; + const value = 15n; + const nextLeafIndex = 2n; + const expectedSiblings = [leaves[0], hashNode(leaves[2], leaves[3])]; + const proof = await indexedMT.getProof(index, value); + + expect(proof.root).to.be.eq(await indexedMT.getRoot()); + expect(proof.existence).to.be.false; + expect(proof.index).to.be.eq(index); + expect(proof.value).to.be.eq(values[Number(index)]); + expect(proof.nextLeafIndex).to.be.eq(nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedSiblings); + }); + + it("should get exception if pass invalid index", async () => { + const index = 1n; + const value = 5n; + + await expect(indexedMT.getProof(index, value)) + .to.be.revertedWithCustomError(indexedMT, "InvalidProofIndex") + .withArgs(index, value); + }); + }); + + describe("verifyProof", () => { + const values: bigint[] = [0n, 10n, 20n, 30n, 40n, 50n]; + let leaves: string[] = []; + + beforeEach("setup", async () => { + for (let i = 0; i < values.length; i++) { + leaves.push( + hashLeaf(BigInt(i), encodeBytes32Value(values[i]), i == values.length - 1 ? 0n : BigInt(i + 1), true), + ); + + if (i > 0) { + await indexedMT.addUint(values[i], i - 1); + } + } + }); + + afterEach("clean", async () => { + leaves = []; + }); + + it("should correctly verify inclusion proofs", async () => { + let index = 1n; + let proof = await indexedMT.getProof(index, values[Number(index)]); + + expect( + await indexedMT.verifyProof({ + root: proof.root, + existence: proof.existence, + siblings: [...proof.siblings], + index: proof.index, + value: proof.value, + nextLeafIndex: proof.nextLeafIndex, + }), + ).to.be.true; + + index = 2n; + proof = await indexedMT.getProof(index, values[Number(index)]); + + expect( + await indexedMT.verifyProof({ + root: proof.root, + existence: proof.existence, + siblings: [...proof.siblings], + index: proof.index, + value: proof.value, + nextLeafIndex: proof.nextLeafIndex, + }), + ).to.be.true; + + index = 5n; + proof = await indexedMT.getProof(index, values[Number(index)]); + + expect( + await indexedMT.verifyProof({ + root: proof.root, + existence: proof.existence, + siblings: [...proof.siblings], + index: proof.index, + value: proof.value, + nextLeafIndex: proof.nextLeafIndex, + }), + ).to.be.true; + }); + }); }); From 00cb6d033560456af57f63609cf57429545c46e9 Mon Sep 17 00:00:00 2001 From: Oleh Komendant Date: Tue, 4 Nov 2025 10:12:59 +0200 Subject: [PATCH 4/6] Fix getProof function --- contracts/libs/data-structures/IndexedMerkleTree.sol | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/contracts/libs/data-structures/IndexedMerkleTree.sol b/contracts/libs/data-structures/IndexedMerkleTree.sol index 9f04773f..9bbbb76f 100644 --- a/contracts/libs/data-structures/IndexedMerkleTree.sol +++ b/contracts/libs/data-structures/IndexedMerkleTree.sol @@ -229,7 +229,9 @@ library IndexedMerkleTree { ? parentIndex_ + 1 : parentIndex_ - 1; - proof_.siblings[i] = tree.nodes[i][currentLevelIndex_]; + proof_.siblings[i] = currentLevelIndex_ < _getLevelNodesCount(tree, i) + ? tree.nodes[i][currentLevelIndex_] + : _getZeroNodeHash(i); parentIndex_ /= 2; } From 917846a078bd6e956f2667f3cf52e2d51fe024eb Mon Sep 17 00:00:00 2001 From: Oleh Komendant Date: Thu, 6 Nov 2025 12:23:20 +0200 Subject: [PATCH 5/6] Finish TS IndexedMerkleTree implementation & add some unit tests --- test/helpers/indexed-merkle-tree.ts | 385 ++++++++++++++++++ .../data-structures/IndexedMerkleTree.test.ts | 346 +++++++++++----- 2 files changed, 634 insertions(+), 97 deletions(-) create mode 100644 test/helpers/indexed-merkle-tree.ts diff --git a/test/helpers/indexed-merkle-tree.ts b/test/helpers/indexed-merkle-tree.ts new file mode 100644 index 00000000..a7b797a6 --- /dev/null +++ b/test/helpers/indexed-merkle-tree.ts @@ -0,0 +1,385 @@ +import { ethers } from "ethers"; + +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 hashNode(leftChild: string, rightChild: string): string { + return ethers.solidityPackedKeccak256(["bytes32", "bytes32"], [leftChild, rightChild]); +} + +export function hashIndexedLeaf(leafData: IndexedLeafData): string { + return ethers.solidityPackedKeccak256( + ["bool", "uint256", "bytes32", "uint256"], + [leafData.isActive, leafData.index, leafData.value, leafData.nextIndex], + ); +} + +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; + + public static buildMerkleTree(leavesData?: IndexedLeafData[], maxLevelsCount: number = 256): IndexedMerkleTree { + return new IndexedMerkleTree( + leavesData ?? [ + { + index: 0n, + value: encodeBytes32Value(0n), + nextIndex: 0n, + isActive: true, + }, + ], + maxLevelsCount, + ); + } + + private constructor(leavesData: IndexedLeafData[], maxLevelsCount: number) { + 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.leavesData[lowLeafIndex.toString()].nextLeafIndex = newLeafIndex; + this._updateMerkleHashes(lowLeafIndex); + + this._pushLeaf(newLeafIndex, { value: value, nextLeafIndex: nextLeafIndex }); + + return newLeafIndex; + } + + 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 = hashIndexedLeaf({ + 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 = hashNode(proof.siblings[i], computedHash); + } else { + computedHash = hashNode(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 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()]) || []; + } + + 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 = hashIndexedLeaf({ + 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 = hashIndexedLeaf({ + 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 = hashIndexedLeaf({ + index: ZERO_IDX, + value: encodeBytes32Value(0n), + nextIndex: ZERO_IDX, + isActive: false, + }); + this.zeroHashesCache.push(currentHash); + + for (let i = 1; i <= maxDepth; i++) { + currentHash = hashNode(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 hashIndexedLeaf(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(hashNode(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 hashNode(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 _getLevelNodesCount(level: bigint): bigint { + return BigInt(Object.values(this.levels[level.toString()]).length); + } + + 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 index f899fcde..a0906a65 100644 --- a/test/libs/data-structures/IndexedMerkleTree.test.ts +++ b/test/libs/data-structures/IndexedMerkleTree.test.ts @@ -5,6 +5,8 @@ import { Reverter } from "@test-helpers"; import { IndexedMerkleTreeMock } from "@ethers-v6"; +import { IndexedMerkleTree, encodeBytes32Value, hashIndexedLeaf } from "@/test/helpers/indexed-merkle-tree.ts"; + const { ethers, networkHelpers } = await hre.network.connect(); describe("IndexedMerkleTree", () => { @@ -14,19 +16,11 @@ describe("IndexedMerkleTree", () => { let indexedMT: IndexedMerkleTreeMock; - function hashNode(leftChild: string, rightChild: string): string { - return ethers.solidityPackedKeccak256(["bytes32", "bytes32"], [leftChild, rightChild]); - } - - function hashLeaf(leafIndex: bigint, value: string, nextLeafIndex: bigint, isActive: boolean): string { - return ethers.solidityPackedKeccak256( - ["bool", "uint256", "bytes32", "uint256"], - [isActive, leafIndex, value, nextLeafIndex], - ); - } + function getRandomIntInclusive(min: number, max: number): number { + min = Math.ceil(min); + max = Math.floor(max); - function encodeBytes32Value(value: bigint): string { - return ethers.toBeHex(value, 32); + return Math.floor(Math.random() * (max - min + 1)) + min; } before("setup", async () => { @@ -43,9 +37,16 @@ describe("IndexedMerkleTree", () => { describe("initialize", () => { it("should correctly initialize IndexedMerkleTree", async () => { - const zeroLeafHash = hashLeaf(0n, ethers.ZeroHash, 0n, true); + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + const zeroLeafHash = hashIndexedLeaf({ + index: 0n, + isActive: true, + nextIndex: 0n, + value: ethers.ZeroHash, + }); expect(await indexedMT.getRoot()).to.be.eq(zeroLeafHash); + expect(await indexedMT.getRoot()).to.be.eq(localIndexedMerkleTree.getRoot()); expect(await indexedMT.getTreeLevels()).to.be.eq(1); expect(await indexedMT.getLeavesCount()).to.be.eq(1); expect(await indexedMT.getNodeHash(0, LEAVES_LEVEL)).to.be.eq(zeroLeafHash); @@ -68,17 +69,25 @@ describe("IndexedMerkleTree", () => { 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.getLeafData(currentIndex); expect(leafData.value).to.be.eq(value); expect(leafData.nextLeafIndex).to.be.eq(0n); - const leafHash = hashLeaf(currentIndex, encodeBytes32Value(value), 0n, true); + const leafHash = hashIndexedLeaf({ + index: currentIndex, + value: encodeBytes32Value(value), + nextIndex: 0n, + isActive: true, + }); expect(await indexedMT.getNodeHash(currentIndex, LEAVES_LEVEL)).to.be.eq(leafHash); const lowLeafData = await indexedMT.getLeafData(lowLeafIndex); @@ -86,12 +95,19 @@ describe("IndexedMerkleTree", () => { expect(lowLeafData.value).to.be.eq(lowLeafValue); expect(lowLeafData.nextLeafIndex).to.be.eq(currentIndex); - const lowLeafNewHash = hashLeaf(lowLeafIndex, encodeBytes32Value(lowLeafValue), currentIndex, true); + const lowLeafNewHash = hashIndexedLeaf({ + index: lowLeafIndex, + value: encodeBytes32Value(lowLeafValue), + nextIndex: currentIndex, + isActive: true, + }); expect(await indexedMT.getNodeHash(lowLeafIndex, LEAVES_LEVEL)).to.be.eq(lowLeafNewHash); lowLeafIndex = currentIndex; lowLeafValue = value; value *= 2n; + + expect(await indexedMT.getRoot()).to.be.eq(localIndexedMerkleTree.getRoot()); } const expectedLevelsCount = Math.ceil(Math.log2(count + 1)) + 1; @@ -99,6 +115,30 @@ describe("IndexedMerkleTree", () => { expect(await indexedMT.getTreeLevels()).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.getRoot()).to.be.eq(localIndexedMerkleTree.getRoot()); + + const leafData = await indexedMT.getLeafData(index); + + expect(leafData.value).to.be.eq(currentValue); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + + expect((await indexedMT.getLeafData(lowLeafIndex)).nextLeafIndex).to.be.eq(index); + } + }); + it("should get exception if pass invalid low leaf index", async () => { await indexedMT.addUint(10n, 0n); await indexedMT.addUint(20n, 1n); @@ -110,65 +150,142 @@ describe("IndexedMerkleTree", () => { describe("getProof", () => { const values: bigint[] = [0n, 10n, 20n, 30n]; - let leaves: string[] = []; + let localIndexedMerkleTree: IndexedMerkleTree; beforeEach("setup", async () => { - for (let i = 0; i < values.length; i++) { - leaves.push( - hashLeaf(BigInt(i), encodeBytes32Value(values[i]), i == values.length - 1 ? 0n : BigInt(i + 1), true), - ); - - if (i > 0) { - await indexedMT.addUint(values[i], i - 1); - } - } - }); + localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + + for (let i = 1; i < values.length; i++) { + localIndexedMerkleTree.add(encodeBytes32Value(values[i])); - afterEach("clean", async () => { - leaves = []; + await indexedMT.addUint(values[i], i - 1); + } }); it("should return correct inclusion proof", async () => { let index = 2n; let value = 20n; - let nextLeafIndex = 3n; - let expectedSiblings = [leaves[3], hashNode(leaves[0], leaves[1])]; + let expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); let proof = await indexedMT.getProof(index, value); - expect(proof.root).to.be.eq(await indexedMT.getRoot()); + expect(proof.root).to.be.eq(expectedProof.root); expect(proof.existence).to.be.true; - expect(proof.index).to.be.eq(index); - expect(proof.value).to.be.eq(value); - expect(proof.nextLeafIndex).to.be.eq(nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedSiblings); + expect(proof.existence).to.be.eq(expectedProof.existence); + expect(proof.index).to.be.eq(expectedProof.index); + expect(proof.value).to.be.eq(expectedProof.value); + expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); index = 1n; value = 10n; - nextLeafIndex = 2n; - expectedSiblings = [leaves[0], hashNode(leaves[2], leaves[3])]; + expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); proof = await indexedMT.getProof(index, value); - expect(proof.root).to.be.eq(await indexedMT.getRoot()); + expect(proof.root).to.be.eq(expectedProof.root); expect(proof.existence).to.be.true; - expect(proof.index).to.be.eq(index); - expect(proof.value).to.be.eq(value); - expect(proof.nextLeafIndex).to.be.eq(nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedSiblings); + expect(proof.existence).to.be.eq(expectedProof.existence); + expect(proof.index).to.be.eq(expectedProof.index); + expect(proof.value).to.be.eq(expectedProof.value); + expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); + }); + + it("should return correct inclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); + 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 expectedProof = newLocalIndexedMT.getProof(index, valueToProve); + const proof = await newIndexedMT.getProof(index, valueToProve); + + expect(proof.root).to.be.eq(expectedProof.root); + expect(proof.existence).to.be.true; + expect(proof.existence).to.be.eq(expectedProof.existence); + expect(proof.index).to.be.eq(expectedProof.index); + expect(proof.value).to.be.eq(expectedProof.value); + expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); + } }); it("should return correct exclusion proof", async () => { const index = 1n; const value = 15n; - const nextLeafIndex = 2n; - const expectedSiblings = [leaves[0], hashNode(leaves[2], leaves[3])]; + const expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); const proof = await indexedMT.getProof(index, value); expect(proof.root).to.be.eq(await indexedMT.getRoot()); expect(proof.existence).to.be.false; - expect(proof.index).to.be.eq(index); - expect(proof.value).to.be.eq(values[Number(index)]); - expect(proof.nextLeafIndex).to.be.eq(nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedSiblings); + expect(proof.existence).to.be.eq(expectedProof.existence); + expect(proof.index).to.be.eq(expectedProof.index); + expect(proof.value).to.be.eq(expectedProof.value); + expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); + }); + + it("should return correct exclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); + 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 expectedProof = newLocalIndexedMT.getProof(index, valueToProve); + const proof = await newIndexedMT.getProof(index, valueToProve); + + expect(proof.root).to.be.eq(expectedProof.root); + expect(proof.existence).to.be.false; + expect(proof.existence).to.be.eq(expectedProof.existence); + expect(proof.index).to.be.eq(expectedProof.index); + expect(proof.value).to.be.eq(expectedProof.value); + expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); + expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); + } }); it("should get exception if pass invalid index", async () => { @@ -183,66 +300,101 @@ describe("IndexedMerkleTree", () => { describe("verifyProof", () => { const values: bigint[] = [0n, 10n, 20n, 30n, 40n, 50n]; - let leaves: string[] = []; + let localIndexedMerkleTree: IndexedMerkleTree; beforeEach("setup", async () => { - for (let i = 0; i < values.length; i++) { - leaves.push( - hashLeaf(BigInt(i), encodeBytes32Value(values[i]), i == values.length - 1 ? 0n : BigInt(i + 1), true), - ); - - if (i > 0) { - await indexedMT.addUint(values[i], i - 1); - } - } - }); + localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + + for (let i = 1; i < values.length; i++) { + const lowIndex = localIndexedMerkleTree.getLowLeafIndex(encodeBytes32Value(values[i])); + localIndexedMerkleTree.add(encodeBytes32Value(values[i])); - afterEach("clean", async () => { - leaves = []; + await indexedMT.addUint(values[i], lowIndex); + } }); it("should correctly verify inclusion proofs", async () => { let index = 1n; - let proof = await indexedMT.getProof(index, values[Number(index)]); - - expect( - await indexedMT.verifyProof({ - root: proof.root, - existence: proof.existence, - siblings: [...proof.siblings], - index: proof.index, - value: proof.value, - nextLeafIndex: proof.nextLeafIndex, - }), - ).to.be.true; + let proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + + expect(await indexedMT.verifyProof(proof)).to.be.true; index = 2n; - proof = await indexedMT.getProof(index, values[Number(index)]); - - expect( - await indexedMT.verifyProof({ - root: proof.root, - existence: proof.existence, - siblings: [...proof.siblings], - index: proof.index, - value: proof.value, - nextLeafIndex: proof.nextLeafIndex, - }), - ).to.be.true; + proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + + expect(await indexedMT.verifyProof(proof)).to.be.true; index = 5n; - proof = await indexedMT.getProof(index, values[Number(index)]); - - expect( - await indexedMT.verifyProof({ - root: proof.root, - existence: proof.existence, - siblings: [...proof.siblings], - index: proof.index, - value: proof.value, - nextLeafIndex: proof.nextLeafIndex, - }), - ).to.be.true; + proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + + expect(await indexedMT.verifyProof(proof)).to.be.true; + }); + + it("should correctly verify inclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); + 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.verifyProof(proof)).to.be.true; + } + }); + + it("should correctly verify exclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); + 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.verifyProof(proof)).to.be.true; + } }); }); }); From 6e8cda60e05af19a2721c914a1dfbdd38f8a8672 Mon Sep 17 00:00:00 2001 From: Oleh Komendant Date: Fri, 7 Nov 2025 15:40:10 +0200 Subject: [PATCH 6/6] Add Bytes32IndexedMerkleTree and AddressIndexedMerkleTree --- .../data-structures/IndexedMerkleTree.sol | 137 ++- .../data-structures/IndexedMerkleTreeMock.sol | 112 ++- .../data-structures/IndexedMerkleTree.test.ts | 859 ++++++++++++------ 3 files changed, 836 insertions(+), 272 deletions(-) diff --git a/contracts/libs/data-structures/IndexedMerkleTree.sol b/contracts/libs/data-structures/IndexedMerkleTree.sol index 9bbbb76f..e3d23331 100644 --- a/contracts/libs/data-structures/IndexedMerkleTree.sol +++ b/contracts/libs/data-structures/IndexedMerkleTree.sol @@ -71,6 +71,136 @@ library IndexedMerkleTree { return _getLevelNodesCount(tree._indexedMT, level_); } + struct Bytes32IndexedMT { + IndexedMT _indexedMT; + } + + function initialize(Bytes32IndexedMT storage tree) internal { + _initialize(tree._indexedMT); + } + + function add( + Bytes32IndexedMT storage tree, + bytes32 value_, + uint256 lowLeafIndex_ + ) internal returns (uint256) { + return _add(tree._indexedMT, value_, lowLeafIndex_); + } + + function getProof( + Bytes32IndexedMT storage tree, + uint256 index_, + bytes32 value_ + ) internal view returns (Proof memory) { + return _proof(tree._indexedMT, index_, value_); + } + + function verifyProof( + Bytes32IndexedMT storage tree, + Proof memory proof_ + ) internal view returns (bool) { + return _verifyProof(tree._indexedMT, proof_); + } + + function getRoot(Bytes32IndexedMT storage tree) internal view returns (bytes32) { + return _getRoot(tree._indexedMT); + } + + function getTreeLevels(Bytes32IndexedMT storage tree) internal view returns (uint256) { + return _getTreeLevels(tree._indexedMT); + } + + function getLeafData( + Bytes32IndexedMT storage tree, + uint256 leafIndex_ + ) internal view returns (LeafData memory) { + return _getLeafData(tree._indexedMT, leafIndex_); + } + + function getNodeHash( + Bytes32IndexedMT storage tree, + uint256 index_, + uint256 level_ + ) internal view returns (bytes32) { + return _getNodeHash(tree._indexedMT, index_, level_); + } + + function getLeavesCount(Bytes32IndexedMT storage tree) internal view returns (uint256) { + return _getLeavesCount(tree._indexedMT); + } + + function getLevelNodesCount( + Bytes32IndexedMT storage tree, + uint256 level_ + ) internal view returns (uint256) { + return _getLevelNodesCount(tree._indexedMT, level_); + } + + struct AddressIndexedMT { + IndexedMT _indexedMT; + } + + function initialize(AddressIndexedMT storage tree) internal { + _initialize(tree._indexedMT); + } + + function add( + AddressIndexedMT storage tree, + address value_, + uint256 lowLeafIndex_ + ) internal returns (uint256) { + return _add(tree._indexedMT, bytes32(uint256(uint160(value_))), lowLeafIndex_); + } + + function getProof( + AddressIndexedMT storage tree, + uint256 index_, + address value_ + ) internal view returns (Proof memory) { + return _proof(tree._indexedMT, index_, bytes32(uint256(uint160(value_)))); + } + + function verifyProof( + AddressIndexedMT storage tree, + Proof memory proof_ + ) internal view returns (bool) { + return _verifyProof(tree._indexedMT, proof_); + } + + function getRoot(AddressIndexedMT storage tree) internal view returns (bytes32) { + return _getRoot(tree._indexedMT); + } + + function getTreeLevels(AddressIndexedMT storage tree) internal view returns (uint256) { + return _getTreeLevels(tree._indexedMT); + } + + function getLeafData( + AddressIndexedMT storage tree, + uint256 leafIndex_ + ) internal view returns (LeafData memory) { + return _getLeafData(tree._indexedMT, leafIndex_); + } + + function getNodeHash( + AddressIndexedMT storage tree, + uint256 index_, + uint256 level_ + ) internal view returns (bytes32) { + return _getNodeHash(tree._indexedMT, index_, level_); + } + + function getLeavesCount(AddressIndexedMT storage tree) internal view returns (uint256) { + return _getLeavesCount(tree._indexedMT); + } + + function getLevelNodesCount( + AddressIndexedMT storage tree, + uint256 level_ + ) internal view returns (uint256) { + return _getLevelNodesCount(tree._indexedMT, level_); + } + uint256 internal constant LEAVES_LEVEL = 0; uint64 internal constant ZERO_IDX = 0; @@ -101,6 +231,7 @@ library IndexedMerkleTree { error InvalidProofIndex(uint256 index, bytes32 value); error NotANodeLevel(); error IndexedMerkleTreeNotInitialized(); + error IndexedMerkleTreeAlreadyInitialized(); modifier onlyInitialized(IndexedMT storage tree) { if (!_isInitialized(tree)) revert IndexedMerkleTreeNotInitialized(); @@ -108,7 +239,7 @@ library IndexedMerkleTree { } function _initialize(IndexedMT storage tree) private { - if (_isInitialized(tree)) revert IndexedMerkleTreeNotInitialized(); + 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)); @@ -301,10 +432,6 @@ library IndexedMerkleTree { uint256 index_, uint256 level_ ) private view returns (bytes32) { - if (level_ == LEAVES_LEVEL) { - revert NotANodeLevel(); - } - uint256 childrenLevel_ = level_ - 1; uint256 leftChild_ = index_ * 2; uint256 rightChild_ = index_ * 2 + 1; diff --git a/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol index 4f9e2c2b..fd3e2081 100644 --- a/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol +++ b/contracts/mock/libs/data-structures/IndexedMerkleTreeMock.sol @@ -8,53 +8,149 @@ 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 addUint(uint256 value_, uint256 lowLeafIndex_) external returns (uint256) { return _uintTree.add(value_, lowLeafIndex_); } - function getProof( + 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 getProofUint( uint256 index_, uint256 value_ ) external view returns (IndexedMerkleTree.Proof memory) { return _uintTree.getProof(index_, value_); } - function verifyProof(IndexedMerkleTree.Proof memory proof_) external view returns (bool) { + 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 pure returns (bytes32) { return IndexedMerkleTree.processProof(proof_); } - function getRoot() external view returns (bytes32) { + function getRootUint() external view returns (bytes32) { return _uintTree.getRoot(); } - function getTreeLevels() external view returns (uint256) { + 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 getLeafData( + 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 getNodeHash(uint256 index_, uint256 level_) external view returns (bytes32) { + 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 getLeavesCount() external view returns (uint256) { + 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 getLevelNodesCount(uint256 level_) external view returns (uint256) { + 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_); + } } diff --git a/test/libs/data-structures/IndexedMerkleTree.test.ts b/test/libs/data-structures/IndexedMerkleTree.test.ts index a0906a65..f522eb46 100644 --- a/test/libs/data-structures/IndexedMerkleTree.test.ts +++ b/test/libs/data-structures/IndexedMerkleTree.test.ts @@ -5,7 +5,8 @@ import { Reverter } from "@test-helpers"; import { IndexedMerkleTreeMock } from "@ethers-v6"; -import { IndexedMerkleTree, encodeBytes32Value, hashIndexedLeaf } from "@/test/helpers/indexed-merkle-tree.ts"; +import { IndexedMerkleTree as IndexedMerkleTreeLib } from "../../../generated-types/ethers/mock/libs/data-structures/IndexedMerkleTreeMock.ts"; +import { IndexedMerkleTree, Proof, encodeBytes32Value, hashIndexedLeaf } from "@/test/helpers/indexed-merkle-tree.ts"; const { ethers, networkHelpers } = await hre.network.connect(); @@ -23,11 +24,27 @@ describe("IndexedMerkleTree", () => { 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]); + } + before("setup", async () => { indexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); - await indexedMT.initializeUintTree(); - await reverter.snapshot(); }); @@ -35,366 +52,690 @@ describe("IndexedMerkleTree", () => { await reverter.revert(); }); - describe("initialize", () => { - it("should correctly initialize IndexedMerkleTree", async () => { - const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); - const zeroLeafHash = hashIndexedLeaf({ - index: 0n, - isActive: true, - nextIndex: 0n, - value: ethers.ZeroHash, + describe("UintIndexedMerkleTree", () => { + beforeEach("setup", async () => { + await indexedMT.initializeUintTree(); + }); + + 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); }); - expect(await indexedMT.getRoot()).to.be.eq(zeroLeafHash); - expect(await indexedMT.getRoot()).to.be.eq(localIndexedMerkleTree.getRoot()); - expect(await indexedMT.getTreeLevels()).to.be.eq(1); - expect(await indexedMT.getLeavesCount()).to.be.eq(1); - expect(await indexedMT.getNodeHash(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", + ); + }); }); - it("should get exception if try to initialize twice", async () => { - await expect(indexedMT.initializeUintTree()).to.be.revertedWithCustomError( - indexedMT, - "IndexedMerkleTreeNotInitialized", - ); + 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()); + } + + 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); + } + + 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"); + + await expect(newIndexedMT.addUint(10n, 0n)).to.be.revertedWithCustomError( + newIndexedMT, + "IndexedMerkleTreeNotInitialized", + ); + }); }); - }); - 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; + describe("getProof", () => { + const values: bigint[] = [0n, 10n, 20n, 30n]; + let localIndexedMerkleTree: IndexedMerkleTree; - const count = 10; + beforeEach("setup", async () => { + localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); - const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + for (let i = 1; i < values.length; i++) { + localIndexedMerkleTree.add(encodeBytes32Value(values[i])); - for (let i = 0; i < count; ++i) { - const currentIndex = startIndex + BigInt(i); + await indexedMT.addUint(values[i], i - 1); + } + }); - await indexedMT.addUint(value, lowLeafIndex); - localIndexedMerkleTree.add(encodeBytes32Value(value)); + 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); - const leafData = await indexedMT.getLeafData(currentIndex); + compareProofs(proof, expectedProof, true); - expect(leafData.value).to.be.eq(value); - expect(leafData.nextLeafIndex).to.be.eq(0n); + index = 1n; + value = 10n; + expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); + proof = await indexedMT.getProofUint(index, value); - const leafHash = hashIndexedLeaf({ - index: currentIndex, - value: encodeBytes32Value(value), - nextIndex: 0n, - isActive: true, - }); - expect(await indexedMT.getNodeHash(currentIndex, LEAVES_LEVEL)).to.be.eq(leafHash); + compareProofs(proof, expectedProof, true); + }); - const lowLeafData = await indexedMT.getLeafData(lowLeafIndex); + it("should return correct inclusion proofs with the random tree elements", async () => { + const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); + const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); - expect(lowLeafData.value).to.be.eq(lowLeafValue); - expect(lowLeafData.nextLeafIndex).to.be.eq(currentIndex); + await newIndexedMT.initializeUintTree(); - const lowLeafNewHash = hashIndexedLeaf({ - index: lowLeafIndex, - value: encodeBytes32Value(lowLeafValue), - nextIndex: currentIndex, - isActive: true, - }); - expect(await indexedMT.getNodeHash(lowLeafIndex, LEAVES_LEVEL)).to.be.eq(lowLeafNewHash); + 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); + } - lowLeafIndex = currentIndex; - lowLeafValue = value; - value *= 2n; + const proofsCount = 100n; - expect(await indexedMT.getRoot()).to.be.eq(localIndexedMerkleTree.getRoot()); - } + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, 99); + const valueToProve = values[randIndex]; - const expectedLevelsCount = Math.ceil(Math.log2(count + 1)) + 1; + const index = newLocalIndexedMT.getLeafIndex(valueToProve); - expect(await indexedMT.getTreeLevels()).to.be.eq(expectedLevelsCount); + 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"); + 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 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); + }); }); - it("should correctly add 100 random elements", async () => { - const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); - const elementsCount = 100n; + describe("verifyProof", () => { + const values: bigint[] = [0n, 10n, 20n, 30n, 40n, 50n]; + let localIndexedMerkleTree: IndexedMerkleTree; + + beforeEach("setup", async () => { + localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); - for (let i = 0; i < elementsCount; ++i) { - const currentValue = ethers.hexlify(ethers.randomBytes(32)); - const lowLeafIndex = localIndexedMerkleTree.getLowLeafIndex(currentValue); + 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); + } + }); - const expectedNextLeafIndex = localIndexedMerkleTree.getLeafData(lowLeafIndex).nextLeafIndex; + it("should correctly verify inclusion proofs", async () => { + let index = 1n; + let proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); - const index = localIndexedMerkleTree.add(currentValue, lowLeafIndex); - await indexedMT.addUint(BigInt(currentValue), lowLeafIndex); + expect(await indexedMT.verifyProofUint(proof)).to.be.true; - expect(await indexedMT.getRoot()).to.be.eq(localIndexedMerkleTree.getRoot()); + index = 2n; + proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); - const leafData = await indexedMT.getLeafData(index); + expect(await indexedMT.verifyProofUint(proof)).to.be.true; - expect(leafData.value).to.be.eq(currentValue); - expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); + index = 5n; + proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); - expect((await indexedMT.getLeafData(lowLeafIndex)).nextLeafIndex).to.be.eq(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"); + 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"); + 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; + } + }); }); - it("should get exception if pass invalid low leaf index", async () => { - await indexedMT.addUint(10n, 0n); - await indexedMT.addUint(20n, 1n); + describe("getters", () => { + it("should get exception if pass invalid index", async () => { + const invalidIndex = 120n; - await expect(indexedMT.addUint(5n, 1n)).to.be.revertedWithCustomError(indexedMT, "InvalidLowLeaf"); - await expect(indexedMT.addUint(25n, 1n)).to.be.revertedWithCustomError(indexedMT, "InvalidLowLeaf"); + await expect(indexedMT.getLeafDataUint(invalidIndex)) + .to.be.revertedWithCustomError(indexedMT, "IndexOutOfBounds") + .withArgs(invalidIndex, LEAVES_LEVEL); + }); }); }); - describe("getProof", () => { - const values: bigint[] = [0n, 10n, 20n, 30n]; - let localIndexedMerkleTree: IndexedMerkleTree; - + describe("Bytes32IndexedMerkleTree", () => { beforeEach("setup", async () => { - localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + await indexedMT.initializeBytes32Tree(); + }); - for (let i = 1; i < values.length; i++) { - localIndexedMerkleTree.add(encodeBytes32Value(values[i])); + describe("initialize", () => { + it("should correctly initialize Bytes32IndexedMerkleTree", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + const zeroLeafHash = hashIndexedLeaf({ + 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); + }); - await indexedMT.addUint(values[i], i - 1); - } + it("should get exception if try to initialize twice", async () => { + await expect(indexedMT.initializeBytes32Tree()).to.be.revertedWithCustomError( + indexedMT, + "IndexedMerkleTreeAlreadyInitialized", + ); + }); }); - it("should return correct inclusion proof", async () => { - let index = 2n; - let value = 20n; - let expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); - let proof = await indexedMT.getProof(index, value); - - expect(proof.root).to.be.eq(expectedProof.root); - expect(proof.existence).to.be.true; - expect(proof.existence).to.be.eq(expectedProof.existence); - expect(proof.index).to.be.eq(expectedProof.index); - expect(proof.value).to.be.eq(expectedProof.value); - expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); - - index = 1n; - value = 10n; - expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); - proof = await indexedMT.getProof(index, value); - - expect(proof.root).to.be.eq(expectedProof.root); - expect(proof.existence).to.be.true; - expect(proof.existence).to.be.eq(expectedProof.existence); - expect(proof.index).to.be.eq(expectedProof.index); - expect(proof.value).to.be.eq(expectedProof.value); - expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); + describe("add", () => { + 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.addBytes32(currentValue, lowLeafIndex); + + expect(await indexedMT.getRootBytes32()).to.be.eq(localIndexedMerkleTree.getRoot()); + + const leafData = await indexedMT.getLeafDataBytes32(index); + + expect(leafData.value).to.be.eq(currentValue); + 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); + }); }); - it("should return correct inclusion proofs with the random tree elements", async () => { - const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); - const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + describe("getProof", () => { + it("should return correct exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(); - await newIndexedMT.initializeUintTree(); + const valuesCount = 100n; + const values = []; - const valuesCount = 100n; - const values = []; + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(currentValue); - for (let i = 0; i < valuesCount; ++i) { - const currentValue = ethers.hexlify(ethers.randomBytes(32)); - const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + localIndexedMT.add(currentValue, lowLeafIndex); + await indexedMT.addBytes32(currentValue, lowLeafIndex); - newLocalIndexedMT.add(currentValue, lowLeafIndex); - await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + values.push(currentValue); + } - values.push(currentValue); - } + const proofsCount = 100n; - const proofsCount = 100n; + for (let i = 0; i < proofsCount; i++) { + let valueToProve: string; - for (let i = 0; i < proofsCount; i++) { - const randIndex = getRandomIntInclusive(0, 99); - const valueToProve = values[randIndex]; + do { + valueToProve = ethers.hexlify(ethers.randomBytes(32)); + } while (values.includes(valueToProve)); - const index = newLocalIndexedMT.getLeafIndex(valueToProve); + const index = localIndexedMT.getLowLeafIndex(valueToProve); - const expectedProof = newLocalIndexedMT.getProof(index, valueToProve); - const proof = await newIndexedMT.getProof(index, valueToProve); + const expectedProof = localIndexedMT.getProof(index, valueToProve); + const proof = await indexedMT.getProofBytes32(index, valueToProve); - expect(proof.root).to.be.eq(expectedProof.root); - expect(proof.existence).to.be.true; - expect(proof.existence).to.be.eq(expectedProof.existence); - expect(proof.index).to.be.eq(expectedProof.index); - expect(proof.value).to.be.eq(expectedProof.value); - expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); - } + compareProofs(proof, expectedProof, false); + } + }); }); - it("should return correct exclusion proof", async () => { - const index = 1n; - const value = 15n; - const expectedProof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(value)); - const proof = await indexedMT.getProof(index, value); - - expect(proof.root).to.be.eq(await indexedMT.getRoot()); - expect(proof.existence).to.be.false; - expect(proof.existence).to.be.eq(expectedProof.existence); - expect(proof.index).to.be.eq(expectedProof.index); - expect(proof.value).to.be.eq(expectedProof.value); - expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); - }); + describe("verifyProof", () => { + it("should correctly verify inclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(); - it("should return correct exclusion proofs with the random tree elements", async () => { - const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); - const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + const valuesCount = 100n; + const values = []; - await newIndexedMT.initializeUintTree(); + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(currentValue); - const valuesCount = 100n; - const values = []; + localIndexedMT.add(currentValue, lowLeafIndex); + await indexedMT.addBytes32(currentValue, lowLeafIndex); - for (let i = 0; i < valuesCount; ++i) { - const currentValue = ethers.hexlify(ethers.randomBytes(32)); - const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + values.push(currentValue); + } - newLocalIndexedMT.add(currentValue, lowLeafIndex); - await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + const proofsCount = 100n; - values.push(currentValue); - } + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, 99); + const valueToProve = values[randIndex]; - const proofsCount = 100n; + const index = localIndexedMT.getLeafIndex(valueToProve); + const proof = localIndexedMT.getProof(index, valueToProve); - for (let i = 0; i < proofsCount; i++) { - let valueToProve: string; + expect(await indexedMT.verifyProofBytes32(proof)).to.be.true; + } + }); - do { - valueToProve = ethers.hexlify(ethers.randomBytes(32)); - } while (values.includes(valueToProve)); + it("should correctly verify exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(); - const index = newLocalIndexedMT.getLowLeafIndex(valueToProve); + const valuesCount = 100n; + const values = []; - const expectedProof = newLocalIndexedMT.getProof(index, valueToProve); - const proof = await newIndexedMT.getProof(index, valueToProve); + for (let i = 0; i < valuesCount; ++i) { + const currentValue = ethers.hexlify(ethers.randomBytes(32)); + const lowLeafIndex = localIndexedMT.getLowLeafIndex(currentValue); - expect(proof.root).to.be.eq(expectedProof.root); - expect(proof.existence).to.be.false; - expect(proof.existence).to.be.eq(expectedProof.existence); - expect(proof.index).to.be.eq(expectedProof.index); - expect(proof.value).to.be.eq(expectedProof.value); - expect(proof.nextLeafIndex).to.be.eq(expectedProof.nextLeafIndex); - expect(proof.siblings).to.be.deep.eq(expectedProof.siblings); - } - }); + localIndexedMT.add(currentValue, lowLeafIndex); + await indexedMT.addBytes32(currentValue, lowLeafIndex); + + values.push(currentValue); + } + + const proofsCount = 100n; - it("should get exception if pass invalid index", async () => { - const index = 1n; - const value = 5n; + for (let i = 0; i < proofsCount; i++) { + let valueToProve: string; - await expect(indexedMT.getProof(index, value)) - .to.be.revertedWithCustomError(indexedMT, "InvalidProofIndex") - .withArgs(index, value); + do { + valueToProve = ethers.hexlify(ethers.randomBytes(32)); + } while (values.includes(valueToProve)); + + const index = localIndexedMT.getLowLeafIndex(valueToProve); + const proof = localIndexedMT.getProof(index, valueToProve); + + expect(await indexedMT.verifyProofBytes32(proof)).to.be.true; + } + }); }); }); - describe("verifyProof", () => { - const values: bigint[] = [0n, 10n, 20n, 30n, 40n, 50n]; - let localIndexedMerkleTree: IndexedMerkleTree; - + describe("AddressIndexedMerkleTree", () => { beforeEach("setup", async () => { - localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + await indexedMT.initializeAddressTree(); + }); - for (let i = 1; i < values.length; i++) { - const lowIndex = localIndexedMerkleTree.getLowLeafIndex(encodeBytes32Value(values[i])); - localIndexedMerkleTree.add(encodeBytes32Value(values[i])); + describe("initialize", () => { + it("should correctly initialize AddressIndexedMerkleTree", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + const zeroLeafHash = hashIndexedLeaf({ + 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); + }); - await indexedMT.addUint(values[i], lowIndex); - } + it("should get exception if try to initialize twice", async () => { + await expect(indexedMT.initializeAddressTree()).to.be.revertedWithCustomError( + indexedMT, + "IndexedMerkleTreeAlreadyInitialized", + ); + }); }); - it("should correctly verify inclusion proofs", async () => { - let index = 1n; - let proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + describe("add", () => { + it("should correctly add 100 random elements", async () => { + const localIndexedMerkleTree = IndexedMerkleTree.buildMerkleTree(); + const elementsCount = 100n; + + 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; - expect(await indexedMT.verifyProof(proof)).to.be.true; + const index = localIndexedMerkleTree.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); - index = 2n; - proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + expect(await indexedMT.getRootAddress()).to.be.eq(localIndexedMerkleTree.getRoot()); - expect(await indexedMT.verifyProof(proof)).to.be.true; + const leafData = await indexedMT.getLeafDataAddress(index); - index = 5n; - proof = localIndexedMerkleTree.getProof(index, encodeBytes32Value(values[Number(index)])); + expect(leafData.value).to.be.eq(encodedAddress); + expect(leafData.nextLeafIndex).to.be.eq(expectedNextLeafIndex); - expect(await indexedMT.verifyProof(proof)).to.be.true; + expect((await indexedMT.getLeafDataAddress(lowLeafIndex)).nextLeafIndex).to.be.eq(index); + } + + expect(await indexedMT.getLevelNodesCountAddress(LEAVES_LEVEL)).to.be.eq(elementsCount + 1n); + }); }); - it("should correctly verify inclusion proofs with the random tree elements", async () => { - const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); - const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + describe("getProof", () => { + it("should return correct exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(); - await newIndexedMT.initializeUintTree(); + const valuesCount = 100n; + const values = []; - const valuesCount = 100n; - 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); - for (let i = 0; i < valuesCount; ++i) { - const currentValue = ethers.hexlify(ethers.randomBytes(32)); - const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + localIndexedMT.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); - newLocalIndexedMT.add(currentValue, lowLeafIndex); - await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + values.push(randomAddress); + } - values.push(currentValue); - } + const proofsCount = 100n; - const proofsCount = 100n; + for (let i = 0; i < proofsCount; i++) { + let addressToProve: string; - for (let i = 0; i < proofsCount; i++) { - const randIndex = getRandomIntInclusive(0, 99); - const valueToProve = values[randIndex]; + do { + addressToProve = ethers.hexlify(ethers.randomBytes(20)); + } while (values.includes(addressToProve)); - const index = newLocalIndexedMT.getLeafIndex(valueToProve); - const proof = newLocalIndexedMT.getProof(index, valueToProve); + const encodedAddressToProve = encodeAddressValue(addressToProve); - expect(await newIndexedMT.verifyProof(proof)).to.be.true; - } + const index = localIndexedMT.getLowLeafIndex(encodedAddressToProve); + + const expectedProof = localIndexedMT.getProof(index, encodedAddressToProve); + const proof = await indexedMT.getProofAddress(index, addressToProve); + + compareProofs(proof, expectedProof, false); + } + }); }); - it("should correctly verify exclusion proofs with the random tree elements", async () => { - const newIndexedMT = await ethers.deployContract("IndexedMerkleTreeMock"); - const newLocalIndexedMT = IndexedMerkleTree.buildMerkleTree(); + describe("verifyProof", () => { + it("should correctly verify inclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(); - await newIndexedMT.initializeUintTree(); + const valuesCount = 100n; + const values = []; - const valuesCount = 100n; - 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); - for (let i = 0; i < valuesCount; ++i) { - const currentValue = ethers.hexlify(ethers.randomBytes(32)); - const lowLeafIndex = newLocalIndexedMT.getLowLeafIndex(currentValue); + localIndexedMT.add(encodedAddress, lowLeafIndex); + await indexedMT.addAddress(randomAddress, lowLeafIndex); - newLocalIndexedMT.add(currentValue, lowLeafIndex); - await newIndexedMT.addUint(BigInt(currentValue), lowLeafIndex); + values.push(randomAddress); + } - values.push(currentValue); - } + const proofsCount = 100n; - const proofsCount = 100n; + for (let i = 0; i < proofsCount; i++) { + const randIndex = getRandomIntInclusive(0, 99); + const addressToProve = values[randIndex]; + const encodedAddressToProve = encodeAddressValue(addressToProve); - for (let i = 0; i < proofsCount; i++) { - let valueToProve: string; + const index = localIndexedMT.getLeafIndex(encodedAddressToProve); + const proof = localIndexedMT.getProof(index, encodedAddressToProve); + + expect(await indexedMT.verifyProofAddress(proof)).to.be.true; + } + }); - do { - valueToProve = ethers.hexlify(ethers.randomBytes(32)); - } while (values.includes(valueToProve)); + it("should correctly verify exclusion proofs with the random tree elements", async () => { + const localIndexedMT = IndexedMerkleTree.buildMerkleTree(); - const index = newLocalIndexedMT.getLowLeafIndex(valueToProve); - const proof = newLocalIndexedMT.getProof(index, valueToProve); + const valuesCount = 100n; + const values = []; - expect(await newIndexedMT.verifyProof(proof)).to.be.true; - } + 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 = 100n; + + 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; + } + }); }); }); });