Add insertMany function in LeanIMT

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 {

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,

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 {

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.

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);
     }

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])