Implemented updateMany method as required in issue #117 (#314)

* feat(lean-imt): added `updateMany` method to package re #117 * feat(lean-imt): implemented some tests on lean-imt re #117 * feat(lean-imt): added more precondition checks re #117 * feat(lean-imt): finished testing on `updateMany` method re #117 * feat(lean-imt): added test to the case when passing repeated indices re #117 * feat(lean-imt): added complexity documentation for `updateMany` method re #117 * feat(lean-imt): added test of several updates re #117 * feat(lean-imt): added repeated indices check re #117 * feat(lean-imt): changed error message to be more accurate re #117 * feat(lean-imt): added complexity in terms only of n re #117 * feat(lean-imt): changed documentation to add discussion in another issue re #117 * feat(lean-imt): fixed typo on documentation re #117 * Update packages/lean-imt/src/lean-imt.ts Co-authored-by: Vivian Plasencia <[email protected]> --------- Co-authored-by: Vivian Plasencia <[email protected]>
privacy-scaling-explorations · Sep 9, 2024 · f910628 · f910628
1 parent 253277c
commit f910628
Show file tree

Hide file tree

Showing 2 changed files with 167 additions and 0 deletions.
diff --git a/packages/lean-imt/src/lean-imt.ts b/packages/lean-imt/src/lean-imt.ts
@@ -233,6 +233,59 @@ export default class LeanIMT<N = bigint> {
         this._nodes[this.depth] = [node]
     }
 
+    /**
+     * Updates m leaves all at once.
+     * It is more efficient than using the {@link LeanIMT#update} method m times because it
+     * prevents updating middle nodes several times. This would happen when updating leaves
+     * with common ancestors. The naive approach of calling 'update' m times has complexity
+     * O(m*log(n)) (where n is the number of leaves of the tree), which ends up in
+     * O(n*log(n)) when m ~ n. With this new approach, this ends up being O(n) because every
+     * node is updated at most once and there are around 2*n nodes in the tree.
+     * @param indices The list of indices of the respective leaves.
+     * @param leaves The list of leaves to be updated.
+     */
+    public updateMany(indices: number[], leaves: N[]) {
+        requireDefined(leaves, "leaves")
+        requireDefined(indices, "indices")
+        requireArray(leaves, "leaves")
+        requireArray(indices, "indices")
+
+        if (leaves.length !== indices.length) {
+            throw new Error("There is no correspondence between indices and leaves")
+        }
+        // This will keep track of the outdated nodes of each level.
+        let modifiedIndices = new Set<number>()
+        for (let i = 0; i < indices.length; i += 1) {
+            requireNumber(indices[i], `index ${i}`)
+            if (indices[i] < 0 || indices[i] >= this.size) {
+                throw new Error(`Index ${i} is out of range`)
+            }
+            if (modifiedIndices.has(indices[i])) {
+                throw new Error(`Leaf ${indices[i]} is repeated`)
+            }
+            modifiedIndices.add(indices[i])
+        }
+
+        modifiedIndices.clear()
+        // First, modify the first level, which consists only of raw, un-hashed values
+        for (let leaf = 0; leaf < indices.length; leaf += 1) {
+            this._nodes[0][indices[leaf]] = leaves[leaf]
+            modifiedIndices.add(indices[leaf] >> 1)
+        }
+
+        // Now update each node of the corresponding levels
+        for (let level = 1; level <= this.depth; level += 1) {
+            const newModifiedIndices: number[] = []
+            for (const index of modifiedIndices) {
+                const leftChild = this._nodes[level - 1][2 * index]
+                const rightChild = this._nodes[level - 1][2 * index + 1]
+                this._nodes[level][index] = rightChild ? this._hash(leftChild, rightChild) : leftChild
+                newModifiedIndices.push(index >> 1)
+            }
+            modifiedIndices = new Set<number>(newModifiedIndices)
+        }
+    }
+
     /**
      * It generates a {@link LeanIMTMerkleProof} for a leaf of the tree.
      * That proof can be verified by this tree using the same hash function.

diff --git a/packages/lean-imt/tests/lean-imt.test.ts b/packages/lean-imt/tests/lean-imt.test.ts
@@ -220,6 +220,120 @@ describe("Lean IMT", () => {
         })
     })
 
+    describe("# updateMany", () => {
+        it(`Should not update any leaf if one of the parameters is not defined`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const fun1 = () => tree.updateMany([1], undefined as any)
+            const fun2 = () => tree.updateMany(undefined as any, [BigInt(1)])
+
+            expect(fun1).toThrow("Parameter 'leaves' is not defined")
+            expect(fun2).toThrow("Parameter 'indices' is not defined")
+        })
+
+        it(`Should not update any leaf if the parameters are not arrays`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const fun1 = () => tree.updateMany([3], BigInt(1) as any)
+            const fun2 = () => tree.updateMany(3 as any, [BigInt(1)])
+
+            expect(fun1).toThrow("Parameter 'leaves' is not an Array instance")
+            expect(fun2).toThrow("Parameter 'indices' is not an Array instance")
+        })
+
+        it(`Should not update any leaf if the parameters are of different size`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const fun1 = () => tree.updateMany([1, 2, 3], [BigInt(1), BigInt(2)])
+            const fun2 = () => tree.updateMany([1], [])
+
+            expect(fun1).toThrow("There is no correspondence between indices and leaves")
+            expect(fun2).toThrow("There is no correspondence between indices and leaves")
+        })
+
+        it(`Should not update any leaf if some index is not a number`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const fun1 = () => tree.updateMany([1, "hello" as any, 3], [BigInt(1), BigInt(2), BigInt(3)])
+            const fun2 = () => tree.updateMany([1, 2, undefined as any], [BigInt(1), BigInt(2), BigInt(3)])
+
+            expect(fun1).toThrow("Parameter 'index 1' is not a number")
+            expect(fun2).toThrow("Parameter 'index 2' is not a number")
+        })
+
+        it(`Should not update any leaf if some index is out of range`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const fun1 = () => tree.updateMany([-1, 2, 3], [BigInt(1), BigInt(2), BigInt(3)])
+            const fun2 = () => tree.updateMany([1, 200000, 3], [BigInt(1), BigInt(2), BigInt(3)])
+            const fun3 = () => tree.updateMany([1, 2, tree.size], [BigInt(1), BigInt(2), BigInt(3)])
+
+            expect(fun1).toThrow("Index 0 is out of range")
+            expect(fun2).toThrow("Index 1 is out of range")
+            expect(fun3).toThrow("Index 2 is out of range")
+        })
+
+        it(`Should not update any leaf when passing an empty list`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+            const previousRoot = tree.root
+
+            tree.updateMany([], [])
+
+            expect(tree.root).toBe(previousRoot)
+        })
+
+        it(`'updateMany' with 1 change should be the same as 'update'`, () => {
+            const tree1 = new LeanIMT(poseidon, leaves)
+            const tree2 = new LeanIMT(poseidon, leaves)
+
+            tree1.update(4, BigInt(-100))
+            tree2.updateMany([4], [BigInt(-100)])
+            expect(tree1.root).toBe(tree2.root)
+
+            tree1.update(0, BigInt(24))
+            tree2.updateMany([0], [BigInt(24)])
+            expect(tree1.root).toBe(tree2.root)
+        })
+
+        it(`'updateMany' should be the same as executing the 'update' function multiple times`, () => {
+            const tree1 = new LeanIMT(poseidon, leaves)
+            const tree2 = new LeanIMT(poseidon, leaves)
+
+            const indices = [0, 2, 4]
+
+            const nodes = [BigInt(10), BigInt(11), BigInt(12)]
+
+            for (let i = 0; i < indices.length; i += 1) {
+                tree1.update(indices[i], nodes[i])
+            }
+            tree2.updateMany(indices, nodes)
+
+            expect(tree1.root).toBe(tree2.root)
+        })
+
+        it(`'updateMany' with repeated indices should fail`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const fun = () => tree.updateMany([4, 1, 4], [BigInt(-100), BigInt(-17), BigInt(1)])
+
+            expect(fun).toThrow("Leaf 4 is repeated")
+        })
+
+        it(`Should update leaves correctly`, () => {
+            const tree = new LeanIMT(poseidon, leaves)
+
+            const updateLeaves = [BigInt(24), BigInt(-10), BigInt(100000)]
+            tree.updateMany([0, 1, 4], updateLeaves)
+
+            const h1_0 = poseidon(updateLeaves[0], updateLeaves[1])
+            const h1_1 = poseidon(leaves[2], leaves[3])
+            const h2_0 = poseidon(h1_0, h1_1)
+            const updatedRoot = poseidon(h2_0, updateLeaves[2])
+
+            expect(tree.root).toBe(updatedRoot)
+        })
+    })
+
     describe("# generateProof", () => {
         it(`Should not generate any proof if the index is not defined`, () => {
             const tree = new LeanIMT(poseidon, leaves)