Implemented smart hashing of batch updates for LMT

packages/common/src/trees/lmt/AbstractLinkedMerkleTree.ts

@@ -55,6 +55,8 @@ export interface AbstractLinkedMerkleTree {
    */
   setLeaf(path: bigint, value?: bigint): LinkedOperationWitnessValue;
 
+  setLeaves(batch: { path: bigint; value: bigint }[]): void;
+
   /**
    * Returns a leaf which lives at a given path.
    * Errors otherwise.

packages/common/src/trees/lmt/LinkedMerkleTree.ts

@@ -12,6 +12,17 @@ import {
   AbstractLinkedMerkleTreeClass,
 } from "./AbstractLinkedMerkleTree";
 
+type LeafOperationInstruction = {
+  witness: bigint;
+  witnessLeaf: LinkedLeafStruct;
+  write: { index: bigint; leaf: Field };
+};
+
+type SetLeafMetadata = {
+  leafPrevious: LeafOperationInstruction | "dummy";
+  leafCurrent: LeafOperationInstruction;
+};
+
 export function createLinkedMerkleTree(
   height: number
 ): AbstractLinkedMerkleTreeClass {
@@ -109,21 +120,20 @@ export function createLinkedMerkleTree(
       return this.tree.getRoot().toConstant();
     }
 
-    private setMerkleLeaf(index: bigint, leaf: LinkedLeaf) {
+    private writeLeaf(index: bigint, leaf: LinkedLeaf) {
       this.leafStore.setLeaf(index, leaf);
 
       const leafHash = new LinkedLeafStruct(
         LinkedLeafStruct.fromValue(leaf)
       ).hash();
-      this.tree.setLeaf(index, leafHash);
+
+      return {
+        index,
+        leaf: leafHash,
+      };
     }
 
-    /**
-     * Sets the value of a node at a given index to a given value.
-     * @param path Position of the leaf node.
-     * @param value New value.
-     */
-    public setLeaf(path: bigint, value: bigint): LinkedOperationWitness {
+    private setLeafInternal(path: bigint, value: bigint): SetLeafMetadata {
       const storedLeaf = this.leafStore.getLeaf(path);
 
       if (storedLeaf === undefined) {
@@ -145,54 +155,107 @@ export function createLinkedMerkleTree(
           throw Error(`Prev leaf shouldn't be undefined (path ${path})`);
         }
 
-        const previousLeafMerkleWitness = this.tree.getWitness(
-          previousLeaf.index
-        );
-
         const newPrevLeaf = {
           ...previousLeaf.leaf,
           nextPath: path,
         };
-        this.setMerkleLeaf(previousLeaf.index, newPrevLeaf);
-
-        const currentMerkleWitness = this.tree.getWitness(nextFreeIndex);
+        const treeWrite1 = this.writeLeaf(previousLeaf.index, newPrevLeaf);
 
         const newLeaf = {
           path,
           value,
           nextPath: previousLeaf.leaf.nextPath,
         };
-        this.setMerkleLeaf(nextFreeIndex, newLeaf);
+        const treeWrite2 = this.writeLeaf(nextFreeIndex, newLeaf);
 
-        return new LinkedOperationWitness({
-          leafPrevious: new LinkedLeafAndMerkleWitness({
-            leaf: new LinkedLeafStruct(
+        return {
+          leafPrevious: {
+            witness: previousLeaf.index,
+            witnessLeaf: new LinkedLeafStruct(
               LinkedLeafStruct.fromValue(previousLeaf.leaf)
             ),
-            merkleWitness: previousLeafMerkleWitness,
-          }),
-          leafCurrent: new LinkedLeafAndMerkleWitness({
-            leaf: LinkedLeafStruct.dummy(),
-            merkleWitness: currentMerkleWitness,
-          }),
-        });
+            write: treeWrite1,
+          },
+          leafCurrent: {
+            witness: nextFreeIndex,
+            witnessLeaf: LinkedLeafStruct.dummy(),
+            write: treeWrite2,
+          },
+        };
       } else {
         // Update case
-        const witnessPrevious =
-          AbstractLinkedRollupMerkleTree.dummyReadWitness();
-
-        // TODO This makes an unnecessary leafstore lookup currently, reuse storedLeaf instead
-        const current = this.getReadWitness(storedLeaf.leaf.path);
-
-        this.setMerkleLeaf(storedLeaf.index, {
+        const updatedLeaf = {
           ...storedLeaf.leaf,
           value: value,
-        });
+        };
+
+        const treeWrite = this.writeLeaf(storedLeaf.index, updatedLeaf);
 
-        return new LinkedOperationWitness({
-          leafPrevious: witnessPrevious,
-          leafCurrent: current,
-        });
+        return {
+          leafPrevious: "dummy",
+          leafCurrent: {
+            witness: storedLeaf.index,
+            witnessLeaf: new LinkedLeafStruct(
+              LinkedLeafStruct.fromValue(storedLeaf.leaf)
+            ),
+            write: treeWrite,
+          },
+        };
+      }
+    }
+
+    private applyOperationInstruction(
+      instruction: LeafOperationInstruction | "dummy"
+    ): LinkedLeafAndMerkleWitness {
+      if (instruction === "dummy") {
+        return AbstractLinkedRollupMerkleTree.dummyReadWitness();
+      }
+
+      const merkleWitness = this.tree.getWitness(instruction.witness);
+
+      this.tree.setLeaf(instruction.write.index, instruction.write.leaf);
+
+      return new LinkedLeafAndMerkleWitness({
+        merkleWitness,
+        leaf: instruction.witnessLeaf,
+      });
+    }
+
+    /**
+     * Sets the value of a node at a given index to a given value.
+     * @param path Position of the leaf node.
+     * @param value New value.
+     */
+    public setLeaf(path: bigint, value: bigint): LinkedOperationWitness {
+      const {
+        leafPrevious: previousLeafInstruction,
+        leafCurrent: currentLeafInstruction,
+      } = this.setLeafInternal(path, value);
+
+      const leafPrevious = this.applyOperationInstruction(
+        previousLeafInstruction
+      );
+      const leafCurrent = this.applyOperationInstruction(
+        currentLeafInstruction
+      );
+
+      return { leafPrevious, leafCurrent };
+    }
+
+    public setLeaves(batch: { path: bigint; value: bigint }[]) {
+      if (batch.length > 0) {
+        const witnesses = batch.map(({ path, value }) =>
+          this.setLeafInternal(path, value)
+        );
+
+        // tree.setLeafBatch internally takes care of making the writes unique to optimize
+        this.tree.setLeaves(
+          witnesses.flatMap(({ leafPrevious, leafCurrent }) =>
+            (leafPrevious === "dummy" ? [] : [leafPrevious.write]).concat(
+              leafCurrent.write
+            )
+          )
+        );
       }
     }
 
@@ -203,18 +266,18 @@ export function createLinkedMerkleTree(
     private setLeafInitialisation() {
       // This is the maximum value of the hash
       const MAX_FIELD_VALUE: bigint = Field.ORDER - 1n;
-      this.leafStore.setLeaf(0n, {
+      const zeroLeaf = {
         value: 0n,
         path: 0n,
         nextPath: MAX_FIELD_VALUE,
-      });
+      };
+      this.leafStore.setLeaf(0n, zeroLeaf);
       // We now set the leafs in the merkle tree to cascade the values up
       // the tree.
-      this.setMerkleLeaf(0n, {
-        value: 0n,
-        path: 0n,
-        nextPath: MAX_FIELD_VALUE,
-      });
+      this.tree.setLeaf(
+        0n,
+        new LinkedLeafStruct(LinkedLeafStruct.fromValue(zeroLeaf)).hash()
+      );
     }
 
     /**

packages/common/src/trees/sparse/RollupMerkleTree.ts

@@ -1,4 +1,5 @@
 import { Bool, Field, Poseidon, Provable, Struct } from "o1js";
+import uniqBy from "lodash/uniqBy";
 
 import { range } from "../../utils";
 import { TypedClass } from "../../types";
@@ -67,6 +68,8 @@ export interface AbstractMerkleTree {
    */
   setLeaf(index: bigint, leaf: Field): void;
 
+  setLeaves(updates: { index: bigint; leaf: Field }[]): void;
+
   /**
    * Returns the witness (also known as
    * [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof))
@@ -294,6 +297,67 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
       }
     }
 
+    public setLeaves(updates: { index: bigint; leaf: Field }[]) {
+      updates.forEach(({ index }) => this.assertIndexRange(index));
+
+      type Change = { level: number; index: bigint; value: Field };
+      const changes: Change[] = [];
+
+      // We have to reverse here, because uniqBy only takes the first occurrence of every entry,
+      // but we need the last one (since that semantically overwrites the previous one,
+      // so we can ignore it)
+      let levelChanges = uniqBy(updates.reverse(), "index")
+        // we can assume no index is in this list twice, so we don't care about the 0 case
+        // This is in reverse order, so its a queue
+        .sort((a, b) => (a.index < b.index ? 1 : -1));
+
+      changes.push(
+        ...levelChanges
+          .map(({ leaf, index }) => ({
+            level: 0,
+            index,
+            value: leaf,
+          }))
+          .reverse()
+      );
+
+      for (let level = 1; level < AbstractRollupMerkleTree.HEIGHT; level += 1) {
+        const nextLevelChanges: typeof levelChanges = [];
+        while (levelChanges.length > 0) {
+          const node = levelChanges.pop()!;
+
+          let newNode;
+          if (node.index % 2n === 0n) {
+            let sibling;
+            const potentialSibling = levelChanges.at(-1);
+            if (
+              potentialSibling !== undefined &&
+              potentialSibling.index === node.index + 1n
+            ) {
+              sibling = potentialSibling.leaf;
+              levelChanges.pop();
+            } else {
+              sibling = Field(this.getNode(level - 1, node.index + 1n));
+            }
+            newNode = Poseidon.hash([node.leaf, sibling]);
+          } else {
+            const sibling = Field(this.getNode(level - 1, node.index - 1n));
+            newNode = Poseidon.hash([sibling, node.leaf]);
+          }
+
+          const nextLevelIndex = node.index / 2n;
+          changes.push({ level, index: nextLevelIndex, value: newNode });
+          nextLevelChanges.push({ index: nextLevelIndex, leaf: newNode });
+        }
+
+        levelChanges = nextLevelChanges.reverse();
+      }
+
+      changes.forEach(({ level, index, value }) => {
+        this.setNode(level, index, value);
+      });
+    }
+
     /**
      * Returns the witness (also known as
      * [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof))