diff --git a/packages/imt.sol/contracts/LazyIMT.sol b/packages/imt.sol/contracts/LazyIMT.sol index 71d869dbc..32084ad54 100644 --- a/packages/imt.sol/contracts/LazyIMT.sol +++ b/packages/imt.sol/contracts/LazyIMT.sol @@ -1,7 +1,6 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.4; -import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; import {InternalLazyIMT, LazyIMTData} from "./internal/InternalLazyIMT.sol"; library LazyIMT { diff --git a/packages/imt.sol/contracts/LeanIMT.sol b/packages/imt.sol/contracts/LeanIMT.sol index 1382c1151..b3753c69f 100644 --- a/packages/imt.sol/contracts/LeanIMT.sol +++ b/packages/imt.sol/contracts/LeanIMT.sol @@ -1,7 +1,6 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.4; -import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol"; import {InternalLeanIMT, LeanIMTData} from "./internal/InternalLeanIMT.sol"; library LeanIMT { @@ -11,6 +10,10 @@ library LeanIMT { return InternalLeanIMT._insert(self, leaf); } + function insertMany(LeanIMTData storage self, uint256[] calldata leaves) public returns (uint256) { + return InternalLeanIMT._insertMany(self, leaves); + } + function update( LeanIMTData storage self, uint256 oldLeaf, diff --git a/packages/imt.sol/contracts/QuinaryIMT.sol b/packages/imt.sol/contracts/QuinaryIMT.sol index ae0849b11..e549108b5 100644 --- a/packages/imt.sol/contracts/QuinaryIMT.sol +++ b/packages/imt.sol/contracts/QuinaryIMT.sol @@ -1,7 +1,6 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.4; -import {PoseidonT6} from "poseidon-solidity/PoseidonT6.sol"; import {InternalQuinaryIMT, QuinaryIMTData} from "./internal/InternalQuinaryIMT.sol"; library QuinaryIMT { diff --git a/packages/imt.sol/contracts/internal/InternalLeanIMT.sol b/packages/imt.sol/contracts/internal/InternalLeanIMT.sol index 063e7e685..f7637114d 100644 --- a/packages/imt.sol/contracts/internal/InternalLeanIMT.sol +++ b/packages/imt.sol/contracts/internal/InternalLeanIMT.sol @@ -73,6 +73,132 @@ library InternalLeanIMT { return node; } + /// @dev Inserts many leaves into the incremental merkle tree. + /// The function ensures that the leaves are valid according to the + /// constraints of the tree and then updates the tree's structure accordingly. + /// @param self: A storage reference to the 'LeanIMTData' struct. + /// @param leaves: The values of the new leaves to be inserted into the tree. + /// @return The root after the leaves have been inserted. + function _insertMany(LeanIMTData storage self, uint256[] calldata leaves) internal returns (uint256) { + // Check that all the new values are correct to be added. + for (uint256 i = 0; i < leaves.length; ) { + if (leaves[i] >= SNARK_SCALAR_FIELD) { + revert LeafGreaterThanSnarkScalarField(); + } else if (leaves[i] == 0) { + revert LeafCannotBeZero(); + } else if (_has(self, leaves[i])) { + revert LeafAlreadyExists(); + } + + self.leaves[leaves[i]] = self.size + i; + + unchecked { + ++i; + } + } + + // Array to save the nodes that will be used to create the next level of the tree. + uint256[] memory currentLevel; + + currentLevel = leaves; + + // Calculate the depth of the tree after adding the new values. + while (2 ** self.depth < self.size + leaves.length) { + self.depth += 1; + } + + // First index to change in every level. + uint256 currentLevelStartIndex = self.size; + + // Size of the level used to create the next level. + uint256 currentLevelSize = self.size + leaves.length; + + // The index where changes begin at the next level. + uint256 nextLevelStartIndex = currentLevelStartIndex >> 1; + + // The size of the next level. + uint256 nextLevelSize = ((currentLevelSize - 1) >> 1) + 1; + + for (uint256 level = 0; level < self.depth; ) { + // The number of nodes for the new level that will be created, + // only the new values, not the entire level. + uint256 numberOfNodes = nextLevelSize - nextLevelStartIndex; + uint256[] memory nextLevel = new uint256[](numberOfNodes); + for (uint256 i = 0; i < numberOfNodes; ) { + uint256 rightNode; + uint256 leftNode; + + // Assign the right node if the value exists. + if ((i + nextLevelStartIndex) * 2 + 1 < currentLevelSize) { + rightNode = currentLevel[(i + nextLevelStartIndex) * 2 + 1 - currentLevelStartIndex]; + } + + // Assign the left node using the saved path or the position in the array. + if ((i + nextLevelStartIndex) * 2 < currentLevelStartIndex) { + leftNode = self.sideNodes[level]; + } else { + leftNode = currentLevel[(i + nextLevelStartIndex) * 2 - currentLevelStartIndex]; + } + + uint256 parentNode; + + // Assign the parent node. + // If it has a right child the result will be the hash(leftNode, rightNode) if not, + // it will be the leftNode. + if (rightNode != 0) { + parentNode = PoseidonT3.hash([leftNode, rightNode]); + } else { + parentNode = leftNode; + } + + nextLevel[i] = parentNode; + + unchecked { + ++i; + } + } + + // Update the `sideNodes` variable. + // If `currentLevelSize` is odd, the saved value will be the last value of the array + // if it is even and there are more than 1 element in `currentLevel`, the saved value + // will be the value before the last one. + // If it is even and there is only one element, there is no need to save anything because + // the correct value for this level was already saved before. + if (currentLevelSize & 1 == 1) { + self.sideNodes[level] = currentLevel[currentLevel.length - 1]; + } else if (currentLevel.length > 1) { + self.sideNodes[level] = currentLevel[currentLevel.length - 2]; + } + + currentLevelStartIndex = nextLevelStartIndex; + + // Calculate the next level startIndex value. + // It is the position of the parent node which is pos/2. + nextLevelStartIndex >>= 1; + + // Update the next array that will be used to calculate the next level. + currentLevel = nextLevel; + + currentLevelSize = nextLevelSize; + + // Calculate the size of the next level. + // The size of the next level is (currentLevelSize - 1) / 2 + 1. + nextLevelSize = ((nextLevelSize - 1) >> 1) + 1; + + unchecked { + ++level; + } + } + + // Update tree size + self.size += leaves.length; + + // Update tree root + self.sideNodes[self.depth] = currentLevel[0]; + + return currentLevel[0]; + } + /// @dev Updates the value of an existing leaf and recalculates hashes /// to maintain tree integrity. /// @param self: A storage reference to the 'LeanIMTData' struct. diff --git a/packages/imt.sol/contracts/test/LeanIMTTest.sol b/packages/imt.sol/contracts/test/LeanIMTTest.sol index 9f0d99dfd..b39f8d1ca 100644 --- a/packages/imt.sol/contracts/test/LeanIMTTest.sol +++ b/packages/imt.sol/contracts/test/LeanIMTTest.sol @@ -11,6 +11,10 @@ contract LeanIMTTest { LeanIMT.insert(data, leaf); } + function insertMany(uint256[] calldata leaves) external { + LeanIMT.insertMany(data, leaves); + } + function update(uint256 oldLeaf, uint256 newLeaf, uint256[] calldata siblingNodes) external { LeanIMT.update(data, oldLeaf, newLeaf, siblingNodes); } diff --git a/packages/imt.sol/test/LeanIMT.ts b/packages/imt.sol/test/LeanIMT.ts index 80c49370d..3c4f90c0b 100644 --- a/packages/imt.sol/test/LeanIMT.ts +++ b/packages/imt.sol/test/LeanIMT.ts @@ -64,6 +64,66 @@ describe("LeanIMT", () => { }) }) + describe("# insertMany", () => { + it("Should not insert a leaf if its value is > SNARK_SCALAR_FIELD", async () => { + const transaction = leanIMTTest.insertMany([SNARK_SCALAR_FIELD]) + + await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafGreaterThanSnarkScalarField") + }) + + it("Should not insert a leaf if it is 0", async () => { + const leaf = 0 + + const transaction = leanIMTTest.insertMany([leaf]) + + await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafCannotBeZero") + }) + it("Should not insert a leaf if it was already inserted before", async () => { + await leanIMTTest.insert(1) + + const transaction = leanIMTTest.insertMany([1]) + + await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafAlreadyExists") + }) + it("Should insert a leaf", async () => { + jsLeanIMT.insert(BigInt(1)) + + await leanIMTTest.insertMany([1]) + + const root = await leanIMTTest.root() + + expect(root).to.equal(jsLeanIMT.root) + }) + it("Should insert 10 leaves", async () => { + const elems: bigint[] = [] + for (let i = 0; i < 10; i += 1) { + elems.push(BigInt(i + 1)) + } + + jsLeanIMT.insertMany(elems) + await leanIMTTest.insertMany(elems) + + const root = await leanIMTTest.root() + expect(root).to.equal(jsLeanIMT.root) + }) + it("Should insert many leaves when the tree is not empty", async () => { + jsLeanIMT.insert(BigInt(1)) + + await leanIMTTest.insert(BigInt(1)) + + const elems: bigint[] = [] + for (let i = 1; i < 10; i += 1) { + elems.push(BigInt(i + 1)) + } + + jsLeanIMT.insertMany(elems) + await leanIMTTest.insertMany(elems) + + const root = await leanIMTTest.root() + expect(root).to.equal(jsLeanIMT.root) + }) + }) + describe("# update", () => { it("Should not update a leaf if the leaf does not exist", async () => { const transaction = leanIMTTest.update(2, 1, [1, 2, 3, 4])