From 8722d2bf4e57e79fc240a6f690c23d5e0b4c3056 Mon Sep 17 00:00:00 2001 From: Cedoor Date: Mon, 12 Aug 2024 11:21:51 +0200 Subject: [PATCH] refactor(lean-imt)!: update function to import exported trees (#313) * refactor(lean-imt)!: update function to import exported trees The import function was made static and a parameter was added to convert node types. re #312 * refactor(lean-imt): convert strings to bigints by default re #312 * refactor(lean-imt): make function to convert node types simpler re #312 * test(lean-imt): add more tests with different type conversions re #312 * docs(lean-imt): update description in import function re #312 --- packages/lean-imt/README.md | 29 ++++++----- packages/lean-imt/src/lean-imt.ts | 28 ++++++++--- packages/lean-imt/tests/lean-imt.test.ts | 63 ++++++++++++++++-------- 3 files changed, 81 insertions(+), 39 deletions(-) diff --git a/packages/lean-imt/README.md b/packages/lean-imt/README.md index f9e2d4c1a..4b13a2a41 100644 --- a/packages/lean-imt/README.md +++ b/packages/lean-imt/README.md @@ -84,9 +84,7 @@ or [JSDelivr](https://www.jsdelivr.com/): import { LeanIMT } from "@zk-kit/lean-imt" import { poseidon2 } from "poseidon-lite" -/** - * Hash function used to compute the tree nodes. - */ +// Hash function used to compute the tree nodes. const hash = (a, b) => poseidon2([a, b]) // To create an instance of a LeanIMT, you must provide the hash function. @@ -132,21 +130,28 @@ tree.update(1, 2n) // [1n, 2n] console.log(tree.leaves) -/** - * If you want to delete a leaf with LeanIMT you can use the update function with an - * arbitrary value to be used for the removed leaves. - */ - +// If you want to delete a leaf with LeanIMT you can use the update function with an +// arbitrary value to be used for the removed leaves. // Update the value of the leaf at position 1 to 0n (deletion). tree.update(1, 0n) // [1n, 0n] console.log(tree.leaves) -/** - * Compute a Merkle Inclusion Proof (proof of membership) for the leaf with index 1. - * The proof is only valid if the value 1 is found in a leaf of the tree. - */ +// Compute a Merkle Inclusion Proof (proof of membership) for the leaf with index 1. +// The proof is only valid if the value 1 is found in a leaf of the tree. const proof = tree.generateProof(1) // true console.log(tree.verifyProof(proof)) + +// Export all the tree nodes. +const nodes = tree.export() + +// Import the nodes. +const tree2 = LeanIMT.import(hash, nodes) + +// Node types are converted from strings to bigints by default. +// The third parameter can be used to convert strings to other types. + +// Import the nodes converting their types to numbers. +const tree3 = LeanIMT.import(hash, nodes, Number) ``` diff --git a/packages/lean-imt/src/lean-imt.ts b/packages/lean-imt/src/lean-imt.ts index 318f0a776..0a1f14d71 100644 --- a/packages/lean-imt/src/lean-imt.ts +++ b/packages/lean-imt/src/lean-imt.ts @@ -329,18 +329,34 @@ export default class LeanIMT { /** * It imports an entire tree by initializing the nodes without calculating * any hashes. Note that it is crucial to ensure the integrity of the tree - * before or after importing it. - * The tree must be empty before importing. + * before or after importing it. If the map function is not defined, node + * values will be converted to bigints by default. + * @param hash The hash function used to create nodes. * @param nodes The stringified JSON of the tree. + * @param map A function to map each node of the tree and convert their types. + * @returns A LeanIMT instance. */ - public import(nodes: string) { + static import(hash: LeanIMTHashFunction, nodes: string, map?: (value: string) => N): LeanIMT { + requireDefined(hash, "hash") requireDefined(nodes, "nodes") + requireFunction(hash, "hash") requireString(nodes, "nodes") - if (this.size !== 0) { - throw new Error("Import failed: the target tree structure is not empty") + if (map) { + requireDefined(map, "map") + requireFunction(map, "map") } - this._nodes = JSON.parse(nodes) + const tree = new LeanIMT(hash) + + tree._nodes = JSON.parse(nodes, (_, value) => { + if (typeof value === "string") { + return map ? map(value) : BigInt(value) + } + + return value + }) + + return tree } } diff --git a/packages/lean-imt/tests/lean-imt.test.ts b/packages/lean-imt/tests/lean-imt.test.ts index d1ee5ec21..0fb3c57e6 100644 --- a/packages/lean-imt/tests/lean-imt.test.ts +++ b/packages/lean-imt/tests/lean-imt.test.ts @@ -335,47 +335,67 @@ describe("Lean IMT", () => { it("Should export a tree", () => { const tree = new LeanIMT(poseidon, leaves) - const exportedTree = tree.export() + const nodes = tree.export() - expect(typeof exportedTree).toBe("string") - expect(JSON.parse(exportedTree)).toHaveLength(4) - expect(JSON.parse(exportedTree)[0]).toHaveLength(5) + expect(typeof nodes).toBe("string") + expect(JSON.parse(nodes)).toHaveLength(4) + expect(JSON.parse(nodes)[0]).toHaveLength(5) }) - it("Should not import a tree if it the exported tree is not defined", () => { + it("Should not import a tree if the required parameters are not valid", () => { const tree = new LeanIMT(poseidon, leaves) + const nodes = tree.export() - const fun = () => tree.import(undefined as any) + const fun1 = () => LeanIMT.import(undefined as any, nodes) + const fun2 = () => LeanIMT.import(poseidon, undefined as any) + const fun3 = () => LeanIMT.import("string" as any, nodes) + const fun4 = () => LeanIMT.import(poseidon, 1 as any) + const fun5 = () => LeanIMT.import(poseidon, nodes, "string" as any) - expect(fun).toThrow("Parameter 'nodes' is not defined") + expect(fun1).toThrow("Parameter 'hash' is not defined") + expect(fun2).toThrow("Parameter 'nodes' is not defined") + expect(fun3).toThrow("Parameter 'hash' is not a function") + expect(fun4).toThrow("Parameter 'nodes' is not a string") + expect(fun5).toThrow("Parameter 'map' is not a function") }) - it("Should not import a tree if it the exported tree is not a string", () => { - const tree = new LeanIMT(poseidon, leaves) + it("Should import a tree converting node types to number", () => { + const hash = (a: number, b: number) => a + b + const tree1 = new LeanIMT(hash, [1, 2, 3]) + const nodes = tree1.export() + + const tree2 = LeanIMT.import(hash, nodes, Number) - const fun = () => tree.import(1 as any) + tree1.insert(4) + tree2.insert(4) - expect(fun).toThrow("Parameter 'nodes' is not a string") + expect(tree2.depth).toBe(tree1.depth) + expect(tree2.size).toBe(tree1.size) + expect(tree2.root).toBe(tree1.root) + expect(tree2.indexOf(2)).toBe(tree1.indexOf(2)) }) - it("Should not import a tree if it is not empty", () => { - const tree1 = new LeanIMT(poseidon, leaves) - const exportedTree = tree1.export() + it("Should import a tree converting node types to booleans", () => { + const hash = (a: boolean, b: boolean) => a && b + const tree1 = new LeanIMT(hash, [true, false, true]) + const nodes = tree1.export() - const tree2 = new LeanIMT(poseidon, leaves) + const tree2 = LeanIMT.import(hash, nodes, Boolean) - const fun = () => tree2.import(exportedTree) + tree1.insert(true) + tree2.insert(true) - expect(fun).toThrow("Import failed: the target tree structure is not empty") + expect(tree2.depth).toBe(tree1.depth) + expect(tree2.size).toBe(tree1.size) + expect(tree2.root).toBe(tree1.root) + expect(tree2.indexOf(false)).toBe(tree1.indexOf(false)) }) it("Should import a tree", () => { const tree1 = new LeanIMT(poseidon, leaves) - const exportedTree = tree1.export() - - const tree2 = new LeanIMT(poseidon) + const nodes = tree1.export() - tree2.import(exportedTree) + const tree2 = LeanIMT.import(poseidon, nodes) tree1.insert(BigInt(4)) tree2.insert(BigInt(4)) @@ -383,6 +403,7 @@ describe("Lean IMT", () => { expect(tree2.depth).toBe(tree1.depth) expect(tree2.size).toBe(tree1.size) expect(tree2.root).toBe(tree1.root) + expect(tree2.indexOf(2n)).toBe(tree1.indexOf(2n)) }) }) })