Implemented unified state decorators

rpanic · rpanic · commit 80d54ff73555 · 2025-01-28T12:56:10.000+01:00
diff --git a/packages/common/src/utils.ts b/packages/common/src/utils.ts
@@ -255,3 +255,24 @@ export function batch<T>(
 
   return range(0, numBatches).map((i) => partitioned[i].map((x) => x[0]));
 }
+
+export type Reference<T> = {
+  set value(value: T);
+  get value(): T;
+};
+
+class ReferenceObject<T> {
+  public constructor(private internalValue: T) {}
+
+  get value() {
+    return this.internalValue;
+  }
+
+  set value(t: T) {
+    this.internalValue = t;
+  }
+}
+
+export function createReference<T>(initial: T): Reference<T> {
+  return new ReferenceObject(initial);
+}
diff --git a/packages/library/src/runtime/Balances.ts b/packages/library/src/runtime/Balances.ts
@@ -1,11 +1,6 @@
 import { EventsRecord, NoConfig } from "@proto-kit/common";
-import {
-  RuntimeModule,
-  runtimeMethod,
-  state,
-  runtimeModule,
-} from "@proto-kit/module";
-import { StateMap, assert } from "@proto-kit/protocol";
+import { RuntimeModule, runtimeMethod, runtimeModule } from "@proto-kit/module";
+import { StateMap, assert, state } from "@proto-kit/protocol";
 import { Field, PublicKey, Struct, Provable } from "o1js";
 
 import { UInt64 } from "../math/UInt64";
diff --git a/packages/library/src/runtime/Withdrawals.ts b/packages/library/src/runtime/Withdrawals.ts
@@ -1,10 +1,5 @@
-import {
-  RuntimeEvents,
-  runtimeModule,
-  RuntimeModule,
-  state,
-} from "@proto-kit/module";
-import { StateMap, Withdrawal } from "@proto-kit/protocol";
+import { RuntimeEvents, runtimeModule, RuntimeModule } from "@proto-kit/module";
+import { StateMap, Withdrawal, state } from "@proto-kit/protocol";
 import { Field, PublicKey, Struct } from "o1js";
 import { inject } from "tsyringe";
 
diff --git a/packages/module/src/index.ts b/packages/module/src/index.ts
@@ -4,7 +4,6 @@ export * from "./runtime/RuntimeModule";
 export * from "./runtime/RuntimeEnvironment";
 export * from "./runtime/Runtime";
 export * from "./state/InMemoryStateService";
-export * from "./state/decorator";
 export * from "./method/MethodParameterEncoder";
 export * from "./runtime/MethodIdResolver";
 export * from "./factories/MethodIdFactory";
diff --git a/packages/module/src/state/decorator.ts b/packages/module/src/state/decorator.ts
diff --git a/packages/protocol/src/hashing/mina-prefixes.ts b/packages/protocol/src/hashing/mina-prefixes.ts
@@ -0,0 +1,10 @@
+export const MINA_PREFIXES = {
+  event: "MinaZkappEvent******",
+  events: "MinaZkappEvents*****",
+  sequenceEvents: "MinaZkappSeqEvents**",
+} as const;
+
+export const MINA_SALTS = {
+  empty_actions: "MinaZkappActionsEmpty",
+  empty_events: "MinaZkappEventsEmpty",
+};
diff --git a/packages/protocol/src/hashing/protokit-prefixes.ts b/packages/protocol/src/hashing/protokit-prefixes.ts
@@ -0,0 +1,23 @@
+import padEnd from "lodash/padEnd";
+import mapValues from "lodash/mapValues";
+
+const length = 20;
+function padToHashPrefix(s: string): string {
+  if (s.length > 20) {
+    throw new Error(`Prefix string ${s} is too long (max ${length})`);
+  }
+  return padEnd(s, length, "*");
+}
+
+function padPrefixRecord<T extends Record<string, string>>(
+  record: T
+): {
+  [Key in keyof T]: string;
+} {
+  return mapValues(record, padToHashPrefix);
+}
+
+export const PROTOKIT_PREFIXES = padPrefixRecord({
+  STATE_PROTOCOL: "pk-protocol-state",
+  STATE_RUNTIME: "pk-runtime-state",
+});
diff --git a/packages/protocol/src/hooks/AccountStateHook.ts b/packages/protocol/src/hooks/AccountStateHook.ts
@@ -3,7 +3,7 @@ import { injectable } from "tsyringe";
 import { noop } from "@proto-kit/common";
 
 import { StateMap } from "../state/StateMap";
-import { protocolState } from "../state/protocol/ProtocolState";
+import { state } from "../state/protocol/ProtocolState";
 import {
   ProvableTransactionHook,
   BeforeTransactionHookArguments,
@@ -16,7 +16,7 @@ export class AccountState extends Struct({
 
 @injectable()
 export class AccountStateHook extends ProvableTransactionHook {
-  @protocolState() public accountState = StateMap.from<PublicKey, AccountState>(
+  @state() public accountState = StateMap.from<PublicKey, AccountState>(
     PublicKey,
     AccountState
   );
diff --git a/packages/protocol/src/index.ts b/packages/protocol/src/index.ts
@@ -60,3 +60,5 @@ export * from "./settlement/modules/NetworkStateSettlementModule";
 export * from "./settlement/messages/Deposit";
 export * from "./settlement/messages/Withdrawal";
 export { constants as ProtocolConstants } from "./Constants";
+export * from "./hashing/protokit-prefixes";
+export * from "./hashing/mina-prefixes";
diff --git a/packages/protocol/src/model/Path.ts b/packages/protocol/src/model/Path.ts
@@ -1,4 +1,5 @@
 import { Field, type FlexibleProvablePure, Poseidon } from "o1js";
+import { hashWithPrefix } from "@proto-kit/common";
 
 import { stringToField } from "../utils/utils";
 
@@ -21,10 +22,15 @@ export class Path {
    *
    * @param className
    * @param propertyKey
+   * @param prefix
    * @returns Field representation of class name + property name
    */
-  public static fromProperty(className: string, propertyKey: string): Field {
-    return Poseidon.hash([
+  public static fromProperty(
+    className: string,
+    propertyKey: string,
+    prefix: string
+  ): Field {
+    return hashWithPrefix(prefix, [
       Path.toField(className),
       Path.toField(propertyKey),
       Field(0),
diff --git a/packages/protocol/src/protocol/ProvableBlockHook.ts b/packages/protocol/src/protocol/ProvableBlockHook.ts
@@ -1,4 +1,5 @@
 import { Field } from "o1js";
+import { NoConfig } from "@proto-kit/common";
 
 import { NetworkState } from "../model/network/NetworkState";
 import { MethodPublicOutput } from "../model/MethodPublicOutput";
@@ -53,7 +54,7 @@ export function toAfterTransactionHookArgument(
 
 // Purpose is to build transition from -> to network state
 export abstract class ProvableBlockHook<
-  Config,
+  Config = NoConfig,
 > extends TransitioningProtocolModule<Config> {
   public abstract beforeBlock(
     networkState: NetworkState,
diff --git a/packages/protocol/src/settlement/contracts/BridgeContract.ts b/packages/protocol/src/settlement/contracts/BridgeContract.ts
@@ -23,6 +23,7 @@ import {
 } from "../messages/OutgoingMessageArgument";
 import { Path } from "../../model/Path";
 import { Withdrawal } from "../messages/Withdrawal";
+import { PROTOKIT_PREFIXES } from "../../hashing/protokit-prefixes";
 
 import type { SettlementContractType } from "./SettlementSmartContract";
 
@@ -137,7 +138,11 @@ export abstract class BridgeContractBase extends TokenContractV2 {
 
     const [withdrawalModule, withdrawalStateName] =
       BridgeContractBase.args.withdrawalStatePath;
-    const mapPath = Path.fromProperty(withdrawalModule, withdrawalStateName);
+    const mapPath = Path.fromProperty(
+      withdrawalModule,
+      withdrawalStateName,
+      PROTOKIT_PREFIXES.STATE_RUNTIME
+    );
 
     // Count account creation fee to return later, so that the sender can fund
     // those accounts with a separate AU
diff --git a/packages/protocol/src/state/protocol/ProtocolState.ts b/packages/protocol/src/state/protocol/ProtocolState.ts
@@ -1,5 +1,9 @@
+import { createReference, Reference } from "@proto-kit/common";
+
 import { State } from "../State";
 import { Path } from "../../model/Path";
+import { StateServiceProvider } from "../StateServiceProvider";
+import { PROTOKIT_PREFIXES } from "../../hashing/protokit-prefixes";
 import { TransitioningProtocolModule } from "../../protocol/TransitioningProtocolModule";
 
 const errors = {
@@ -9,54 +13,89 @@ const errors = {
       Did you forget to extend your block module with 'extends ...Hook'?`
     ),
 
-  missingProtocol: (className: string) =>
+  missingParent: (className: string, type: string, moduleType: string) =>
     new Error(
-      `Unable to provide 'procotol' for state, ${className} is missing a name. 
-      Did you forget to extend your block module with 'extends ...Hook'?`
+      `Unable to provide parent '${type}' for state, ${className} is missing a name. 
+      Did you forget to extend your module with 'extends ${moduleType}'?`
     ),
 };
 
-// TODO Homogenize @protocolState() towards a single @state() API
+export interface StatefulModule {
+  name?: string;
+  parent?: {
+    stateServiceProvider: StateServiceProvider;
+  };
+}
+
+export function createStateGetter<TargetModule extends StatefulModule>(
+  target: TargetModule,
+  propertyKey: string,
+  valueReference: Reference<State<unknown> | undefined>,
+  prefix: string,
+  debugInfo: { parentName: string; baseModuleNames: string }
+) {
+  return () => {
+    const { value } = valueReference;
+    // Short-circuit this to return the state in case its already initialized
+    if (value !== undefined && value.path !== undefined) {
+      return value;
+    }
+
+    if (target.name === undefined) {
+      throw errors.missingName(target.constructor.name);
+    }
+
+    if (!target.parent) {
+      throw errors.missingParent(
+        target.constructor.name,
+        debugInfo.parentName,
+        debugInfo.baseModuleNames
+      );
+    }
+
+    const path = Path.fromProperty(target.name, propertyKey, prefix);
+    if (value) {
+      value.path = path;
+      value.stateServiceProvider = target.parent.stateServiceProvider;
+    }
+    return value;
+  };
+}
 
 /**
  * Decorates a runtime module property as state, passing down some
  * underlying values to improve developer experience.
  */
-export function protocolState() {
-  return <
-    TargetTransitioningModule extends TransitioningProtocolModule<unknown>,
-  >(
+export function state() {
+  return <TargetTransitioningModule extends StatefulModule>(
     target: TargetTransitioningModule,
     propertyKey: string
   ) => {
-    let value: State<unknown> | undefined;
+    const stateReference = createReference<State<unknown> | undefined>(
+      undefined
+    );
+
+    const isProtocol = target instanceof TransitioningProtocolModule;
+    const statePrefix = isProtocol
+      ? PROTOKIT_PREFIXES.STATE_PROTOCOL
+      : PROTOKIT_PREFIXES.STATE_RUNTIME;
+    const debugInfo = isProtocol
+      ? { parentName: "protocol", baseModuleNames: "...Hook" }
+      : { parentName: "runtime", baseModuleNames: "RuntimeModule" };
 
     Object.defineProperty(target, propertyKey, {
       enumerable: true,
 
-      get: function get() {
-        // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
-        const self = this as TargetTransitioningModule;
-
-        if (self.name === undefined) {
-          throw errors.missingName(self.constructor.name);
-        }
-
-        if (!self.parent) {
-          throw errors.missingProtocol(self.constructor.name);
-        }
-
-        // TODO Add Prefix?
-        const path = Path.fromProperty(self.name, propertyKey);
-        if (value) {
-          value.path = path;
-          value.stateServiceProvider = self.parent.stateServiceProvider;
-        }
-        return value;
-      },
+      get: createStateGetter(
+        target,
+        propertyKey,
+        stateReference,
+        statePrefix,
+        debugInfo
+      ),
 
       set: (newValue: State<unknown>) => {
-        value = newValue;
+        stateReference.value = newValue;
       },
     });
   };
diff --git a/packages/protocol/src/utils/MinaPrefixedProvableHashList.ts b/packages/protocol/src/utils/MinaPrefixedProvableHashList.ts
diff --git a/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts b/packages/sequencer/src/protocol/production/sequencing/TransactionExecutionService.ts
diff --git a/packages/sequencer/src/settlement/BridgingModule.ts b/packages/sequencer/src/settlement/BridgingModule.ts
diff --git a/packages/stack/src/scripts/graphql/server.ts b/packages/stack/src/scripts/graphql/server.ts
