From 42738809ad92457548e8b0ed2681a97cdf23af55 Mon Sep 17 00:00:00 2001
From: Raphael Panic <raphael.panic@gmx.at>
Date: Wed, 26 Mar 2025 11:47:22 +0100
Subject: [PATCH 1/3] Implemented smart hashing of batch updates

---
 .../src/trees/sparse/RollupMerkleTree.ts      | 54 +++++++++++++++
 packages/common/test/trees/MerkleTree.test.ts | 65 ++++++++++++++++++-
 2 files changed, 118 insertions(+), 1 deletion(-)

diff --git a/packages/common/src/trees/sparse/RollupMerkleTree.ts b/packages/common/src/trees/sparse/RollupMerkleTree.ts
index e5673ec95..bb6612337 100644
--- a/packages/common/src/trees/sparse/RollupMerkleTree.ts
+++ b/packages/common/src/trees/sparse/RollupMerkleTree.ts
@@ -67,6 +67,8 @@ export interface AbstractMerkleTree {
    */
   setLeaf(index: bigint, leaf: Field): void;
 
+  setLeafBatch(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 +296,58 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
       }
     }
 
+    public setLeafBatch(updates: { index: bigint; leaf: Field }[]): number {
+      updates.forEach(({ index }) => this.assertIndexRange(index));
+
+      type Change = { level: number; index: bigint; value: Field };
+      const changes: Change[] = [];
+
+      let numberOfHashes = 0;
+
+      // 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
+      let levelChanges = updates.sort((a, b) => (a.index < b.index ? 1 : -1));
+
+      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(0);
+            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]);
+          }
+          numberOfHashes++;
+
+          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)
+      );
+
+      return numberOfHashes;
+    }
+
     /**
      * Returns the witness (also known as
      * [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof))
diff --git a/packages/common/test/trees/MerkleTree.test.ts b/packages/common/test/trees/MerkleTree.test.ts
index 2452960e8..b278ad044 100644
--- a/packages/common/test/trees/MerkleTree.test.ts
+++ b/packages/common/test/trees/MerkleTree.test.ts
@@ -1,7 +1,70 @@
 import { beforeEach } from "@jest/globals";
 import { Field } from "o1js";
 
-import { createMerkleTree, InMemoryMerkleTreeStorage, log } from "../../src";
+import {
+  createMerkleTree,
+  InMemoryMerkleTreeStorage,
+  log,
+  RollupMerkleTree,
+  range,
+} from "../../src";
+
+describe("batch setLeaf", () => {
+  function generateBatch(size: number) {
+    return range(0, size).map(() => ({
+      index: Field.random().toBigInt(),
+      leaf: Field.random(),
+    }));
+  }
+
+  function captureTime<R>(f: () => R): [number, R] {
+    const start = Date.now();
+    const ret = f();
+    return [Date.now() - start, ret];
+  }
+
+  it("correctness", () => {
+    const Tree = createMerkleTree(256);
+    const tree1 = new Tree(new InMemoryMerkleTreeStorage());
+    const tree2 = new Tree(new InMemoryMerkleTreeStorage());
+
+    tree1.setLeaf(1n, Field(5));
+    tree1.setLeaf(Field.ORDER - 1n, Field(7));
+
+    tree2.setLeafBatch([
+      { index: 1n, leaf: Field(5) },
+      { index: Field.ORDER - 1n, leaf: Field(7) },
+    ]);
+
+    expect(tree1.getRoot().toString()).toStrictEqual(
+      tree2.getRoot().toString()
+    );
+  });
+
+  it.each([10, 100])("test speedup", (batchSize) => {
+    const tree1 = new RollupMerkleTree(new InMemoryMerkleTreeStorage());
+    const tree2 = new RollupMerkleTree(new InMemoryMerkleTreeStorage());
+
+    const batch = generateBatch(batchSize);
+
+    const slice = batch.slice();
+    const [time1, numHashes] = captureTime(() => tree1.setLeafBatch(slice));
+    const [time2] = captureTime(() =>
+      batch.forEach(({ index, leaf }) => tree2.setLeaf(index, leaf))
+    );
+
+    console.log(`Speedup for batch size ${batchSize}`);
+    console.log(time1);
+    console.log(time2);
+
+    console.log(`numHashes1 ${numHashes}`);
+    console.log(`numHashes2 ${255 * batchSize}`);
+
+    expect(tree1.getRoot().toString()).toStrictEqual(
+      tree2.getRoot().toString()
+    );
+  });
+});
 
 describe.each([4, 16, 256])("cachedMerkleTree - %s", (height) => {
   class RollupMerkleTree extends createMerkleTree(height) {}

From 91f67fef32790590a1585601f6f241c50eb89d43 Mon Sep 17 00:00:00 2001
From: Raphael Panic <raphael.panic@gmx.at>
Date: Sun, 13 Apr 2025 21:02:57 +0200
Subject: [PATCH 2/3] Finished batch leaf update, fixed a few bugs there

---
 .../src/trees/lmt/AbstractLinkedMerkleTree.ts |   2 +
 .../common/src/trees/lmt/LinkedMerkleTree.ts  | 151 ++++++++++++------
 .../src/trees/sparse/RollupMerkleTree.ts      |  36 +++--
 packages/common/test/trees/MerkleTree.test.ts |  64 ++++++--
 .../sequencing/BlockResultService.ts          |   9 +-
 .../tracing/StateTransitionTracingService.ts  |   4 -
 6 files changed, 183 insertions(+), 83 deletions(-)

diff --git a/packages/common/src/trees/lmt/AbstractLinkedMerkleTree.ts b/packages/common/src/trees/lmt/AbstractLinkedMerkleTree.ts
index d5209ad06..e951f83d9 100644
--- a/packages/common/src/trees/lmt/AbstractLinkedMerkleTree.ts
+++ b/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.
diff --git a/packages/common/src/trees/lmt/LinkedMerkleTree.ts b/packages/common/src/trees/lmt/LinkedMerkleTree.ts
index 93bc395d8..39e710990 100644
--- a/packages/common/src/trees/lmt/LinkedMerkleTree.ts
+++ b/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,55 +155,106 @@ 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 {
+          leafPrevious: "dummy",
+          leafCurrent: {
+            witness: storedLeaf.index,
+            witnessLeaf: new LinkedLeafStruct(
+              LinkedLeafStruct.fromValue(storedLeaf.leaf)
+            ),
+            write: treeWrite,
+          },
+        };
+      }
+    }
 
-        return new LinkedOperationWitness({
-          leafPrevious: witnessPrevious,
-          leafCurrent: current,
-        });
+    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 }[]) {
+      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 +264,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()
+      );
     }
 
     /**
diff --git a/packages/common/src/trees/sparse/RollupMerkleTree.ts b/packages/common/src/trees/sparse/RollupMerkleTree.ts
index bb6612337..5bbafeb65 100644
--- a/packages/common/src/trees/sparse/RollupMerkleTree.ts
+++ b/packages/common/src/trees/sparse/RollupMerkleTree.ts
@@ -2,10 +2,10 @@ import { Bool, Field, Poseidon, Provable, Struct } from "o1js";
 
 import { range } from "../../utils";
 import { TypedClass } from "../../types";
+import uniqBy from "lodash/uniqBy";
 
 import { MerkleTreeStore } from "./MerkleTreeStore";
 import { InMemoryMerkleTreeStorage } from "./InMemoryMerkleTreeStorage";
-
 export class StructTemplate extends Struct({
   path: Provable.Array(Field, 0),
   isLeft: Provable.Array(Bool, 0),
@@ -67,7 +67,7 @@ export interface AbstractMerkleTree {
    */
   setLeaf(index: bigint, leaf: Field): void;
 
-  setLeafBatch(updates: { index: bigint; leaf: Field }[]): void;
+  setLeaves(updates: { index: bigint; leaf: Field }[]): void;
 
   /**
    * Returns the witness (also known as
@@ -296,17 +296,26 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
       }
     }
 
-    public setLeafBatch(updates: { index: bigint; leaf: Field }[]): number {
+    public setLeaves(updates: { index: bigint; leaf: Field }[]) {
       updates.forEach(({ index }) => this.assertIndexRange(index));
 
       type Change = { level: number; index: bigint; value: Field };
       const changes: Change[] = [];
 
-      let numberOfHashes = 0;
-
-      // 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
-      let levelChanges = updates.sort((a, b) => (a.index < b.index ? 1 : -1));
+      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 = [];
@@ -316,7 +325,7 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
           let newNode;
           if (node.index % 2n === 0n) {
             let sibling;
-            const potentialSibling = levelChanges.at(0);
+            const potentialSibling = levelChanges.at(-1);
             if (
               potentialSibling !== undefined &&
               potentialSibling.index === node.index + 1n
@@ -331,7 +340,6 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
             const sibling = Field(this.getNode(level - 1, node.index + 1n));
             newNode = Poseidon.hash([sibling, node.leaf]);
           }
-          numberOfHashes++;
 
           const nextLevelIndex = node.index / 2n;
           changes.push({ level, index: nextLevelIndex, value: newNode });
@@ -341,11 +349,9 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
         levelChanges = nextLevelChanges.reverse();
       }
 
-      changes.forEach(({ level, index, value }) =>
-        this.setNode(level, index, value)
-      );
-
-      return numberOfHashes;
+      changes.forEach(({ level, index, value }) => {
+        this.setNode(level, index, value);
+      });
     }
 
     /**
diff --git a/packages/common/test/trees/MerkleTree.test.ts b/packages/common/test/trees/MerkleTree.test.ts
index b278ad044..1cd3a9023 100644
--- a/packages/common/test/trees/MerkleTree.test.ts
+++ b/packages/common/test/trees/MerkleTree.test.ts
@@ -10,56 +10,86 @@ import {
 } from "../../src";
 
 describe("batch setLeaf", () => {
-  function generateBatch(size: number) {
+  function generateBatch(
+    size: number,
+    height: number,
+    generator = () => Field.random().toBigInt()
+  ) {
+    const max = 2n ** BigInt(height - 1);
     return range(0, size).map(() => ({
-      index: Field.random().toBigInt(),
+      index: generator() % max,
       leaf: Field.random(),
     }));
   }
 
+  function generateBatchAdjacent(size: number, height: number) {
+    let start = 0n;
+    return generateBatch(size, height, () => {
+      start += 1n;
+      return start - 1n;
+    });
+  }
+
   function captureTime<R>(f: () => R): [number, R] {
     const start = Date.now();
     const ret = f();
     return [Date.now() - start, ret];
   }
+  const height = 10;
+  const max_index = 2n ** BigInt(height - 1) - 1n;
 
-  it("correctness", () => {
-    const Tree = createMerkleTree(256);
+  it.each([
+    [
+      { index: 1n, leaf: Field(5) },
+      { index: max_index, leaf: Field(7) },
+    ],
+    [
+      { index: max_index, leaf: Field(7) },
+      { index: max_index - 1n, leaf: Field(7) },
+      { index: 50n, leaf: Field(7) },
+      { index: 1n, leaf: Field(5) },
+    ],
+    generateBatch(5, height),
+    generateBatch(10, height),
+    generateBatch(50, height),
+    generateBatch(300, height),
+  ])("correctness", (...writes) => {
+    const Tree = createMerkleTree(height);
     const tree1 = new Tree(new InMemoryMerkleTreeStorage());
     const tree2 = new Tree(new InMemoryMerkleTreeStorage());
 
-    tree1.setLeaf(1n, Field(5));
-    tree1.setLeaf(Field.ORDER - 1n, Field(7));
+    writes.forEach(({ index, leaf }) => {
+      tree1.setLeaf(index, leaf);
+    });
 
-    tree2.setLeafBatch([
-      { index: 1n, leaf: Field(5) },
-      { index: Field.ORDER - 1n, leaf: Field(7) },
-    ]);
+    tree2.setLeaves(writes);
 
     expect(tree1.getRoot().toString()).toStrictEqual(
       tree2.getRoot().toString()
     );
   });
 
-  it.each([10, 100])("test speedup", (batchSize) => {
+  it.each([
+    ["random", 10, generateBatch],
+    ["random", 100, generateBatch],
+    ["adjacent", 100, generateBatchAdjacent],
+    ["adjacent", 10, generateBatchAdjacent],
+  ])("test speedup: %s (%i leaves)", (label, batchSize, generateFunction) => {
     const tree1 = new RollupMerkleTree(new InMemoryMerkleTreeStorage());
     const tree2 = new RollupMerkleTree(new InMemoryMerkleTreeStorage());
 
-    const batch = generateBatch(batchSize);
+    const batch = generateFunction(batchSize, RollupMerkleTree.HEIGHT);
 
     const slice = batch.slice();
-    const [time1, numHashes] = captureTime(() => tree1.setLeafBatch(slice));
+    const [time1] = captureTime(() => tree1.setLeaves(slice));
     const [time2] = captureTime(() =>
       batch.forEach(({ index, leaf }) => tree2.setLeaf(index, leaf))
     );
 
-    console.log(`Speedup for batch size ${batchSize}`);
+    console.log(`Speedup for batch size ${batchSize}, mode ${label}`);
     console.log(time1);
     console.log(time2);
 
-    console.log(`numHashes1 ${numHashes}`);
-    console.log(`numHashes2 ${255 * batchSize}`);
-
     expect(tree1.getRoot().toString()).toStrictEqual(
       tree2.getRoot().toString()
     );
diff --git a/packages/sequencer/src/protocol/production/sequencing/BlockResultService.ts b/packages/sequencer/src/protocol/production/sequencing/BlockResultService.ts
index 969a8e838..5419344bb 100644
--- a/packages/sequencer/src/protocol/production/sequencing/BlockResultService.ts
+++ b/packages/sequencer/src/protocol/production/sequencing/BlockResultService.ts
@@ -170,10 +170,15 @@ export class BlockResultService {
 
     const tree = new LinkedMerkleTree(store.treeStore, store);
 
-    Object.entries(stateDiff).forEach(([key, state]) => {
+    const writes = Object.entries(stateDiff).map(([key, state]) => {
       const treeValue = state !== undefined ? Poseidon.hash(state) : Field(0);
-      tree.setLeaf(BigInt(key), treeValue.toBigInt());
+      return { path: BigInt(key), value: treeValue.toBigInt() };
     });
+    tree.setLeaves(writes);
+    // Object.entries(stateDiff).forEach(([key, state]) => {
+    //   const treeValue = state !== undefined ? Poseidon.hash(state) : Field(0);
+    //   tree.setLeaf(BigInt(key), treeValue.toBigInt());
+    // });
 
     return tree;
   }
diff --git a/packages/sequencer/src/protocol/production/tracing/StateTransitionTracingService.ts b/packages/sequencer/src/protocol/production/tracing/StateTransitionTracingService.ts
index 4f0d2194b..03621892b 100644
--- a/packages/sequencer/src/protocol/production/tracing/StateTransitionTracingService.ts
+++ b/packages/sequencer/src/protocol/production/tracing/StateTransitionTracingService.ts
@@ -142,10 +142,6 @@ export class StateTransitionTracingService {
         async (transitionInfo) => {
           const { stateTransition, type, witnessRoot } = transitionInfo;
 
-          // const merkleWitness = tree.getWitness(
-          //   stateTransition.path.toBigInt()
-          // );
-
           let witness: LinkedMerkleTreeWitness;
 
           if (stateTransition.to.isSome.toBoolean()) {

From 6df6b8e78daa453b9272d933e3fa3ab4a0c1f58b Mon Sep 17 00:00:00 2001
From: Raphael Panic <raphael.panic@gmx.at>
Date: Mon, 14 Apr 2025 13:14:21 +0200
Subject: [PATCH 3/3] Fixed a bug for blockproduction and setLeaves

---
 .../common/src/trees/lmt/LinkedMerkleTree.ts  | 22 +++++----
 .../src/trees/sparse/RollupMerkleTree.ts      |  8 ++-
 packages/common/test/trees/MerkleTree.test.ts | 49 +++++++++++++------
 .../sequencing/TransactionExecutionService.ts |  6 +--
 4 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/packages/common/src/trees/lmt/LinkedMerkleTree.ts b/packages/common/src/trees/lmt/LinkedMerkleTree.ts
index 39e710990..dc209a9b6 100644
--- a/packages/common/src/trees/lmt/LinkedMerkleTree.ts
+++ b/packages/common/src/trees/lmt/LinkedMerkleTree.ts
@@ -243,18 +243,20 @@ export function createLinkedMerkleTree(
     }
 
     public setLeaves(batch: { path: bigint; value: bigint }[]) {
-      const witnesses = batch.map(({ path, value }) =>
-        this.setLeafInternal(path, value)
-      );
+      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
+        // 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
+            )
           )
-        )
-      );
+        );
+      }
     }
 
     /**
diff --git a/packages/common/src/trees/sparse/RollupMerkleTree.ts b/packages/common/src/trees/sparse/RollupMerkleTree.ts
index 5bbafeb65..1974e7078 100644
--- a/packages/common/src/trees/sparse/RollupMerkleTree.ts
+++ b/packages/common/src/trees/sparse/RollupMerkleTree.ts
@@ -1,11 +1,12 @@
 import { Bool, Field, Poseidon, Provable, Struct } from "o1js";
+import uniqBy from "lodash/uniqBy";
 
 import { range } from "../../utils";
 import { TypedClass } from "../../types";
-import uniqBy from "lodash/uniqBy";
 
 import { MerkleTreeStore } from "./MerkleTreeStore";
 import { InMemoryMerkleTreeStorage } from "./InMemoryMerkleTreeStorage";
+
 export class StructTemplate extends Struct({
   path: Provable.Array(Field, 0),
   isLeft: Provable.Array(Bool, 0),
@@ -302,6 +303,9 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
       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
@@ -337,7 +341,7 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
             }
             newNode = Poseidon.hash([node.leaf, sibling]);
           } else {
-            const sibling = Field(this.getNode(level - 1, node.index + 1n));
+            const sibling = Field(this.getNode(level - 1, node.index - 1n));
             newNode = Poseidon.hash([sibling, node.leaf]);
           }
 
diff --git a/packages/common/test/trees/MerkleTree.test.ts b/packages/common/test/trees/MerkleTree.test.ts
index 1cd3a9023..221c8d70e 100644
--- a/packages/common/test/trees/MerkleTree.test.ts
+++ b/packages/common/test/trees/MerkleTree.test.ts
@@ -36,16 +36,16 @@ describe("batch setLeaf", () => {
     return [Date.now() - start, ret];
   }
   const height = 10;
-  const max_index = 2n ** BigInt(height - 1) - 1n;
+  const maxIndex = 2n ** BigInt(height - 1) - 1n;
 
   it.each([
     [
       { index: 1n, leaf: Field(5) },
-      { index: max_index, leaf: Field(7) },
+      { index: maxIndex, leaf: Field(7) },
     ],
     [
-      { index: max_index, leaf: Field(7) },
-      { index: max_index - 1n, leaf: Field(7) },
+      { index: maxIndex, leaf: Field(7) },
+      { index: maxIndex - 1n, leaf: Field(7) },
       { index: 50n, leaf: Field(7) },
       { index: 1n, leaf: Field(5) },
     ],
@@ -69,6 +69,31 @@ describe("batch setLeaf", () => {
     );
   });
 
+  it.each([
+    // This tests the correct retrieval of previously-set siblings (vs. above
+    // where always fetch zero-siblings)
+    [[{ index: 1n, leaf: Field(5) }], [{ index: 4n, leaf: Field(1) }]],
+    [[{ index: 4n, leaf: Field(5) }], [{ index: 1n, leaf: Field(1) }]],
+  ])("correctness - batches", (...writes) => {
+    expect.assertions(writes.length);
+
+    const Tree = createMerkleTree(height);
+    const tree1 = new Tree(new InMemoryMerkleTreeStorage());
+    const tree2 = new Tree(new InMemoryMerkleTreeStorage());
+
+    writes.forEach((writes2) => {
+      writes2.forEach(({ index, leaf }) => {
+        tree1.setLeaf(index, leaf);
+      });
+
+      tree2.setLeaves(writes2);
+
+      expect(tree1.getRoot().toString()).toStrictEqual(
+        tree2.getRoot().toString()
+      );
+    });
+  });
+
   it.each([
     ["random", 10, generateBatch],
     ["random", 100, generateBatch],
@@ -97,26 +122,24 @@ describe("batch setLeaf", () => {
 });
 
 describe.each([4, 16, 256])("cachedMerkleTree - %s", (height) => {
-  class RollupMerkleTree extends createMerkleTree(height) {}
+  class MerkleTree extends createMerkleTree(height) {}
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  class RollupMerkleTreeWitness extends RollupMerkleTree.WITNESS {}
+  class MerkleTreeWitness extends MerkleTree.WITNESS {}
 
   let store: InMemoryMerkleTreeStorage;
-  let tree: RollupMerkleTree;
+  let tree: MerkleTree;
 
   beforeEach(() => {
     log.setLevel("INFO");
 
     store = new InMemoryMerkleTreeStorage();
-    tree = new RollupMerkleTree(store);
+    tree = new MerkleTree(store);
   });
 
   it("should have the same root when empty", () => {
     expect.assertions(1);
 
-    expect(tree.getRoot().toBigInt()).toStrictEqual(
-      RollupMerkleTree.EMPTY_ROOT
-    );
+    expect(tree.getRoot().toBigInt()).toStrictEqual(MerkleTree.EMPTY_ROOT);
   });
 
   it("should have a different root when not empty", () => {
@@ -124,9 +147,7 @@ describe.each([4, 16, 256])("cachedMerkleTree - %s", (height) => {
 
     tree.setLeaf(1n, Field(1));
 
-    expect(tree.getRoot().toBigInt()).not.toStrictEqual(
-      RollupMerkleTree.EMPTY_ROOT
-    );
+    expect(tree.getRoot().toBigInt()).not.toStrictEqual(MerkleTree.EMPTY_ROOT);
   });
 
   it("should have the same root after adding and removing item", () => {
diff --git a/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts b/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts
index 10b8da93b..05d25c279 100644
--- a/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts
+++ b/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts
@@ -173,7 +173,7 @@ export async function executeWithExecutionContext<MethodResult>(
   };
 }
 
-function traceSTs(msg: string, stateTransitions: StateTransition<any>[]) {
+function traceLogSTs(msg: string, stateTransitions: StateTransition<any>[]) {
   log.trace(
     msg,
     JSON.stringify(
@@ -265,7 +265,7 @@ export class TransactionExecutionService {
       throw error;
     }
 
-    traceSTs(`${hookName} STs:`, result.stateTransitions);
+    traceLogSTs(`${hookName} STs:`, result.stateTransitions);
 
     return result;
   }
@@ -344,7 +344,7 @@ export class TransactionExecutionService {
       args,
       runtimeContextInputs
     );
-    traceSTs("STs:", runtimeResult.stateTransitions);
+    traceLogSTs("STs:", runtimeResult.stateTransitions);
 
     // Apply runtime STs (only if the tx succeeded)
     if (runtimeResult.status.toBoolean()) {