Refactor/separate tree services

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

@@ -6,6 +6,7 @@ import {
 } from "@proto-kit/common";
 import {
   AsyncLinkedLeafStore,
+  AsyncMerkleTreeStore,
   CachedLinkedLeafStore,
 } from "@proto-kit/sequencer";
 
@@ -42,7 +43,9 @@ export class LinkedMerkleWitnessDTO {
 export class LinkedMerkleWitnessResolver extends GraphqlModule<object> {
   public constructor(
     @inject("AsyncLinkedLeafStore")
-    private readonly treeStore: AsyncLinkedLeafStore
+    private readonly leafStore: AsyncLinkedLeafStore,
+    @inject("AsyncTreeStore")
+    private readonly treeStore: AsyncMerkleTreeStore
   ) {
     super();
   }
@@ -52,7 +55,10 @@ export class LinkedMerkleWitnessResolver 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 = await CachedLinkedLeafStore.new(this.treeStore);
+    const syncStore = await CachedLinkedLeafStore.new(
+      this.leafStore,
+      this.treeStore
+    );
 
     const tree = new LinkedMerkleTree(syncStore.treeStore, syncStore);
     await syncStore.preloadKey(BigInt(path));

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

@@ -9,6 +9,7 @@ import {
 } from "@proto-kit/common";
 import {
   AsyncLinkedLeafStore,
+  AsyncMerkleTreeStore,
   CachedLinkedLeafStore,
 } from "@proto-kit/sequencer";
 
@@ -92,7 +93,9 @@ export class LinkedTreeWitnessDTO {
 export class MerkleWitnessResolver extends GraphqlModule<object> {
   public constructor(
     @inject("AsyncLinkedLeafStore")
-    private readonly treeStore: AsyncLinkedLeafStore
+    private readonly leafStore: AsyncLinkedLeafStore,
+    @inject("AsyncTreeStore")
+    private readonly treeStore: AsyncMerkleTreeStore
   ) {
     super();
   }
@@ -102,7 +105,10 @@ 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 = await CachedLinkedLeafStore.new(this.treeStore);
+    const syncStore = await CachedLinkedLeafStore.new(
+      this.leafStore,
+      this.treeStore
+    );
     await syncStore.preloadKey(BigInt(path));
 
     const tree = new LinkedMerkleTree(syncStore.treeStore, syncStore);

packages/persistance/src/PrismaDatabaseConnection.ts

@@ -13,6 +13,7 @@ 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 { PrismaLinkedLeafStore } from "./services/prisma/PrismaLinkedLeafStore";
 
 export interface PrismaDatabaseConfig {
   // Either object-based config or connection string
@@ -57,7 +58,7 @@ export class PrismaDatabaseConnection
     StorageDependencyMinimumDependencies<{
       readonly prisma: PrismaDatabaseConnection;
     }>,
-    "blockTreeStore" | "asyncLinkedLeafStore" | "unprovenLinkedLeafStore"
+    "blockTreeStore" | "asyncTreeStore" | "unprovenTreeStore"
   > {
     return {
       asyncStateService: {
@@ -94,6 +95,26 @@ export class PrismaDatabaseConnection
       transactionStorage: {
         useClass: PrismaTransactionStorage,
       },
+
+      asyncLinkedLeafStore: {
+        useGenerated: (module) => {
+          return new PrismaLinkedLeafStore(
+            module.prisma,
+            module.prisma.tracer,
+            "batch"
+          );
+        },
+      },
+
+      unprovenLinkedLeafStore: {
+        useGenerated: (module) => {
+          return new PrismaLinkedLeafStore(
+            module.prisma,
+            module.prisma.tracer,
+            "block"
+          );
+        },
+      },
     };
   }
 

packages/persistance/src/PrismaRedisDatabase.ts

@@ -22,7 +22,6 @@ import {
   RedisConnectionModule,
   RedisTransaction,
 } from "./RedisConnection";
-import { PrismaLinkedLeafStore } from "./services/prisma/PrismaLinkedLeafStore";
 
 export interface PrismaRedisCombinedConfig {
   prisma: PrismaDatabaseConfig;
@@ -68,28 +67,6 @@ export class PrismaRedisDatabase
     return {
       ...PrismaDatabaseConnection.dependencies(),
       ...RedisConnectionModule.dependencies(),
-
-      asyncLinkedLeafStore: {
-        useGenerated: (module) => {
-          return new PrismaLinkedLeafStore(
-            module.prisma,
-            module.redis,
-            module.tracer,
-            "batch"
-          );
-        },
-      },
-
-      unprovenLinkedLeafStore: {
-        useGenerated: (module) => {
-          return new PrismaLinkedLeafStore(
-            module.prisma,
-            module.redis,
-            module.tracer,
-            "block"
-          );
-        },
-      },
     };
   }
 

packages/persistance/src/RedisConnection.ts

@@ -43,14 +43,14 @@ export class RedisConnectionModule
 
   public static dependencies(): Pick<
     StorageDependencyMinimumDependencies<{ redis: RedisConnectionModule }>,
-    "asyncMerkleStore" | "blockTreeStore" | "unprovenMerkleStore"
+    "blockTreeStore" | "asyncTreeStore" | "unprovenTreeStore"
   > {
     return {
-      asyncMerkleStore: {
+      asyncTreeStore: {
         useGenerated: ({ redis }) =>
           new RedisMerkleTreeStore(redis, redis.tracer),
       },
-      unprovenMerkleStore: {
+      unprovenTreeStore: {
         useGenerated: ({ redis }) =>
           new RedisMerkleTreeStore(redis, redis.tracer, "unproven"),
       },

packages/persistance/src/services/prisma/PrismaLinkedLeafStore.ts

@@ -1,36 +1,21 @@
 import { noop, StoredLeaf } from "@proto-kit/common";
-import { AsyncLinkedLeafStore, Tracer } from "@proto-kit/sequencer";
+import { AsyncLinkedLeafStore, trace, Tracer } from "@proto-kit/sequencer";
 import { injectable } from "tsyringe";
 import { Prisma } from "@prisma/client";
 
-import { PrismaConnection } from "../../PrismaDatabaseConnection";
-import { RedisMerkleTreeStore } from "../redis/RedisMerkleTreeStore";
-import { RedisConnection } from "../../RedisConnection";
+import type { PrismaConnection } from "../../PrismaDatabaseConnection";
 
 import { Decimal } from "./PrismaStateService";
 
 @injectable()
 export class PrismaLinkedLeafStore implements AsyncLinkedLeafStore {
   private cache: StoredLeaf[] = [];
 
-  private readonly redisMerkleStore: RedisMerkleTreeStore;
-
   public constructor(
     private readonly connection: PrismaConnection,
-    redisConnection: RedisConnection,
-    tracer: Tracer,
+    public readonly tracer: Tracer,
     private readonly mask: string = "base"
-  ) {
-    this.redisMerkleStore = new RedisMerkleTreeStore(
-      redisConnection,
-      tracer,
-      mask
-    );
-  }
-
-  public get treeStore() {
-    return this.redisMerkleStore;
-  }
+  ) {}
 
   private assertCacheEmpty() {
     if (this.cache.length > 0) {
@@ -42,6 +27,7 @@ export class PrismaLinkedLeafStore implements AsyncLinkedLeafStore {
     noop();
   }
 
+  @trace("LinkedLeafStore.commit")
   public async commit(): Promise<void> {
     if (this.cache.length > 0) {
       const data = this.cache.map((entry) => ({
@@ -126,6 +112,7 @@ export class PrismaLinkedLeafStore implements AsyncLinkedLeafStore {
       : undefined;
   }
 
+  @trace("getPreviousLeaves", ([paths]) => ({ numPaths: paths.length }))
   public async getPreviousLeavesAsync(paths: bigint[]) {
     this.assertCacheEmpty();
 

packages/sdk/src/query/StateServiceQueryModule.ts

@@ -4,6 +4,7 @@ import {
   CachedLinkedLeafStore,
   AsyncLinkedLeafStore,
   AppChainModule,
+  AsyncMerkleTreeStore,
 } from "@proto-kit/sequencer";
 import { Field } from "o1js";
 import { inject, injectable } from "tsyringe";
@@ -30,18 +31,25 @@ export class StateServiceQueryModule
     );
   }
 
-  public get treeStore(): AsyncLinkedLeafStore {
+  public get leafStore(): AsyncLinkedLeafStore {
     return this.sequencer.dependencyContainer.resolve("AsyncLinkedLeafStore");
   }
 
+  public get treeStore(): AsyncMerkleTreeStore {
+    return this.sequencer.dependencyContainer.resolve("AsyncTreeStore");
+  }
+
   public get(key: Field) {
     return this.asyncStateService.get(key);
   }
 
   public async merkleWitness(
     path: Field
   ): Promise<LinkedMerkleTreeReadWitness | undefined> {
-    const syncStore = await CachedLinkedLeafStore.new(this.treeStore);
+    const syncStore = await CachedLinkedLeafStore.new(
+      this.leafStore,
+      this.treeStore
+    );
     await syncStore.preloadKey(path.toBigInt());
 
     const tree = new LinkedMerkleTree(syncStore.treeStore, syncStore);

packages/sdk/test/stprover-emit-sts.test.ts

@@ -64,7 +64,7 @@ describe("StateTransition", () => {
 
     await appChain.start();
     appChain.setSigner(senderKey);
-  });
+  }, 30_000);
 
   afterEach(async () => {
     await appChain.close();

packages/sequencer/src/protocol/production/BatchProducerModule.ts

@@ -18,6 +18,7 @@ import type { Database } from "../../storage/Database";
 import { AsyncLinkedLeafStore } from "../../state/async/AsyncLinkedLeafStore";
 import { CachedLinkedLeafStore } from "../../state/lmt/CachedLinkedLeafStore";
 import { ensureNotBusy } from "../../helpers/BusyGuard";
+import { AsyncMerkleTreeStore } from "../../state/async/AsyncMerkleTreeStore";
 
 import { BlockProofSerializer } from "./tasks/serializers/BlockProofSerializer";
 import { BatchTracingService } from "./tracing/BatchTracingService";
@@ -47,7 +48,9 @@ const errors = {
 export class BatchProducerModule extends SequencerModule {
   public constructor(
     @inject("AsyncLinkedLeafStore")
-    private readonly merkleStore: AsyncLinkedLeafStore,
+    private readonly leafStore: AsyncLinkedLeafStore,
+    @inject("AsyncTreeStore")
+    private readonly merkleStore: AsyncMerkleTreeStore,
     @inject("BatchStorage") private readonly batchStorage: BatchStorage,
     @inject("Database")
     private readonly database: Database,
@@ -149,7 +152,10 @@ export class BatchProducerModule extends SequencerModule {
       throw errors.blockWithoutTxs();
     }
 
-    const merkleTreeStore = await CachedLinkedLeafStore.new(this.merkleStore);
+    const merkleTreeStore = await CachedLinkedLeafStore.new(
+      this.leafStore,
+      this.merkleStore
+    );
 
     const trace = await this.batchTraceService.traceBatch(
       blocks.map((block) => block),

packages/sequencer/src/protocol/production/sequencing/BlockProducerModule.ts

@@ -44,6 +44,8 @@ export class BlockProducerModule extends SequencerModule<BlockConfig> {
     private readonly unprovenStateService: AsyncStateService,
     @inject("UnprovenLinkedLeafStore")
     private readonly unprovenLinkedLeafStore: AsyncLinkedLeafStore,
+    @inject("UnprovenTreeStore")
+    private readonly unprovenTreeStore: AsyncMerkleTreeStore,
     @inject("BlockQueue")
     private readonly blockQueue: BlockQueue,
     @inject("TransactionStorage")
@@ -118,6 +120,7 @@ export class BlockProducerModule extends SequencerModule<BlockConfig> {
       await this.resultService.generateMetadataForNextBlock(
         block,
         this.unprovenLinkedLeafStore,
+        this.unprovenTreeStore,
         this.blockTreeStore,
         this.unprovenStateService
       );

packages/sequencer/src/protocol/production/sequencing/BlockResultService.ts

@@ -188,7 +188,8 @@ export class BlockResultService {
   }))
   public async generateMetadataForNextBlock(
     block: Block,
-    merkleTreeStore: AsyncLinkedLeafStore,
+    leafStore: AsyncLinkedLeafStore,
+    treeStore: AsyncMerkleTreeStore,
     blockHashTreeStore: AsyncMerkleTreeStore,
     stateService: AsyncStateService
   ): Promise<{
@@ -202,7 +203,7 @@ export class BlockResultService {
       block.beforeBlockStateTransitions
     );
 
-    const inMemoryStore = await CachedLinkedLeafStore.new(merkleTreeStore);
+    const inMemoryStore = await CachedLinkedLeafStore.new(leafStore, treeStore);
 
     const tree = await this.applyStateDiff(inMemoryStore, combinedDiff);
 

packages/sequencer/src/settlement/BridgingModule.ts

@@ -56,6 +56,7 @@ import { AsyncLinkedLeafStore } from "../state/async/AsyncLinkedLeafStore";
 import { CachedLinkedLeafStore } from "../state/lmt/CachedLinkedLeafStore";
 import { SettleableBatch } from "../storage/model/Batch";
 import { SequencerModule } from "../sequencer/builder/SequencerModule";
+import { AsyncMerkleTreeStore } from "../state/async/AsyncMerkleTreeStore";
 
 import type { SettlementModule } from "./SettlementModule";
 import { SettlementUtils } from "./utils/SettlementUtils";
@@ -114,6 +115,8 @@ export class BridgingModule extends SequencerModule<BridgingModuleConfig> {
     private readonly outgoingMessageCollector: OutgoingMessageCollector,
     @inject("AsyncLinkedLeafStore")
     private readonly linkedLeafStore: AsyncLinkedLeafStore,
+    @inject("AsyncTreeStore")
+    private readonly treeStore: AsyncMerkleTreeStore,
     @inject("FeeStrategy")
     private readonly feeStrategy: FeeStrategy,
     @inject("BaseLayer") private readonly baseLayer: MinaBaseLayer,
@@ -590,7 +593,10 @@ export class BridgingModule extends SequencerModule<BridgingModuleConfig> {
 
     const bridgeContract = this.createBridgeContract(bridgeAddress, tokenId);
 
-    const cachedStore = await CachedLinkedLeafStore.new(this.linkedLeafStore);
+    const cachedStore = await CachedLinkedLeafStore.new(
+      this.linkedLeafStore,
+      this.treeStore
+    );
     const tree = new LinkedMerkleTree(cachedStore.treeStore, cachedStore);
 
     // Create withdrawal batches and send them as L1 transactions

packages/sequencer/src/state/async/AsyncLinkedLeafStore.ts

@@ -1,10 +1,6 @@
 import { StoredLeaf } from "@proto-kit/common";
 
-import { AsyncMerkleTreeStore } from "./AsyncMerkleTreeStore";
-
 export interface AsyncLinkedLeafStore {
-  treeStore: AsyncMerkleTreeStore;
-
   openTransaction: () => Promise<void>;
 
   commit: () => Promise<void>;