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

@@ -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<number>(hash, nodes, Number)
 ```

packages/lean-imt/src/lean-imt.ts

@@ -329,18 +329,34 @@ export default class LeanIMT<N = bigint> {
     /**
      * 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<N = bigint>(hash: LeanIMTHashFunction<N>, nodes: string, map?: (value: string) => N): LeanIMT<N> {
+        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<N>(hash)
+
+        tree._nodes = JSON.parse(nodes, (_, value) => {
+            if (typeof value === "string") {
+                return map ? map(value) : BigInt(value)
+            }
+
+            return value
+        })
+
+        return tree
     }
 }

packages/lean-imt/tests/lean-imt.test.ts

@@ -335,54 +335,75 @@ 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<number>(hash, [1, 2, 3])
+            const nodes = tree1.export()
+
+            const tree2 = LeanIMT.import<number>(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<boolean>(hash, [true, false, true])
+            const nodes = tree1.export()
 
-            const tree2 = new LeanIMT(poseidon, leaves)
+            const tree2 = LeanIMT.import<boolean>(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))
 
             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))
         })
     })
 })