Native masking for State and Tree masks

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/persistance/src/PrismaDatabaseConnection.ts

@@ -7,12 +7,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
@@ -54,12 +54,9 @@ export class PrismaDatabaseConnection
 
   public dependencies(): OmitKeys<
     StorageDependencyMinimumDependencies,
-    "asyncMerkleStore" | "blockTreeStore" | "unprovenMerkleStore"
+    "blockTreeStore" | "treeStoreCreator"
   > {
     return {
-      asyncStateService: {
-        useFactory: () => new PrismaStateService(this, this.tracer, "batch"),
-      },
       batchStorage: {
         useClass: PrismaBatchStore,
       },
@@ -78,8 +75,8 @@ export class PrismaDatabaseConnection
       transactionStorage: {
         useClass: PrismaTransactionStorage,
       },
-      unprovenStateService: {
-        useFactory: () => new PrismaStateService(this, this.tracer, "block"),
+      stateServiceCreator: {
+        useClass: PrismaStateServiceCreator,
       },
     };
   }

packages/persistance/src/RedisConnection.ts

@@ -8,6 +8,7 @@ import { DependencyFactory } from "@proto-kit/common";
 import isArray from "lodash/isArray";
 
 import { RedisMerkleTreeStore } from "./services/redis/RedisMerkleTreeStore";
+import { RedisTreeStoreCreator } from "./creators/RedisTreeStoreCreator";
 
 export interface RedisConnectionConfig {
   host: string;
@@ -44,20 +45,16 @@ export class RedisConnectionModule
 
   public dependencies(): Pick<
     StorageDependencyMinimumDependencies,
-    "asyncMerkleStore" | "blockTreeStore" | "unprovenMerkleStore"
+    "blockTreeStore" | "treeStoreCreator"
   > {
     return {
-      asyncMerkleStore: {
-        useFactory: () => new RedisMerkleTreeStore(this, this.tracer),
-      },
-      unprovenMerkleStore: {
-        useFactory: () =>
-          new RedisMerkleTreeStore(this, this.tracer, "unproven"),
-      },
       blockTreeStore: {
         useFactory: () =>
           new RedisMerkleTreeStore(this, this.tracer, "blockHash"),
       },
+      treeStoreCreator: {
+        useClass: RedisTreeStoreCreator,
+      },
     };
   }
 

packages/persistance/src/creators/PrismaStateServiceCreator.ts

@@ -0,0 +1,38 @@
+import {
+  AsyncStateService,
+  StateServiceCreator,
+  Tracer,
+} from "@proto-kit/sequencer";
+import { inject, injectable } from "tsyringe";
+
+import { PrismaStateService } from "../services/prisma/PrismaStateService";
+import type { PrismaConnection } from "../PrismaDatabaseConnection";
+
+@injectable()
+export class PrismaStateServiceCreator implements StateServiceCreator {
+  public constructor(
+    @inject("Database") private readonly connection: PrismaConnection,
+    @inject("Tracer") private readonly tracer: Tracer
+  ) {}
+
+  public async createMask(
+    name: string,
+    parent: string
+  ): Promise<AsyncStateService> {
+    return new PrismaStateService(this.connection, this.tracer, name, parent);
+  }
+
+  public getMask(name: string): AsyncStateService {
+    return new PrismaStateService(this.connection, this.tracer, name);
+  }
+
+  public async mergeIntoParent(name: string): Promise<void> {
+    const service = new PrismaStateService(this.connection, this.tracer, name);
+    await service.mergeIntoParent();
+  }
+
+  public async drop(name: string): Promise<void> {
+    const service = new PrismaStateService(this.connection, this.tracer, name);
+    await service.drop();
+  }
+}

packages/persistance/src/creators/RedisTreeStoreCreator.ts

@@ -0,0 +1,28 @@
+import {
+  AsyncMerkleTreeStore,
+  InMemoryMerkleTreeStoreMask,
+  MaskGraph,
+  Tracer,
+  TreeStoreCreator,
+} from "@proto-kit/sequencer";
+import { inject, injectable } from "tsyringe";
+
+import { RedisMerkleTreeStore } from "../services/redis/RedisMerkleTreeStore";
+import type { RedisConnection } from "../RedisConnection";
+
+@injectable()
+export class RedisTreeStoreCreator
+  extends MaskGraph<
+    AsyncMerkleTreeStore,
+    RedisMerkleTreeStore,
+    InMemoryMerkleTreeStoreMask
+  >
+  implements TreeStoreCreator
+{
+  public constructor(
+    @inject("Database") connection: RedisConnection,
+    @inject("Tracer") tracer: Tracer
+  ) {
+    super(new RedisMerkleTreeStore(connection, tracer, "base"));
+  }
+}

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

@@ -3,6 +3,7 @@ import {
   StateEntry,
   Tracer,
   trace,
+  MaskName,
 } from "@proto-kit/sequencer";
 import { Field } from "o1js";
 import { Prisma } from "@prisma/client";
@@ -25,56 +26,88 @@ const Decimal = Prisma.Decimal.clone({
 export class PrismaStateService implements AsyncStateService {
   private cache: StateEntry[] = [];
 
-  private maskId?: number;
+  private maskId?: { id: number; parentId?: number };
 
   /**
    * @param connection
    * @param tracer
    * @param mask A indicator to which masking level the values belong.
    * This name has to be unique
-   * @param parent
+   * @param parentName
    */
   public constructor(
     private readonly connection: PrismaConnection,
     public readonly tracer: Tracer,
     private readonly mask: string,
-    private readonly parent?: number
+    private readonly parentName?: string
   ) {}
 
-  private async getMaskId(): Promise<number> {
+  private async getMaskId(): Promise<{ id: number; parentId?: number }> {
     if (this.maskId === undefined) {
-      this.maskId = await this.initializeMask(this.mask, this.parent);
+      this.maskId = await this.initializeMask(this.mask, this.parentName);
     }
     return this.maskId;
   }
 
-  private async initializeMask(mask: string, parent?: number): Promise<number> {
+  private async initializeMask(
+    mask: string,
+    parentName?: string
+  ): Promise<{ id: number; parentId?: number }> {
     const { prismaClient } = this.connection;
 
     const found = await prismaClient.mask.findFirst({
       where: {
         name: mask,
-        parent,
+      },
+      include: {
+        parentMask: true,
       },
     });
 
     if (found === null) {
+      // Find parent id
+      let parentId: number | undefined = undefined;
+      if (parentName !== undefined) {
+        const parent = await prismaClient.mask.findFirst({
+          where: {
+            name: parentName,
+          },
+        });
+        if (parent === null) {
+          throw new Error(`Parent mask with name ${parentName} not found`);
+        }
+
+        parentId = parent.id;
+      } else if (mask !== MaskName.base()) {
+        throw new Error(
+          "Can't initialize mask that's not the base using a null-parent "
+        );
+      }
+
+      // Create mask
       const createdMask = await prismaClient.mask.create({
         data: {
-          parent,
+          parent: parentId,
           name: mask,
         },
       });
-      return createdMask.id;
+      return {
+        id: createdMask.id,
+        parentId: parentId,
+      };
     }
-    return found.id;
+
+    return {
+      id: found.id,
+      parentId: found.parent ?? undefined,
+    };
   }
 
   @trace("db.state.commit")
   public async commit(): Promise<void> {
     const { prismaClient } = this.connection;
 
-    const maskId = await this.getMaskId();
+    const { id: maskId } = await this.getMaskId();
 
     const data = this.cache
       .filter((entry) => entry.value !== undefined)
@@ -100,19 +133,19 @@ export class PrismaStateService implements AsyncStateService {
   }
 
   public async getMany(keys: Field[]): Promise<StateEntry[]> {
-    const maskId = await this.getMaskId();
+    const { id: maskId } = await this.getMaskId();
     const paths = keys.map((key) => new Decimal(key.toString()));
 
     const records: {
       path: Prisma.Decimal;
+      // TODO This could potentially be non-null, but should be tested
       values: Prisma.Decimal[] | null;
     }[] = await this.connection.prismaClient.$queryRaw(
       readState(maskId, paths)
     );
 
     return records.map((record) => ({
       key: Field(record.path.toFixed()),
-      // TODO Figure out why that is nullable
       value: record.values?.map((x) => Field(x.toFixed())) ?? [],
     }));
   }
@@ -131,16 +164,23 @@ export class PrismaStateService implements AsyncStateService {
   }
 
   public async createMask(name: string): Promise<AsyncStateService> {
-    const maskId = await this.getMaskId();
-    return new PrismaStateService(this.connection, this.tracer, name, maskId);
+    // We only call this to make sure this mask actually exists, therefore that the
+    // relation can be satisfied
+    await this.getMaskId();
+    return new PrismaStateService(
+      this.connection,
+      this.tracer,
+      name,
+      this.mask
+    );
   }
 
   public async mergeIntoParent(): Promise<void> {
-    const maskId = await this.getMaskId();
+    const { id: maskId, parentId } = await this.getMaskId();
 
     const client = this.connection.prismaClient;
 
-    if (this.parent !== undefined) {
+    if (parentId !== undefined) {
       // Rough strategy here:
       // 1. Delete all entries that are bound to be overwritten from the parent mask
       // 2. Update this mask's entries to parent mask id
@@ -158,7 +198,7 @@ export class PrismaStateService implements AsyncStateService {
             parent: maskId,
           },
           data: {
-            parent: this.parent,
+            parent: parentId,
           },
         }),
         client.mask.delete({
@@ -173,7 +213,7 @@ export class PrismaStateService implements AsyncStateService {
   }
 
   public async drop(): Promise<void> {
-    const maskId = await this.getMaskId();
+    const { id: maskId } = await this.getMaskId();
 
     await this.connection.prismaClient.state.deleteMany({
       where: {

packages/persistance/src/services/redis/RedisMerkleTreeStore.ts

@@ -1,5 +1,6 @@
 import {
   AsyncMerkleTreeStore,
+  InMemoryMerkleTreeStoreMask,
   MerkleTreeNode,
   MerkleTreeNodeQuery,
   trace,
@@ -12,10 +13,12 @@ import type { RedisConnection } from "../../RedisConnection";
 export class RedisMerkleTreeStore implements AsyncMerkleTreeStore {
   private cache: MerkleTreeNode[] = [];
 
+  public readonly name = "base";
+
   public constructor(
     private readonly connection: RedisConnection,
     public readonly tracer: Tracer,
-    private readonly mask: string = "base"
+    private readonly mask: string
   ) {}
 
   private getKey(node: MerkleTreeNodeQuery): string {
@@ -85,4 +88,8 @@ export class RedisMerkleTreeStore implements AsyncMerkleTreeStore {
     // });
     // console.log(`Reduced ${concat.length} to ${this.cache.length} items to write`)
   }
+
+  public async createMask(name: string) {
+    return new InMemoryMerkleTreeStoreMask(this, name);
+  }
 }