fix: lazy imt depth overflow (#138)

* fix(imt.sol): lazy imt depth overflow

* fix(imt.sol): typo in error message

* refactor(imt.sol): use uint40 everywhere instead of uint32

* fix(imt.sol): test error case

* test: failing test for imt tree root

* chore(imt.sol): add assertion to test

---------

Co-authored-by: alrevuelta <alvrevuelta@gmail.com>

packages/imt.sol/contracts/internal/InternalLazyIMT.sol

@@ -5,7 +5,7 @@ import {PoseidonT3} from "poseidon-solidity/PoseidonT3.sol";
 import {SNARK_SCALAR_FIELD, MAX_DEPTH} from "../Constants.sol";
 
 struct LazyIMTData {
-    uint32 maxIndex;
+    uint40 maxIndex;
     uint40 numberOfLeaves;
     mapping(uint256 => uint256) elements;
 }
@@ -86,7 +86,7 @@ library InternalLazyIMT {
 
     function _init(LazyIMTData storage self, uint8 depth) internal {
         require(depth <= MAX_DEPTH, "LazyIMT: Tree too large");
-        self.maxIndex = uint32((1 << depth) - 1);
+        self.maxIndex = uint40((1 << depth) - 1);
         self.numberOfLeaves = 0;
     }
 
@@ -151,21 +151,24 @@ library InternalLazyIMT {
         uint40 numberOfLeaves = self.numberOfLeaves;
         // dynamically determine a depth
         uint8 depth = 1;
-        while (uint8(2) ** depth < numberOfLeaves) {
+        while (uint40(2) ** uint40(depth) < numberOfLeaves) {
             depth++;
         }
         return _root(self, numberOfLeaves, depth);
     }
 
     function _root(LazyIMTData storage self, uint8 depth) internal view returns (uint256) {
+        require(depth > 0, "LazyIMT: depth must be > 0");
+        require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH");
         uint40 numberOfLeaves = self.numberOfLeaves;
-        require(2 ** depth >= numberOfLeaves, "LazyIMT: ambiguous depth");
-        return _root(self, self.numberOfLeaves, depth);
+        require(uint40(2) ** uint40(depth) >= numberOfLeaves, "LazyIMT: ambiguous depth");
+        return _root(self, numberOfLeaves, depth);
     }
 
+    // Here it's assumed that the depth value is valid. If it is either 0 or > 2^8-1
+    // this function will panic.
     function _root(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth) internal view returns (uint256) {
-        require(depth > 0, "LazyIMT: depth must be > 0");
-        require(depth <= MAX_DEPTH, "LazyIMT: depth must be < MAX_DEPTH");
+        require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH");
         // this should always short circuit if self.numberOfLeaves == 0
         if (numberOfLeaves == 0) return _defaultZero(depth);
         uint40 index = numberOfLeaves - 1;

packages/imt.sol/test/LazyIMT.ts

@@ -113,6 +113,22 @@ describe("LazyIMT", () => {
             })
         }
 
+        it("Should insert multiple leaves", async () => {
+            const depth = 8
+
+            const merkleTree = new IMT(poseidon2, depth, BigInt(0))
+            await lazyIMTTest.init(depth)
+
+            for (let x = 0; x < 130; x += 1) {
+                const e = random()
+                await lazyIMTTest.insert(e)
+                merkleTree.insert(e)
+            }
+
+            const root = await lazyIMTTest.root()
+            expect(root.toString()).to.equal(merkleTree.root.toString())
+        })
+
         it("Should fail to insert too many leaves", async () => {
             const depth = 3
 
@@ -283,6 +299,6 @@ describe("LazyIMT", () => {
             await lazyIMTTest.insert(e)
         }
         await expect(lazyIMTTest.staticRoot(4)).to.be.revertedWith("LazyIMT: ambiguous depth")
-        await expect(lazyIMTTest.staticRoot(33)).to.be.revertedWith("LazyIMT: depth must be < MAX_DEPTH")
+        await expect(lazyIMTTest.staticRoot(33)).to.be.revertedWith("LazyIMT: depth must be <= MAX_DEPTH")
     })
 })