Merge pull request #68 from proto-kit/feature/state-service-refactor

Refactored State Services into more clear structure

packages/protocol/src/utils/merkletree/InMemoryMerkleTreeStorage.ts

@@ -1,7 +1,7 @@
 import { MerkleTreeStore } from "./MerkleTreeStore";
 
 export class InMemoryMerkleTreeStorage implements MerkleTreeStore {
-  protected readonly nodes: {
+  protected nodes: {
     [key: number]: {
       [key: string]: bigint;
     };

packages/protocol/src/utils/merkletree/MerkleTreeStore.ts

@@ -1,13 +1,3 @@
-export interface AsyncMerkleTreeStore {
-  openTransaction: () => void;
-
-  commit: () => void;
-
-  setNodeAsync: (key: bigint, level: number, value: bigint) => Promise<void>;
-
-  getNodeAsync: (key: bigint, level: number) => Promise<bigint | undefined>;
-}
-
 export interface MerkleTreeStore {
   setNode: (key: bigint, level: number, value: bigint) => void;
 

packages/sequencer/src/index.ts

@@ -19,12 +19,6 @@ export * from "./worker/worker/FlowTaskWorker";
 export * from "./worker/worker/LocalTaskWorkerModule";
 export * from "./protocol/baselayer/BaseLayer";
 export * from "./protocol/baselayer/NoopBaseLayer";
-export * from "./protocol/production/execution/CachedStateService";
-export * from "./protocol/production/execution/DummyStateService";
-export * from "./protocol/production/execution/MerkleStoreWitnessProvider";
-export * from "./protocol/production/state/AsyncStateService";
-export * from "./protocol/production/tasks/providers/PreFilledStateService";
-export * from "./protocol/production/tasks/providers/PreFilledWitnessProvider";
 export * from "./protocol/production/tasks/BlockProvingTask";
 export * from "./protocol/production/tasks/CompileRegistry";
 export * from "./protocol/production/tasks/RuntimeProvingTask";
@@ -45,3 +39,13 @@ export * from "./storage/StorageDependencyFactory";
 export * from "./helpers/query/QueryTransportModule";
 export * from "./helpers/query/QueryBuilderFactory";
 export * from "./helpers/query/NetworkStateQuery";
+export * from "./state/prefilled/PreFilledStateService";
+export * from "./state/prefilled/PreFilledWitnessProvider";
+export * from "./state/async/AsyncMerkleTreeStore";
+export * from "./state/async/AsyncStateService";
+export * from "./state/merkle/CachedMerkleTreeStore";
+export * from "./state/merkle/SyncCachedMerkleTreeStore";
+export * from "./state/state/DummyStateService";
+export * from "./state/state/SyncCachedStateService";
+export * from "./state/state/CachedStateService";
+export * from "./state/MerkleStoreWitnessProvider";

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

@@ -1,6 +1,5 @@
 import { inject } from "tsyringe";
 import {
-  AsyncMerkleTreeStore,
   BlockProverPublicInput,
   BlockProverPublicOutput,
   DefaultProvableHashList,
@@ -22,15 +21,16 @@ import {
   ComputedBlock,
   ComputedBlockTransaction,
 } from "../../storage/model/Block";
+import { CachedStateService } from "../../state/state/CachedStateService";
+import { CachedMerkleTreeStore } from "../../state/merkle/CachedMerkleTreeStore";
+import { AsyncStateService } from "../../state/async/AsyncStateService";
+import { AsyncMerkleTreeStore } from "../../state/async/AsyncMerkleTreeStore";
 
-import { AsyncStateService } from "./state/AsyncStateService";
-import { CachedStateService } from "./execution/CachedStateService";
+import { BlockProverParameters } from "./tasks/BlockProvingTask";
 import { StateTransitionProofParameters } from "./tasks/StateTransitionTaskParameters";
 import { RuntimeProofParameters } from "./tasks/RuntimeTaskParameters";
 import { TransactionTraceService } from "./TransactionTraceService";
 import { BlockTaskFlowService } from "./BlockTaskFlowService";
-import { BlockProverParameters } from "./tasks/BlockProvingTask";
-import { CachedMerkleTreeStore } from "./execution/CachedMerkleTreeStore";
 
 export interface StateRecord {
   [key: string]: Field[] | undefined;

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

@@ -33,15 +33,13 @@ import { PendingTransaction } from "../../mempool/PendingTransaction";
 import { distinctByString } from "../../helpers/utils";
 import { ComputedBlockTransaction } from "../../storage/model/Block";
 
-import { CachedStateService } from "./execution/CachedStateService";
 import type { StateRecord, TransactionTrace } from "./BlockProducerModule";
-import { DummyStateService } from "./execution/DummyStateService";
 import { StateTransitionProofParameters } from "./tasks/StateTransitionTaskParameters";
-import { AsyncStateService } from "./state/AsyncStateService";
-import {
-  CachedMerkleTreeStore,
-  SyncCachedMerkleTreeStore,
-} from "./execution/CachedMerkleTreeStore";
+import { DummyStateService } from "../../state/state/DummyStateService";
+import { CachedMerkleTreeStore } from "../../state/merkle/CachedMerkleTreeStore";
+import { CachedStateService } from "../../state/state/CachedStateService";
+import { SyncCachedMerkleTreeStore } from "../../state/merkle/SyncCachedMerkleTreeStore";
+import { AsyncStateService } from "../../state/async/AsyncStateService";
 
 const errors = {
   methodIdNotFound: (methodId: string) =>

packages/sequencer/src/protocol/production/tasks/BlockProvingTask.ts

@@ -6,12 +6,11 @@ import {
   MethodPublicOutput,
   Protocol,
   ProtocolModulesRecord,
-  // eslint-disable-next-line @typescript-eslint/no-shadow
   ReturnType,
   StateTransitionProof,
   StateTransitionProvable,
 } from "@proto-kit/protocol";
-import { Field, Proof, Provable } from "o1js";
+import { Field, Proof } from "o1js";
 import { Runtime } from "@proto-kit/module";
 import { inject, injectable, Lifecycle, scoped } from "tsyringe";
 import { ProvableMethodExecutionContext } from "@proto-kit/common";
@@ -27,7 +26,7 @@ import { Task } from "../../../worker/flow/Task";
 
 import { CompileRegistry } from "./CompileRegistry";
 import { DecodedState, JSONEncodableState } from "./RuntimeTaskParameters";
-import { PreFilledStateService } from "./providers/PreFilledStateService";
+import { PreFilledStateService } from "../../../state/prefilled/PreFilledStateService";
 
 type RuntimeProof = Proof<undefined, MethodPublicOutput>;
 type BlockProof = Proof<BlockProverPublicInput, BlockProverPublicOutput>;
@@ -192,6 +191,22 @@ export class BlockProvingTask
     );
   }
 
+  private async executeWithPrefilledStateService<Return>(
+    startingState: DecodedState,
+    callback: () => Promise<Return>
+  ): Promise<Return> {
+    const prefilledStateService = new PreFilledStateService(startingState);
+    this.protocol.stateServiceProvider.setCurrentStateService(
+      prefilledStateService
+    );
+
+    const returnValue = await callback();
+
+    this.protocol.stateServiceProvider.popCurrentStateService();
+
+    return returnValue;
+  }
+
   public async compute(
     input: PairingDerivedInput<
       StateTransitionProof,
@@ -202,23 +217,24 @@ export class BlockProvingTask
     const stateTransitionProof = input.input1;
     const runtimeProof = input.input2;
 
-    const prefilledStateService = new PreFilledStateService(
-      input.params.startingState
-    );
-    this.protocol.stateServiceProvider.setCurrentStateService(
-      prefilledStateService
+    await this.executeWithPrefilledStateService(
+      input.params.startingState,
+      // eslint-disable-next-line putout/putout
+      async () => {
+        this.blockProver.proveTransaction(
+          input.params.publicInput,
+          stateTransitionProof,
+          runtimeProof,
+          input.params.executionData
+        );
+      }
     );
 
-    this.blockProver.proveTransaction(
-      input.params.publicInput,
-      stateTransitionProof,
-      runtimeProof,
-      input.params.executionData
+    return await this.executeWithPrefilledStateService(
+      input.params.startingState,
+      async () =>
+        await this.executionContext.current().result.prove<BlockProof>()
     );
-
-    this.protocol.stateServiceProvider.popCurrentStateService();
-
-    return await this.executionContext.current().result.prove<BlockProof>();
   }
 
   // eslint-disable-next-line sonarjs/no-identical-functions

packages/sequencer/src/protocol/production/tasks/RuntimeProvingTask.ts

@@ -19,7 +19,7 @@ import {
   RuntimeProofParameters,
   RuntimeProofParametersSerializer,
 } from "./RuntimeTaskParameters";
-import { PreFilledStateService } from "./providers/PreFilledStateService";
+import { PreFilledStateService } from "../../../state/prefilled/PreFilledStateService";
 
 type RuntimeProof = Proof<undefined, MethodPublicOutput>;
 

packages/sequencer/src/protocol/production/tasks/StateTransitionTask.ts

@@ -25,7 +25,7 @@ import {
   StateTransitionParametersSerializer,
   StateTransitionProofParameters,
 } from "./StateTransitionTaskParameters";
-import { PreFilledWitnessProvider } from "./providers/PreFilledWitnessProvider";
+import { PreFilledWitnessProvider } from "../../../state/prefilled/PreFilledWitnessProvider";
 import { CompileRegistry } from "./CompileRegistry";
 
 @injectable()

packages/sequencer/src/state/StateServices.md

@@ -0,0 +1,32 @@
+## State Service Architecture
+
+The state services are built using a abstraction stack, where every service that is below the root is based on the parent.
+All of this services have the ability to base off one another to allow for diff-based changes.
+You can think of it a bit like git, where every commit is only the diff between the parents state and the next state.
+
+*But*, in this architecture there is no such thing as branching, it is strictly linear and behaves almost like a traditional stack.
+That means, for example, that you shouldn't change state in a service that has one or more children, as these changes might not get reflected instantly in the children's perception of state.
+
+Now I will go over all the different kinds of stateservices
+
+Btw all of this also applies to to merkle tree stores
+
+### AsyncStateService
+
+This is always the base for the whole stack and is meant to be implemented by the actual persistence layer. 
+That means primarily the DB (but is also implemented by the in-memory state service).
+It's functions are async-based in order to enable external DB APIs
+
+### CachedStateService
+
+It receives a AsyncStateService as a parent and can build on top of it.
+It "caches" in the sense that it can preload state entries from the parent (asynchronously) and then lets circuits among others perform synchronous operations on them.
+It additionally caches write operations that can then later be merged back into the parent (or thrown away). 
+
+### SyncCachedStateService
+
+These are the same as CachedStateService, with the difference that it accepts a CachedStateService and requires no preloading.
+
+### PreFilledStateService
+
+Pre-filled with only a part of the data needed. This is mostly used in Task implementation where all state needed is passed as args.

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

@@ -0,0 +1,9 @@
+export interface AsyncMerkleTreeStore {
+  openTransaction: () => void;
+
+  commit: () => void;
+
+  setNodeAsync: (key: bigint, level: number, value: bigint) => Promise<void>;
+
+  getNodeAsync: (key: bigint, level: number) => Promise<bigint | undefined>;
+}

packages/sequencer/src/protocol/production/state/AsyncStateService.ts → packages/sequencer/src/state/async/AsyncStateService.ts

@@ -6,6 +6,11 @@ import { Field } from "o1js";
  * CachedStateService to preload keys for In-Circuit usage.
  */
 export interface AsyncStateService {
+  openTransaction: () => void;
+
+  commit: () => void;
+
   setAsync: (key: Field, value: Field[] | undefined) => Promise<void>;
+
   getAsync: (key: Field) => Promise<Field[] | undefined>;
 }

packages/sequencer/src/protocol/production/execution/CachedMerkleTreeStore.ts → packages/sequencer/src/state/merkle/CachedMerkleTreeStore.ts

@@ -1,9 +1,9 @@
 import {
-  AsyncMerkleTreeStore,
-  InMemoryMerkleTreeStorage, MerkleTreeStore,
+  InMemoryMerkleTreeStorage,
   RollupMerkleTree
 } from "@proto-kit/protocol";
 import { log, noop } from "@proto-kit/common";
+import { AsyncMerkleTreeStore } from "../async/AsyncMerkleTreeStore";
 
 export class CachedMerkleTreeStore
   extends InMemoryMerkleTreeStorage
@@ -123,42 +123,4 @@ export class CachedMerkleTreeStore
       this.getNode(key, level) ?? (await this.parent.getNodeAsync(key, level))
     );
   }
-}
-
-export class SyncCachedMerkleTreeStore extends InMemoryMerkleTreeStorage {
-  private writeCache: {
-    [key: number]: {
-      [key: string]: bigint;
-    };
-  } = {};
-
-  public constructor(private readonly parent: MerkleTreeStore) {
-    super();
-  }
-
-  public getNode(key: bigint, level: number): bigint | undefined {
-    return super.getNode(key, level) ?? this.parent.getNode(key, level);
-  }
-
-  public setNode(key: bigint, level: number, value: bigint) {
-    super.setNode(key, level, value);
-    (this.writeCache[level] ??= {})[key.toString()] = value;
-  }
-
-  public mergeIntoParent() {
-    if (Object.keys(this.writeCache).length === 0) {
-      return;
-    }
-
-    const { height } = RollupMerkleTree;
-    const nodes = this.writeCache
-
-    Array.from({ length: height }).forEach((ignored, level) =>
-      Object.entries(nodes[level]).forEach((entry) => {
-        this.parent.setNode(BigInt(entry[0]), level, entry[1]);
-      })
-    );
-
-    this.writeCache = {}
-  }
 }

packages/sequencer/src/state/merkle/SyncCachedMerkleTreeStore.ts

@@ -0,0 +1,35 @@
+import {
+  InMemoryMerkleTreeStorage, MerkleTreeStore,
+  RollupMerkleTree
+} from "@proto-kit/protocol";
+
+export class SyncCachedMerkleTreeStore extends InMemoryMerkleTreeStorage {
+  public constructor(private readonly parent: MerkleTreeStore) {
+    super();
+  }
+
+  public getNode(key: bigint, level: number): bigint | undefined {
+    return super.getNode(key, level) ?? this.parent.getNode(key, level);
+  }
+
+  public setNode(key: bigint, level: number, value: bigint) {
+    super.setNode(key, level, value);
+  }
+
+  public mergeIntoParent() {
+    if (Object.keys(this.nodes).length === 0) {
+      return;
+    }
+
+    const { height } = RollupMerkleTree;
+    const { nodes } = this;
+
+    Array.from({ length: height }).forEach((ignored, level) =>
+      Object.entries(nodes[level]).forEach((entry) => {
+        this.parent.setNode(BigInt(entry[0]), level, entry[1]);
+      })
+    );
+
+    this.nodes = {}
+  }
+}

packages/sequencer/src/protocol/production/execution/CachedStateService.ts → packages/sequencer/src/state/state/CachedStateService.ts

@@ -1,8 +1,8 @@
 import { Field } from "o1js";
-import { log } from "@proto-kit/common";
+import { log, noop } from "@proto-kit/common";
 import { InMemoryStateService } from "@proto-kit/module";
 
-import { AsyncStateService } from "../state/AsyncStateService";
+import { AsyncStateService } from "../async/AsyncStateService";
 
 const errors = {
   parentIsUndefined: () => new Error("Parent StateService is undefined"),
@@ -28,6 +28,14 @@ export class CachedStateService
     }
   }
 
+  public commit(): void {
+    noop();
+  }
+
+  public openTransaction(): void {
+    noop();
+  }
+
   public async preloadKey(key: Field) {
     // Only preload it if it hasn't been preloaded previously
     if (this.parent !== undefined && this.get(key) === undefined) {