@proto-kit/module 0.1.1-develop.165 → 0.1.1-develop.1665

package/src/runtime/MethodIdResolver.ts

@@ -0,0 +1,108 @@
+import { filterNonUndefined } from "@proto-kit/common";
+import { stringToField, RuntimeMethodIdMapping } from "@proto-kit/protocol";
+import { Poseidon } from "o1js";
+import { inject, injectable } from "tsyringe";
+
+import {
+  RuntimeMethodInvocationType,
+  runtimeMethodTypeMetadataKey,
+} from "../method/runtimeMethod";
+
+import type { Runtime, RuntimeModulesRecord } from "./Runtime";
+
+/**
+ * Please see `getMethodId` to learn more about
+ * methodId encoding
+ */
+@injectable()
+export class MethodIdResolver {
+  private readonly dictionary: {
+    [key: string]: { moduleName: string; methodName: string };
+  } = {};
+
+  public constructor(
+    @inject("Runtime") private readonly runtime: Runtime<RuntimeModulesRecord>
+  ) {
+    this.dictionary = runtime.runtimeModuleNames.reduce<
+      Record<string, { moduleName: string; methodName: string }>
+    >((dict, moduleName) => {
+      this.runtime.assertIsValidModuleName(moduleName);
+
+      runtime.resolve(moduleName).runtimeMethodNames.forEach((methodName) => {
+        dict[this.getMethodId(moduleName, methodName).toString()] = {
+          moduleName,
+          methodName,
+        };
+      });
+
+      return dict;
+    }, {});
+  }
+
+  /**
+   * The purpose of this method is to provide a dictionary where
+   * we can look up properties like methodId and invocationType
+   * for each runtimeMethod using their module name and method name
+   */
+  public methodIdMap(): RuntimeMethodIdMapping {
+    const methodIdResolver =
+      this.runtime.dependencyContainer.resolve<MethodIdResolver>(
+        "MethodIdResolver"
+      );
+
+    const rawMappings = this.runtime.moduleNames.flatMap((moduleName) => {
+      const module = this.runtime.resolve(moduleName);
+      return module.runtimeMethodNames.map((method) => {
+        // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+        const type = Reflect.getMetadata(
+          runtimeMethodTypeMetadataKey,
+          module,
+          method
+        ) as RuntimeMethodInvocationType | undefined;
+
+        if (type !== undefined) {
+          return {
+            name: `${moduleName}.${method}`,
+            methodId: methodIdResolver.getMethodId(moduleName, method),
+            type,
+          } as const;
+        }
+
+        return undefined;
+      });
+    });
+
+    return rawMappings
+      .filter(filterNonUndefined)
+      .reduce<RuntimeMethodIdMapping>((acc, entry) => {
+        acc[entry.name] = {
+          methodId: entry.methodId,
+          type: entry.type,
+        };
+        return acc;
+      }, {});
+  }
+
+  public getMethodNameFromId(methodId: bigint): [string, string] | undefined {
+    const methodPath = this.dictionary[methodId.toString()];
+
+    if (methodPath === undefined) {
+      return undefined;
+    }
+
+    const { moduleName, methodName } = methodPath;
+
+    this.runtime.assertIsValidModuleName(moduleName);
+
+    return [moduleName, methodName];
+  }
+
+  public getMethodId(moduleName: string, methodName: string): bigint {
+    this.runtime.assertIsValidModuleName(moduleName);
+
+    return Poseidon.hash([
+      stringToField(moduleName),
+      stringToField(methodName),
+    ]).toBigInt();
+  }
+}

package/src/runtime/Runtime.ts

@@ -1,34 +1,54 @@
-// eslint-disable-next-line max-len
-/* eslint-disable @typescript-eslint/no-explicit-any,@typescript-eslint/no-unsafe-assignment */
-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 {
-  fieldToString,
-  stringToField,
   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";
-
-import { RuntimeModule } from "./RuntimeModule.js";
+  AsyncWrappedMethod,
+} from "../method/runtimeMethod";
+import { MethodIdFactory } from "../factories/MethodIdFactory";
+
+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.
@@ -45,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) => {
@@ -114,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",
@@ -153,30 +166,87 @@ 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 program = ZkProgram({
+        name,
+        publicOutput: MethodPublicOutput,
+        methods: bucket,
+      });
+
+      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,
+      };
+    });
   }
 }
 
@@ -187,51 +257,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.registerValue({
-    //   Runtime: this,
-    // });
   }
 
-  public get appChain(): AreProofsEnabled | undefined {
-    return this.container.resolve<AreProofsEnabled>("AppChain");
+  // TODO Remove after changing DFs to type-based approach
+  public create(childContainerProvider: ChildContainerProvider) {
+    super.create(childContainerProvider);
+
+    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");
   }
 
   /**
@@ -245,52 +323,43 @@ export class Runtime<Modules extends RuntimeModulesRecord>
    * @param methodId The encoded name of the method to call.
    * Encoding: "stringToField(module.name) << 128 + stringToField(method-name)"
    */
-  public getMethodById(methodId: bigint): (...args: unknown[]) => unknown {
-    const [moduleName, methodName] = this.getMethodNameFromId(methodId);
+  public getMethodById(
+    methodId: bigint
+  ): ((...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);
-  }
-
-  public getMethodNameFromId(methodId: bigint): [string, string] {
-    // eslint-disable-next-line @typescript-eslint/no-magic-numbers
-    const moduleName = fieldToString(methodId >> 128n);
-    // eslint-disable-next-line @typescript-eslint/no-magic-numbers
-    const methodName = fieldToString(methodId % 2n ** 128n);
-
-    return [moduleName, methodName];
-  }
-
-  public getMethodId(moduleName: string, methodName: string): bigint {
-    return (
-      // eslint-disable-next-line @typescript-eslint/no-magic-numbers
-      (stringToField(moduleName).toBigInt() << 128n) +
-      stringToField(methodName).toBigInt()
-    );
+    return (method as (...args: unknown[]) => Promise<unknown>).bind(module);
   }
 
   /**
    * Add a name and other respective properties required by RuntimeModules,
    * that come from the current Runtime
    *
-   * @param name - Name of the runtime module to decorate
+   * @param moduleName - Name of the runtime module to decorate
+   * @param containedModule
    */
   public decorateModule(
     moduleName: StringKeyOf<Modules>,
     containedModule: InstanceType<Modules[StringKeyOf<Modules>]>
   ) {
     containedModule.name = moduleName;
-    containedModule.runtime = this;
+    containedModule.parent = this;
 
     super.decorateModule(moduleName, containedModule);
   }
@@ -299,6 +368,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;
+}