Improve serialization for merkle tree

zkdb/src/merkle-tree/merkle-tree-base.ts

@@ -1,6 +1,7 @@
 import { Field, Poseidon } from 'snarkyjs';
 import { MerkleProof } from './common.js';
 import { createExtendedMerkleWitness } from './merkle-tree-extended.js';
+import { BSON } from 'bson';
 
 export type MerkleNodesMap = {
   [level: number]: {
@@ -27,12 +28,12 @@ export abstract class BaseMerkleTree {
    *
    * @param height The height of the Merkle tree.
    */
-  constructor(height: number) {
+  constructor(height: number, nodesMap: MerkleNodesMap = {}) {
     if (height < 1) {
       throw new Error('Merkle tree height must be greater than 0');
     }
     this.height = height;
-    this.nodesMap = {};
+    this.nodesMap = nodesMap;
     this.zeroes = new Array(height);
     this.zeroes[0] = Field(0);
     for (let i = 1; i < height; i += 1) {
@@ -237,4 +238,45 @@ export abstract class BaseMerkleTree {
   public get leafCount() {
     return 2n ** BigInt(this.height - 1);
   }
+
+  toBSON(): Uint8Array {
+    return this._toBSON(this.nodesMap);
+  }
+
+  static fromBSON(data: Uint8Array): MerkleNodesMap {
+    const root = BSON.deserialize(data);
+    const matrix: { index: number; value: string }[][] = root.matrix;
+
+    let nodesMap: MerkleNodesMap = {};
+
+    for (let level = 0; level < matrix.length; level++) {
+      nodesMap[level] = {};
+      const levelMatrix = matrix[level];
+      for (const element of levelMatrix) {
+        nodesMap[level][element.index] = Field(element.value);
+      }
+    }
+
+    return nodesMap;
+  }
+
+  protected _toBSON(nodesMap: MerkleNodesMap): Uint8Array {
+    let matrix: { index: number; value: string }[][] = [];
+
+    for (const level in nodesMap) {
+      matrix[level] = [];
+      const levelNodes = nodesMap[level];
+      for (const nodeIndex in levelNodes) {
+        const nodeValue = levelNodes[nodeIndex];
+        matrix[level].push({
+          index: parseInt(nodeIndex),
+          value: nodeValue.toString(),
+        });
+      }
+    }
+
+    const root = { matrix };
+
+    return BSON.serialize(root);
+  }
 }

zkdb/src/merkle-tree/merkle-tree-ipfs.ts

@@ -1,8 +1,6 @@
 import { BaseMerkleTree, MerkleNodesMap } from './merkle-tree-base.js';
 import { StorageEngineIPFS } from '../storage-engine/ipfs.js';
-import { BSON } from 'bson';
 import { Field } from 'snarkyjs';
-import { mapFieldToString, mapStringToField } from './mappers.js';
 
 export const MERKLE_TREE_COLLECTION_NAME = '.security';
 
@@ -11,8 +9,12 @@ export const MERKLE_TREE_FILE_NAME = 'merkle_tree';
 export default class DistributedMerkleTree extends BaseMerkleTree {
   private ipfs: StorageEngineIPFS;
 
-  constructor(ipfs: StorageEngineIPFS, height: number) {
-    super(height);
+  constructor(
+    ipfs: StorageEngineIPFS,
+    height: number,
+    nodesMap: MerkleNodesMap = {}
+  ) {
+    super(height, nodesMap);
     this.ipfs = ipfs;
     ipfs.use(MERKLE_TREE_COLLECTION_NAME);
   }
@@ -67,7 +69,7 @@ export default class DistributedMerkleTree extends BaseMerkleTree {
 
   protected async writeNodes(nodes: MerkleNodesMap): Promise<void> {
     this.ipfs.use(MERKLE_TREE_COLLECTION_NAME);
-    await this.ipfs.writeMerkleBSON(mapFieldToString(nodes));
+    await this.ipfs.writeMerkleBSON(this._toBSON(nodes));
   }
 
   protected async fetchNodes(): Promise<MerkleNodesMap> {
@@ -76,20 +78,7 @@ export default class DistributedMerkleTree extends BaseMerkleTree {
     }
 
     const nodeBytes = await this.ipfs.readMerkleBSON();
-    const bson = BSON.deserialize(nodeBytes);
-
-    const nodesMap: MerkleNodesMap = bson as MerkleNodesMap;
-
-    for (const level in nodesMap) {
-      const levelNodes = nodesMap[level];
-      for (const nodeIndex in levelNodes) {
-        const nodeValue = levelNodes[nodeIndex];
-        delete levelNodes[nodeIndex];
-        levelNodes[nodeIndex] = Field(nodeValue);
-      }
-    }
-
-    return mapStringToField(nodesMap);
+    return BaseMerkleTree.fromBSON(nodeBytes);
   }
 
   protected async calculateMerklePath(_: bigint): Promise<MerkleNodesMap> {

zkdb/src/merkle-tree/merkle-tree.test.ts

@@ -78,8 +78,6 @@ async function getNodeTest(
 
   const zeroes: Field[] = JSON.parse(JSON.stringify(expectedMerkleTree)).zeroes;
 
-  Circuit.log('zeroes', zeroes);
-
   // Execute
   let isPassed = Bool(true);
 
@@ -133,8 +131,9 @@ async function getSetLeafMerkleTreeTest(
   merkleTree: BaseMerkleTree,
   expectedMerkleTree: MerkleTree
 ) {
+  console.log('set leaf');
   // Setup
-  const LEAF_AMOUNT = 131;
+  const LEAF_AMOUNT = 20;
 
   // Execute
   for (let i = 0; i < LEAF_AMOUNT; i++) {

zkdb/src/storage-engine/ipfs.ts

@@ -402,9 +402,9 @@ export class StorageEngineIPFS implements TIPFSFileSystem, TIPFSFileIndex {
    * @param BSONData Merkle tree BSON data
    * @returns IIPFSEntry
    */
-  async writeMerkleBSON(BSONData: any): Promise<IIPFSEntry> {
+  async writeMerkleBSON(bson: Uint8Array): Promise<IIPFSEntry> {
     this.use(MERKLE_TREE_COLLECTION_NAME);
-    return this.write(MERKLE_TREE_FILE_NAME, BSON.serialize(BSONData));
+    return this.write(MERKLE_TREE_FILE_NAME, bson);
   }
 
   /**