feat: add get merkle proof to lazy imt tree (privacy-scaling-explorat…

…ions#162)

* feat: add get merkle proof to lazy imt tree

* chore: run prettier

* chore: fix lint problems

* chore: remove proof index + misc changes

* chore: run prettier

* chore: move revert inside loop

* chore: take snapshoot in each test

* Lazy merkle tree merkle proof logic (#1)

* chore: fix eslint problems

---------

Co-authored-by: Chance <jchancehud@gmail.com>

packages/imt.sol/contracts/LazyIMT.sol

@@ -38,6 +38,14 @@ library LazyIMT {
         return InternalLazyIMT._root(self, depth);
     }
 
+    function merkleProofElements(
+        LazyIMTData storage self,
+        uint40 index,
+        uint8 depth
+    ) public view returns (uint256[] memory) {
+        return InternalLazyIMT._merkleProofElements(self, index, depth);
+    }
+
     function _root(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth) internal view returns (uint256) {
         return InternalLazyIMT._root(self, numberOfLeaves, depth);
     }

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

@@ -171,9 +171,21 @@ library InternalLazyIMT {
         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;
+        uint256[] memory levels = new uint256[](depth + 1);
+        _levels(self, numberOfLeaves, depth, levels);
+        return levels[depth];
+    }
 
-        uint256[MAX_DEPTH + 1] memory levels;
+    function _levels(
+        LazyIMTData storage self,
+        uint40 numberOfLeaves,
+        uint8 depth,
+        uint256[] memory levels
+    ) internal view {
+        require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH");
+        require(numberOfLeaves > 0, "LazyIMT: number of leaves must be > 0");
+        // this should always short circuit if self.numberOfLeaves == 0
+        uint40 index = numberOfLeaves - 1;
 
         if (index & 1 == 0) {
             levels[0] = self.elements[_indexForElement(0, index)];
@@ -199,6 +211,57 @@ library InternalLazyIMT {
                 i++;
             }
         }
-        return levels[depth];
+    }
+
+    function _merkleProofElements(
+        LazyIMTData storage self,
+        uint40 index,
+        uint8 depth
+    ) internal view returns (uint256[] memory) {
+        uint40 numberOfLeaves = self.numberOfLeaves;
+        require(index < numberOfLeaves, "LazyIMT: leaf must exist");
+
+        uint8 targetDepth = 1;
+        while (uint40(2) ** uint40(targetDepth) < numberOfLeaves) {
+            targetDepth++;
+        }
+        require(depth >= targetDepth, "LazyIMT: proof depth");
+        // pass depth -1 because we don't need the root value
+        uint256[] memory _elements = new uint256[](depth);
+        _levels(self, numberOfLeaves, targetDepth - 1, _elements);
+
+        // unroll the bottom entry of the tree because it will never need to
+        // be pulled from _levels
+        if (index & 1 == 0) {
+            if (index + 1 >= numberOfLeaves) {
+                _elements[0] = _defaultZero(0);
+            } else {
+                _elements[0] = self.elements[_indexForElement(0, index + 1)];
+            }
+        } else {
+            _elements[0] = self.elements[_indexForElement(0, index - 1)];
+        }
+        index >>= 1;
+
+        for (uint8 i = 1; i < depth; ) {
+            uint256 currentLevelCount = numberOfLeaves >> i;
+            if (index & 1 == 0) {
+                // if the element is an uncomputed edge node we'll use the value set
+                // from _levels above
+                // otherwise set as usual below
+                if (index + 1 < currentLevelCount) {
+                    _elements[i] = self.elements[_indexForElement(i, index + 1)];
+                } else if (((numberOfLeaves - 1) >> i) <= index) {
+                    _elements[i] = _defaultZero(i);
+                }
+            } else {
+                _elements[i] = self.elements[_indexForElement(i, index - 1)];
+            }
+            unchecked {
+                index >>= 1;
+                i++;
+            }
+        }
+        return _elements;
     }
 }

packages/imt.sol/contracts/test/LazyIMTTest.sol

@@ -40,4 +40,8 @@ contract LazyIMTTest {
     function staticRoot(uint8 depth) public view returns (uint256) {
         return LazyIMT.root(data, depth);
     }
+
+    function merkleProofElements(uint40 index, uint8 depth) public view returns (uint256[] memory) {
+        return LazyIMT.merkleProofElements(data, index, depth);
+    }
 }

packages/imt.sol/test/LazyIMT.ts

@@ -1,7 +1,8 @@
 import { expect } from "chai"
-import { run } from "hardhat"
+import { run, network } from "hardhat"
 import { poseidon2 } from "poseidon-lite"
 import { IMT } from "@zk-kit/imt"
+import type { BigNumber } from "ethers"
 import { LazyIMT, LazyIMTTest } from "../typechain-types"
 
 const random = () => poseidon2([Math.floor(Math.random() * 2 ** 40), 0])
@@ -289,6 +290,74 @@ describe("LazyIMT", () => {
         })
     })
 
+    describe("# merkleProof", () => {
+        // Given a a merkle proof (elements and indices) and a leaf, calculates the root
+        function calculateRoot(leafIndex: number, leaf: BigNumber, proofElements: BigNumber[]) {
+            let hash = leaf
+            const proofIndices = []
+            for (let x = 0; x < proofElements.length; x += 1) {
+                proofIndices.push((leafIndex >> x) & 1)
+            }
+            for (let i = 0; i < proofElements.length; i += 1) {
+                const proofElement = proofElements[i]
+                const proofIndex = proofIndices[i]
+                if (proofIndex) {
+                    hash = poseidon2([proofElement.toString(), hash.toString()])
+                } else {
+                    hash = poseidon2([hash.toString(), proofElement.toString()])
+                }
+            }
+            return hash
+        }
+
+        it("Should produce valid Merke proofs for different trees", async () => {
+            // Test different depths (key) and leafs (values)
+            const tests: { [key: number]: number[] } = {
+                1: [1],
+                2: [1, 2],
+                5: [6, 7, 8, 15, 16],
+                7: [7, 127],
+                20: [9, 14, 15, 16, 18, 26, 27, 28, 40, 128, 129]
+            }
+
+            // For each depth
+            // eslint-disable-next-line guard-for-in
+            for (const depth in tests) {
+                // For each amount of leafs
+                for (const numLeaf of tests[depth]) {
+                    // Freeze the state
+                    const snapshoot = await network.provider.request({ method: "evm_snapshot", params: [] })
+
+                    // Create the tree
+                    await lazyIMTTest.init(depth)
+                    const elements = []
+                    for (let x = 0; x < numLeaf; x += 1) {
+                        const e = random()
+                        elements.push(e)
+                        await lazyIMTTest.insert(e)
+                    }
+
+                    // Get proofs for every leafs and verify against the root
+                    for (let leafIndex = 0; leafIndex < numLeaf; leafIndex += 1) {
+                        const proofElements = await lazyIMTTest.merkleProofElements(leafIndex, depth)
+
+                        // Calculate the root we arrive at with the proof we got
+                        const calculatedRoot = calculateRoot(leafIndex, elements[leafIndex], proofElements)
+
+                        // Get the root from the contract
+                        const staticRoot = await lazyIMTTest.staticRoot(depth)
+
+                        // If they match, proof is valid
+                        await expect(calculatedRoot).to.be.equal(staticRoot)
+                    }
+
+                    // Done with test, revert the tree state
+                    await network.provider.request({ method: "evm_revert", params: [snapshoot] })
+                }
+            }
+        }).timeout(5 * 60 * 1000)
+    })
+
     it("Should fail to generate out of range static root", async () => {
         await lazyIMTTest.init(10)