Implemented smart hashing of batch updates

nx.json

@@ -17,6 +17,7 @@
     "prisma:generate": {
       "inputs": [
         "{projectRoot}/package.json",
+        "{projectRoot}/prisma/*",
         "{projectRoot}/prisma/**/*"
       ]
     },

packages/api/src/graphql/modules/MerkleWitnessResolver.ts

@@ -3,8 +3,10 @@ import { Length } from "class-validator";
 import { inject } from "tsyringe";
 import { RollupMerkleTree, RollupMerkleTreeWitness } from "@proto-kit/common";
 import {
-  AsyncMerkleTreeStore,
+  BlockStorage,
   CachedMerkleTreeStore,
+  MaskName,
+  TreeStoreCreator,
 } from "@proto-kit/sequencer";
 
 import { GraphqlModule, graphqlModule } from "../GraphqlModule";
@@ -34,7 +36,9 @@ export class MerkleWitnessDTO {
 @graphqlModule()
 export class MerkleWitnessResolver extends GraphqlModule<object> {
   public constructor(
-    @inject("AsyncMerkleStore") private readonly treeStore: AsyncMerkleTreeStore
+    @inject("TreeStoreCreator")
+    private readonly treeStoreCreator: TreeStoreCreator,
+    @inject("BlockStorage") private readonly blockStorage: BlockStorage
   ) {
     super();
   }
@@ -44,7 +48,14 @@ export class MerkleWitnessResolver extends GraphqlModule<object> {
       "Allows retrieval of merkle witnesses corresponding to a specific path in the appchain's state tree. These proves are generally retrieved from the current 'proven' state",
   })
   public async witness(@Arg("path") path: string) {
-    const syncStore = new CachedMerkleTreeStore(this.treeStore);
+    const latestBlock = await this.blockStorage.getLatestBlock();
+    const maskName =
+      latestBlock !== undefined
+        ? MaskName.block(latestBlock.block.height)
+        : MaskName.base();
+    const treeStore = this.treeStoreCreator.getMask(maskName);
+
+    const syncStore = new CachedMerkleTreeStore(treeStore);
     await syncStore.preloadKey(BigInt(path));
 
     const tree = new RollupMerkleTree(syncStore);

packages/common/src/trees/RollupMerkleTree.ts

@@ -65,6 +65,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))
@@ -288,6 +290,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))

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) {}

packages/indexer/package.json

@@ -22,7 +22,7 @@
   "dependencies": {
     "@envelop/extended-validation": "^4.1.0",
     "@inkjs/ui": "^1.0.0",
-    "@prisma/client": "^5.18.0",
+    "@prisma/client": "^5.22.0",
     "@types/yargs": "^17.0.29",
     "figlet": "^1.7.0",
     "ink": "^4.4.1",

packages/persistance/package.json

@@ -20,9 +20,9 @@
     "access": "public"
   },
   "dependencies": {
-    "@prisma/client": "^5.18.0",
+    "@prisma/client": "^5.22.0",
     "lodash": "^4.17.21",
-    "prisma": "^5.18.0",
+    "prisma": "^5.22.0",
     "redis": "^4.6.12",
     "reflect-metadata": "^0.1.13"
   },

packages/persistance/prisma/migrations/20250129151707_state_masking/migration.sql

@@ -0,0 +1,29 @@
+/*
+  Warnings:
+
+  - The primary key for the `State` table will be changed. If it partially fails, the table could be left without primary key constraint.
+  - You are about to drop the column `mask` on the `State` table. All the data in the column will be lost.
+  - Added the required column `maskId` to the `State` table without a default value. This is not possible if the table is not empty.
+
+*/
+-- AlterTable
+TRUNCATE TABLE "State";
+ALTER TABLE "State" DROP CONSTRAINT "State_pkey",
+DROP COLUMN "mask",
+ADD COLUMN     "maskId" INTEGER NOT NULL,
+ADD CONSTRAINT "State_pkey" PRIMARY KEY ("path", "maskId");
+
+-- CreateTable
+CREATE TABLE "Mask" (
+    "id" SERIAL NOT NULL,
+    "name" VARCHAR(256) NOT NULL,
+    "parent" INTEGER,
+
+    CONSTRAINT "Mask_pkey" PRIMARY KEY ("id")
+);
+
+-- AddForeignKey
+ALTER TABLE "State" ADD CONSTRAINT "State_maskId_fkey" FOREIGN KEY ("maskId") REFERENCES "Mask"("id") ON DELETE RESTRICT ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "Mask" ADD CONSTRAINT "Mask_parent_fkey" FOREIGN KEY ("parent") REFERENCES "Mask"("id") ON DELETE SET NULL ON UPDATE CASCADE;

packages/persistance/prisma/schema.prisma

@@ -3,9 +3,6 @@
 
 generator client {
   provider = "prisma-client-js"
-  // output   = "./../node_modules/.prisma/client"
-  // Enable after upgrade to 5.9.0
-  // previewFeatures = ["relationJoins"]
 }
 
 datasource db {
@@ -16,9 +13,21 @@ datasource db {
 model State {
   path   Decimal   @db.Decimal(78, 0)
   values Decimal[] @db.Decimal(78, 0)
-  mask   String    @db.VarChar(256)
 
-  @@id([path, mask])
+  maskId Int
+  mask   Mask @relation(fields: [maskId], references: [id])
+
+  @@id([path, maskId])
+}
+
+model Mask {
+  id   Int    @id @default(autoincrement())
+  name String @db.VarChar(256)
+
+  parent     Int?
+  parentMask Mask?   @relation("ParentMasks", fields: [parent], references: [id])
+  children   Mask[]  @relation("ParentMasks")
+  State      State[]
 }
 
 model Transaction {

packages/persistance/src/PrismaDatabaseConnection.ts

@@ -6,12 +6,12 @@ import {
 } from "@proto-kit/sequencer";
 import { DependencyFactory, OmitKeys } from "@proto-kit/common";
 
-import { PrismaStateService } from "./services/prisma/PrismaStateService";
 import { PrismaBatchStore } from "./services/prisma/PrismaBatchStore";
 import { PrismaBlockStorage } from "./services/prisma/PrismaBlockStorage";
 import { PrismaSettlementStorage } from "./services/prisma/PrismaSettlementStorage";
 import { PrismaMessageStorage } from "./services/prisma/PrismaMessageStorage";
 import { PrismaTransactionStorage } from "./services/prisma/PrismaTransactionStorage";
+import { PrismaStateServiceCreator } from "./creators/PrismaStateServiceCreator";
 
 export interface PrismaDatabaseConfig {
   // Either object-based config or connection string
@@ -49,12 +49,9 @@ export class PrismaDatabaseConnection
 
   public dependencies(): OmitKeys<
     StorageDependencyMinimumDependencies,
-    "asyncMerkleStore" | "blockTreeStore" | "unprovenMerkleStore"
+    "blockTreeStore" | "treeStoreCreator"
   > {
     return {
-      asyncStateService: {
-        useFactory: () => new PrismaStateService(this, "batch"),
-      },
       batchStorage: {
         useClass: PrismaBatchStore,
       },
@@ -64,9 +61,6 @@ export class PrismaDatabaseConnection
       blockStorage: {
         useClass: PrismaBlockStorage,
       },
-      unprovenStateService: {
-        useFactory: () => new PrismaStateService(this, "block"),
-      },
       settlementStorage: {
         useClass: PrismaSettlementStorage,
       },
@@ -76,6 +70,9 @@ export class PrismaDatabaseConnection
       transactionStorage: {
         useClass: PrismaTransactionStorage,
       },
+      stateServiceCreator: {
+        useClass: PrismaStateServiceCreator,
+      },
     };
   }
 
@@ -90,6 +87,7 @@ export class PrismaDatabaseConnection
       "Settlement",
       "IncomingMessageBatch",
       "IncomingMessageBatchTransaction",
+      "Mask",
     ];
 
     await this.prismaClient.$transaction(

packages/persistance/src/PrismaRedisDatabase.ts

@@ -4,10 +4,12 @@ import {
   StorageDependencyMinimumDependencies,
   Database,
   closeable,
+  Sequencer,
 } from "@proto-kit/sequencer";
 import { ChildContainerProvider } from "@proto-kit/common";
 import { PrismaClient } from "@prisma/client";
 import { RedisClientType } from "redis";
+import { inject } from "tsyringe";
 
 import {
   PrismaConnection,
@@ -36,10 +38,10 @@ export class PrismaRedisDatabase
 
   public redis: RedisConnectionModule;
 
-  public constructor() {
+  public constructor(@inject("Sequencer") sequencer: Sequencer<any>) {
     super();
     this.prisma = new PrismaDatabaseConnection();
-    this.redis = new RedisConnectionModule();
+    this.redis = new RedisConnectionModule(sequencer);
   }
 
   public get prismaClient(): PrismaClient {