@proto-kit/module 0.1.1-develop.191 → 0.1.1-develop.2024

package/src/runtime/Runtime.ts

@@ -1,32 +1,54 @@
-// eslint-disable-next-line max-len
-/* eslint-disable @typescript-eslint/no-explicit-any,@typescript-eslint/no-unsafe-assignment,max-lines */
-import { Experimental } from "snarkyjs";
-import { DependencyContainer, injectable } from "tsyringe";
+/* eslint-disable @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-unsafe-argument */
+import { ZkProgram } from "o1js";
+import { container, DependencyContainer, injectable } from "tsyringe";
 import {
   StringKeyOf,
   ModuleContainer,
-  ModulesConfig,
   ModulesRecord,
   TypedClass,
   ZkProgrammable,
   PlainZkProgram,
-  WithZkProgrammable,
   AreProofsEnabled,
+  ChildContainerProvider,
+  CompilableModule,
+  CompileRegistry,
 } from "@proto-kit/common";
-import { MethodPublicOutput } from "@proto-kit/protocol";
+import {
+  MethodPublicOutput,
+  StateServiceProvider,
+  SimpleAsyncStateService,
+  RuntimeMethodExecutionContext,
+  RuntimeTransaction,
+  NetworkState,
+} from "@proto-kit/protocol";
 
 import {
   combineMethodName,
   isRuntimeMethod,
+  runtimeMethodTypeMetadataKey,
   toWrappedMethod,
-  WrappedMethod,
-} from "../method/runtimeMethod.js";
-import { StateService } from "../state/InMemoryStateService.js";
-import { StateServiceProvider } from "../state/StateServiceProvider";
+  AsyncWrappedMethod,
+} from "../method/runtimeMethod";
 import { MethodIdFactory } from "../factories/MethodIdFactory";
 
-import { RuntimeModule } from "./RuntimeModule.js";
+import { RuntimeModule } from "./RuntimeModule";
 import { MethodIdResolver } from "./MethodIdResolver";
+import { RuntimeEnvironment } from "./RuntimeEnvironment";
+
+export function getAllPropertyNames(obj: any) {
+  let currentPrototype: any | undefined = obj;
+  let keys: (string | symbol)[] = [];
+  // if primitive (primitives still have keys) skip the first iteration
+  if (!(obj instanceof Object)) {
+    currentPrototype = Object.getPrototypeOf(obj);
+  }
+  // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
+  while (currentPrototype) {
+    keys = keys.concat(Reflect.ownKeys(currentPrototype));
+    currentPrototype = Object.getPrototypeOf(currentPrototype);
+  }
+  return keys;
+}
 
 /**
  * Record of modules accepted by the Runtime module container.
@@ -43,66 +65,56 @@ const errors = {
     new Error(`Unable to find method with id ${methodKey}`),
 };
 
-/**
- * Definition / required arguments for the Runtime class
- */
-export interface RuntimeDefinition<Modules extends RuntimeModulesRecord> {
-  state: StateService;
-  modules: Modules;
-  config?: ModulesConfig<Modules>;
-}
-
 export class RuntimeZkProgrammable<
-  Modules extends RuntimeModulesRecord
+  Modules extends RuntimeModulesRecord,
 > extends ZkProgrammable<undefined, MethodPublicOutput> {
-  // eslint-disable-next-line @typescript-eslint/no-use-before-define
   public constructor(public runtime: Runtime<Modules>) {
     super();
   }
 
-  public get appChain() {
-    return this.runtime.appChain;
+  public get areProofsEnabled() {
+    return this.runtime.areProofsEnabled;
   }
 
-  public zkProgramFactory(): PlainZkProgram<undefined, MethodPublicOutput> {
+  public zkProgramFactory(): PlainZkProgram<undefined, MethodPublicOutput>[] {
     type Methods = Record<
       string,
       {
         privateInputs: any;
-        method: WrappedMethod;
+        method: AsyncWrappedMethod;
       }
     >;
     // We need to use explicit type annotations here,
     // therefore we can't use destructuring
-    // eslint-disable-next-line max-len
-    // eslint-disable-next-line @typescript-eslint/no-use-before-define,prefer-destructuring,putout/putout
+
+    // eslint-disable-next-line prefer-destructuring
     const runtime: Runtime<Modules> = this.runtime;
 
+    const MAXIMUM_METHODS_PER_ZK_PROGRAM = 8;
+
     const runtimeMethods = runtime.runtimeModuleNames.reduce<Methods>(
       (allMethods, runtimeModuleName) => {
-        runtime.isValidModuleName(
-          runtime.definition.modules,
-          runtimeModuleName
-        );
+        runtime.isValidModuleName(runtime.definition, runtimeModuleName);
 
         /**
          * Couldnt find a better way to circumvent the type assertion
          * regarding resolving only known modules. We assert in the line above
          * but we cast it to any anyways to satisfy the proof system.
          */
-        // eslint-disable-next-line max-len
+
         // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
         const runtimeModule = runtime.resolve(runtimeModuleName as any);
 
-        // eslint-disable-next-line max-len
         // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
         const modulePrototype = Object.getPrototypeOf(runtimeModule) as Record<
           string,
-          (...args: unknown[]) => unknown
+          // Technically not all methods have to be async, but for this context it's ok
+          (...args: unknown[]) => Promise<unknown>
         >;
 
-        const modulePrototypeMethods =
-          Object.getOwnPropertyNames(modulePrototype);
+        const modulePrototypeMethods = getAllPropertyNames(runtimeModule).map(
+          (method) => method.toString()
+        );
 
         const moduleMethods = modulePrototypeMethods.reduce<Methods>(
           (allModuleMethods, methodName) => {
@@ -112,14 +124,17 @@ export class RuntimeZkProgrammable<
                 methodName
               );
               const method = modulePrototype[methodName];
-              const wrappedMethod = Reflect.apply(
-                toWrappedMethod,
+              const invocationType = Reflect.getMetadata(
+                runtimeMethodTypeMetadataKey,
                 runtimeModule,
-                [methodName, method]
+                methodName
               );
 
-              // eslint-disable-next-line no-warning-comments
-              // TODO: find out how to import the Tuple type
+              const wrappedMethod: AsyncWrappedMethod = Reflect.apply(
+                toWrappedMethod,
+                runtimeModule,
+                [methodName, method, { invocationType }]
+              );
 
               const privateInputs = Reflect.getMetadata(
                 "design:paramtypes",
@@ -151,30 +166,100 @@ export class RuntimeZkProgrammable<
     );
 
     const sortedRuntimeMethods = Object.fromEntries(
-      // eslint-disable-next-line @typescript-eslint/require-array-sort-compare
       Object.entries(runtimeMethods).sort()
     );
 
-    const program = Experimental.ZkProgram({
-      publicOutput: MethodPublicOutput,
-      methods: sortedRuntimeMethods,
-    });
-
-    const SelfProof = Experimental.ZkProgram.Proof(program);
-
-    const methods = Object.keys(sortedRuntimeMethods).reduce<
-      Record<string, any>
-    >((boundMethods, methodName) => {
-      boundMethods[methodName] = program[methodName].bind(program);
-      return boundMethods;
-    }, {});
-
-    return {
-      compile: program.compile.bind(program),
-      verify: program.verify.bind(program),
-      Proof: SelfProof,
-      methods,
+    const splitRuntimeMethods = () => {
+      const buckets: Array<
+        Record<
+          string,
+          {
+            privateInputs: any;
+            method: AsyncWrappedMethod;
+          }
+        >
+      > = [];
+      Object.entries(sortedRuntimeMethods).forEach(
+        async ([methodName, method]) => {
+          let methodAdded = false;
+          for (const bucket of buckets) {
+            if (buckets.length === 0) {
+              const record: Record<
+                string,
+                {
+                  privateInputs: any;
+                  method: AsyncWrappedMethod;
+                }
+              > = {};
+              record[methodName] = method;
+              buckets.push(record);
+              methodAdded = true;
+              break;
+            } else if (
+              Object.keys(bucket).length <=
+              MAXIMUM_METHODS_PER_ZK_PROGRAM - 1
+            ) {
+              bucket[methodName] = method;
+              methodAdded = true;
+              break;
+            }
+          }
+          if (!methodAdded) {
+            const record: Record<
+              string,
+              {
+                privateInputs: any;
+                method: AsyncWrappedMethod;
+              }
+            > = {};
+            record[methodName] = method;
+            buckets.push(record);
+          }
+        }
+      );
+      return buckets;
     };
+
+    return splitRuntimeMethods().map((bucket, index) => {
+      const name = `RuntimeProgram-${index}`;
+      const wrappedBucket = Object.fromEntries(
+        Object.entries(bucket).map(([methodName, methodDef]) => [
+          methodName,
+          {
+            privateInputs: methodDef.privateInputs,
+            method: async (...args: any[]) => {
+              const publicOutput = await methodDef.method(...args);
+              return { publicOutput };
+            },
+          },
+        ])
+      );
+
+      const program = ZkProgram({
+        name,
+        publicOutput: MethodPublicOutput,
+        methods: wrappedBucket,
+      });
+
+      const SelfProof = ZkProgram.Proof(program);
+
+      const methods = Object.keys(bucket).reduce<Record<string, any>>(
+        (boundMethods, methodName) => {
+          boundMethods[methodName] = program[methodName].bind(program);
+          return boundMethods;
+        },
+        {}
+      );
+
+      return {
+        name,
+        compile: program.compile.bind(program),
+        verify: program.verify.bind(program),
+        analyzeMethods: program.analyzeMethods.bind(program),
+        Proof: SelfProof,
+        methods,
+      };
+    });
   }
 }
 
@@ -185,59 +270,59 @@ export class RuntimeZkProgrammable<
 @injectable()
 export class Runtime<Modules extends RuntimeModulesRecord>
   extends ModuleContainer<Modules>
-  implements WithZkProgrammable<undefined, MethodPublicOutput>
+  implements RuntimeEnvironment, CompilableModule
 {
   public static from<Modules extends RuntimeModulesRecord>(
-    definition: RuntimeDefinition<Modules>
-  ) {
-    return new Runtime(definition);
+    definition: Modules
+  ): TypedClass<Runtime<Modules>> {
+    return class RuntimeScoped extends Runtime<Modules> {
+      public constructor() {
+        super(definition);
+      }
+    };
   }
 
   // runtime modules composed into a ZkProgram
-  public program?: ReturnType<typeof Experimental.ZkProgram>;
+  public program?: ReturnType<typeof ZkProgram>;
 
-  public definition: RuntimeDefinition<Modules>;
+  // No idea why we have to do this, but if we don't re-define it here,
+  // js can't access it from the superclass somehow
+  public definition: Modules;
 
   public zkProgrammable: ZkProgrammable<undefined, MethodPublicOutput>;
 
-  // eslint-disable-next-line no-warning-comments
-  // TODO DI
-  private readonly stateServiceProviderInstance = new StateServiceProvider(
-    this.definition.state
-  );
-
   /**
    * Creates a new Runtime from the provided config
    *
    * @param modules - Configuration object for the constructed Runtime
    */
-  public constructor(definition: RuntimeDefinition<Modules>) {
+  public constructor(definition: Modules) {
     super(definition);
     this.definition = definition;
     this.zkProgrammable = new RuntimeZkProgrammable<Modules>(this);
-
-    // this.registerDependencyFactories([MethodIdFactory]);
   }
 
-  // eslint-disable-next-line no-warning-comments
   // TODO Remove after changing DFs to type-based approach
-  public start() {
-    this.registerValue({
-      Runtime: this,
-    });
-    this.registerDependencyFactories([MethodIdFactory]);
-  }
+  public create(childContainerProvider: ChildContainerProvider) {
+    super.create(childContainerProvider);
 
-  public get appChain(): AreProofsEnabled | undefined {
-    return this.container.resolve<AreProofsEnabled>("AppChain");
+    this.useDependencyFactory(MethodIdFactory);
   }
 
-  public get stateService(): StateService {
-    return this.stateServiceProviderInstance.stateService;
+  public get areProofsEnabled(): AreProofsEnabled | undefined {
+    return this.container.resolve<AreProofsEnabled>("AreProofsEnabled");
   }
 
   public get stateServiceProvider(): StateServiceProvider {
-    return this.stateServiceProviderInstance;
+    return this.container.resolve<StateServiceProvider>("StateServiceProvider");
+  }
+
+  public get stateService(): SimpleAsyncStateService {
+    return this.stateServiceProvider.stateService;
+  }
+
+  public get methodIdResolver(): MethodIdResolver {
+    return this.container.resolve<MethodIdResolver>("MethodIdResolver");
   }
 
   /**
@@ -253,28 +338,26 @@ export class Runtime<Modules extends RuntimeModulesRecord>
    */
   public getMethodById(
     methodId: bigint
-  ): ((...args: unknown[]) => unknown) | undefined {
-    const methodDescriptor = this.container
-      .resolve<MethodIdResolver>("MethodIdResolver")
-      .getMethodNameFromId(methodId);
+  ): ((...args: unknown[]) => Promise<unknown>) | undefined {
+    const methodDescriptor =
+      this.methodIdResolver.getMethodNameFromId(methodId);
 
     if (methodDescriptor === undefined) {
       return undefined;
     }
     const [moduleName, methodName] = methodDescriptor;
 
-    this.isValidModuleName(this.definition.modules, moduleName);
+    this.assertIsValidModuleName(moduleName);
     const module = this.resolve(moduleName);
 
-    // eslint-disable-next-line max-len
-    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions,@typescript-eslint/no-unsafe-member-access
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
     const method = (module as any)[methodName];
     if (method === undefined) {
       throw errors.methodNotFound(`${moduleName}.${methodName}`);
     }
 
     // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
-    return (method as (...args: unknown[]) => unknown).bind(module);
+    return (method as (...args: unknown[]) => Promise<unknown>).bind(module);
   }
 
   /**
@@ -289,7 +372,7 @@ export class Runtime<Modules extends RuntimeModulesRecord>
     containedModule: InstanceType<Modules[StringKeyOf<Modules>]>
   ) {
     containedModule.name = moduleName;
-    containedModule.runtime = this;
+    containedModule.parent = this;
 
     super.decorateModule(moduleName, containedModule);
   }
@@ -298,6 +381,16 @@ export class Runtime<Modules extends RuntimeModulesRecord>
    * @returns A list of names of all the registered module names
    */
   public get runtimeModuleNames() {
-    return Object.keys(this.definition.modules);
+    return this.moduleNames;
+  }
+
+  public async compile(registry: CompileRegistry) {
+    const context = container.resolve(RuntimeMethodExecutionContext);
+    context.setup({
+      transaction: RuntimeTransaction.dummyTransaction(),
+      networkState: NetworkState.empty(),
+    });
+    return await this.zkProgrammable.compile(registry);
   }
 }
+/* eslint-enable @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-unsafe-argument */

package/src/runtime/RuntimeEnvironment.ts

@@ -0,0 +1,16 @@
+import { AreProofsEnabled, WithZkProgrammable } from "@proto-kit/common";
+import {
+  MethodPublicOutput,
+  SimpleAsyncStateService,
+  StateServiceProvider,
+} from "@proto-kit/protocol";
+
+import { MethodIdResolver } from "./MethodIdResolver";
+
+export interface RuntimeEnvironment
+  extends WithZkProgrammable<undefined, MethodPublicOutput> {
+  get areProofsEnabled(): AreProofsEnabled | undefined;
+  get stateService(): SimpleAsyncStateService;
+  get stateServiceProvider(): StateServiceProvider;
+  get methodIdResolver(): MethodIdResolver;
+}

package/src/runtime/RuntimeModule.ts

@@ -1,40 +1,68 @@
-import { ConfigurableModule, Presets } from "@proto-kit/common";
+import { ConfigurableModule, NoConfig, Presets } from "@proto-kit/common";
 import { container, injectable } from "tsyringe";
-import { NetworkState, RuntimeTransaction } from "@proto-kit/protocol";
-
 import {
+  NetworkState,
+  RuntimeTransaction,
   RuntimeMethodExecutionContext,
   RuntimeMethodExecutionData,
-} from "../method/RuntimeMethodExecutionContext";
-import { StateService } from "../state/InMemoryStateService";
+  RuntimeMethodExecutionDataStruct,
+} from "@proto-kit/protocol";
+import { FlexibleProvablePure, Provable, Bool } from "o1js";
+
 import { runtimeMethodNamesMetadataKey } from "../method/runtimeMethod";
+import {
+  OutgoingMessages,
+  OutgoingMessagesRecord,
+} from "../messages/OutgoingMessages";
 
-import type {
-  Runtime,
-  RuntimeDefinition,
-  RuntimeModulesRecord,
-} from "./Runtime";
+import { RuntimeEnvironment } from "./RuntimeEnvironment";
 
 const errors = {
   inputDataNotSet: () => new Error("Input data for runtime execution not set"),
 };
 
-/**
- * This type exists to carry over certain runtime properties
- * to runtime modules, until we can inject them through DI.
- */
-export interface PartialRuntime
-  extends Pick<Runtime<RuntimeModulesRecord>, "zkProgrammable"> {
-  definition: Pick<RuntimeDefinition<RuntimeModulesRecord>, "state">;
+type EventRecord = Record<string, FlexibleProvablePure<any>>;
+
+type InferProvable<T extends FlexibleProvablePure<any>> =
+  T extends Provable<infer U> ? U : never;
+
+export class RuntimeEvents<Events extends EventRecord> {
+  public constructor(private readonly events: Events) {}
+
+  public emitIf<Key extends keyof Events>(
+    condition: Bool,
+    eventName: Key,
+    event: InferProvable<Events[Key]>
+  ) {
+    if (this.events === undefined) {
+      throw new Error(
+        "'events' property not defined, make sure to define the event types on your runtimemodule"
+      );
+    }
+    const eventType: FlexibleProvablePure<any> = this.events[eventName];
+    if (typeof eventName !== "string") {
+      throw new Error("Only string");
+    }
+    return container
+      .resolve(RuntimeMethodExecutionContext)
+      .addEvent(eventType, event, eventName, condition);
+  }
 
-  get stateService(): StateService;
+  public emit<Key extends keyof Events>(
+    eventName: Key,
+    event: InferProvable<Events[Key]>
+  ) {
+    this.emitIf(Bool(true), eventName, event);
+  }
 }
 
 /**
  * Base class for runtime modules providing the necessary utilities.
  */
 @injectable()
-export class RuntimeModule<Config> extends ConfigurableModule<Config> {
+export class RuntimeModule<
+  Config = NoConfig,
+> extends ConfigurableModule<Config> {
   public static presets: Presets<unknown> = {};
 
   /**
@@ -51,9 +79,11 @@ export class RuntimeModule<Config> extends ConfigurableModule<Config> {
 
   public name?: string;
 
-  public test?: number;
+  public parent?: RuntimeEnvironment;
 
-  public runtime?: Runtime<RuntimeModulesRecord>;
+  public events?: RuntimeEvents<any> = undefined;
+
+  public messages?: OutgoingMessages<OutgoingMessagesRecord> = undefined;
 
   public constructor() {
     super();
@@ -65,15 +95,17 @@ export class RuntimeModule<Config> extends ConfigurableModule<Config> {
     this.runtimeMethodNames = methodNames ?? [];
   }
 
-  private getInputs(): RuntimeMethodExecutionData {
-    const { input } = container.resolve<RuntimeMethodExecutionContext>(
-      RuntimeMethodExecutionContext
-    );
-
-    if (input === undefined) {
-      throw errors.inputDataNotSet();
-    }
-    return input;
+  public getInputs(): RuntimeMethodExecutionData {
+    return Provable.witness(RuntimeMethodExecutionDataStruct, () => {
+      const { input } = container.resolve<RuntimeMethodExecutionContext>(
+        RuntimeMethodExecutionContext
+      );
+      if (input === undefined) {
+        throw errors.inputDataNotSet();
+      }
+
+      return input;
+    });
   }
 
   public get transaction(): RuntimeTransaction {

package/src/state/InMemoryStateService.ts

@@ -1,27 +1,23 @@
-import { Field } from "snarkyjs";
-
-export interface StateService {
-  get: (key: Field) => Field[] | undefined;
-  set: (key: Field, value: Field[] | undefined) => void;
-}
+import { Field } from "o1js";
+import { SimpleAsyncStateService } from "@proto-kit/protocol";
 
 /**
- * Naive implementation of a StateService for testing purposes
+ * Naive implementation of an in-memory variant of the StateService interface
  */
-export class InMemoryStateService implements StateService {
-  public values: Record<string, Field[] | undefined> = {};
+export class InMemoryStateService implements SimpleAsyncStateService {
+  /**
+   * This mapping container null values if the specific entry has been deleted.
+   * This is used by the CachedState service to keep track of deletions
+   */
+  public values: Record<string, Field[] | null> = {};
 
-  public get(key: Field): Field[] | undefined {
-    return this.values[key.toString()];
+  public async get(key: Field): Promise<Field[] | undefined> {
+    return this.values[key.toString()] ?? undefined;
   }
 
-  public set(key: Field, value: Field[] | undefined) {
-    if (
-      value === undefined &&
-      Object.prototype.hasOwnProperty.call(this.values, key.toString())
-    ) {
-      // eslint-disable-next-line @typescript-eslint/no-dynamic-delete
-      delete this.values[key.toString()];
+  public async set(key: Field, value: Field[] | undefined) {
+    if (value === undefined) {
+      this.values[key.toString()] = null;
     } else {
       this.values[key.toString()] = value;
     }