Added Prisma and Redis Creators

Author: rpanic

packages/persistance/src/PrismaDatabaseConnection.ts

@@ -12,6 +12,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 { PrismaStateServiceCreator } from "./creators/PrismaStateServiceCreator";
 
 export interface PrismaDatabaseConfig {
   // Either object-based config or connection string
@@ -49,12 +50,9 @@ export class PrismaDatabaseConnection
 
   public dependencies(): OmitKeys<
     StorageDependencyMinimumDependencies,
-    "asyncMerkleStore" | "blockTreeStore" | "unprovenMerkleStore"
+    "blockTreeStore" | "treeStoreCreator"
   > {
     return {
-      asyncStateService: {
-        useFactory: () => new PrismaStateService(this, "batch"),
-      },
       batchStorage: {
         useClass: PrismaBatchStore,
       },
@@ -73,8 +71,8 @@ export class PrismaDatabaseConnection
       transactionStorage: {
         useClass: PrismaTransactionStorage,
       },
-      unprovenStateService: {
-        useFactory: () => new PrismaStateService(this, "block"),
+      stateServiceCreator: {
+        useClass: PrismaStateServiceCreator,
       },
     };
   }

packages/persistance/src/RedisConnection.ts

@@ -7,6 +7,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;
@@ -39,18 +40,15 @@ export class RedisConnectionModule
 
   public dependencies(): Pick<
     StorageDependencyMinimumDependencies,
-    "asyncMerkleStore" | "blockTreeStore" | "unprovenMerkleStore"
+    "blockTreeStore" | "treeStoreCreator"
   > {
     return {
-      asyncMerkleStore: {
-        useFactory: () => new RedisMerkleTreeStore(this),
-      },
-      unprovenMerkleStore: {
-        useFactory: () => new RedisMerkleTreeStore(this, "unproven"),
-      },
       blockTreeStore: {
         useFactory: () => new RedisMerkleTreeStore(this, "blockHash"),
       },
+      treeStoreCreator: {
+        useClass: RedisTreeStoreCreator,
+      },
     };
   }
 

packages/persistance/src/creators/PrismaStateServiceCreator.ts

@@ -0,0 +1,33 @@
+import { AsyncStateService, StateServiceCreator } 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
+  ) {}
+
+  public async createMask(
+    name: string,
+    parent: string
+  ): Promise<AsyncStateService> {
+    return new PrismaStateService(this.connection, name, parent);
+  }
+
+  public getMask(name: string): AsyncStateService {
+    return new PrismaStateService(this.connection, name);
+  }
+
+  public async mergeIntoParent(name: string): Promise<void> {
+    const service = new PrismaStateService(this.connection, name);
+    await service.mergeIntoParent();
+  }
+
+  public async drop(name: string): Promise<void> {
+    const service = new PrismaStateService(this.connection, name);
+    await service.drop();
+  }
+}

packages/persistance/src/creators/RedisTreeStoreCreator.ts

@@ -0,0 +1,15 @@
+import {
+  AsyncMerkleTreeStore,
+  InMemoryMerkleTreeStoreMask,
+  MaskGraph,
+  TreeStoreCreator,
+} from "@proto-kit/sequencer";
+import { RedisMerkleTreeStore } from "../services/redis/RedisMerkleTreeStore";
+
+export class RedisTreeStoreCreator
+  extends MaskGraph<
+    AsyncMerkleTreeStore,
+    RedisMerkleTreeStore,
+    InMemoryMerkleTreeStoreMask
+  >
+  implements TreeStoreCreator {}

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

@@ -1,4 +1,4 @@
-import { AsyncStateService, StateEntry } from "@proto-kit/sequencer";
+import { AsyncStateService, MaskName, StateEntry } from "@proto-kit/sequencer";
 import { Field } from "o1js";
 import { Prisma } from "@prisma/client";
 import { noop } from "@proto-kit/common";
@@ -20,53 +20,85 @@ 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 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,
     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,
+    };
   }
 
   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)
@@ -92,7 +124,7 @@ 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: {
@@ -123,16 +155,18 @@ export class PrismaStateService implements AsyncStateService {
   }
 
   public async createMask(name: string): Promise<AsyncStateService> {
-    const maskId = await this.getMaskId();
-    return new PrismaStateService(this.connection, 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, 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
@@ -150,7 +184,7 @@ export class PrismaStateService implements AsyncStateService {
             parent: maskId,
           },
           data: {
-            parent: this.parent,
+            parent: parentId,
           },
         }),
         client.mask.delete({
@@ -165,7 +199,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,
 } from "@proto-kit/sequencer";
@@ -10,9 +11,11 @@ import type { RedisConnection } from "../../RedisConnection";
 export class RedisMerkleTreeStore implements AsyncMerkleTreeStore {
   private cache: MerkleTreeNode[] = [];
 
+  public readonly name = "base";
+
   public constructor(
     private readonly connection: RedisConnection,
-    private readonly mask: string = "base"
+    private readonly mask: string
   ) {}
 
   private getKey(node: MerkleTreeNodeQuery): string {
@@ -80,4 +83,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);
+  }
 }