@proto-kit/module 0.1.1-develop.153

package/src/method/runtimeMethod.ts

@@ -0,0 +1,192 @@
+import { Bool, Field, FlexibleProvable, Struct } from "snarkyjs";
+import { container } from "tsyringe";
+import {
+  StateTransition,
+  DefaultProvableHashList,
+  ProvableStateTransition,
+  MethodPublicOutput,
+} from "@proto-kit/protocol";
+import { DecoratedMethod, toProver, ZkProgrammable } from "@proto-kit/common";
+
+import type { RuntimeModule } from "../runtime/RuntimeModule.js";
+
+import { RuntimeMethodExecutionContext } from "./RuntimeMethodExecutionContext.js";
+
+const errors = {
+  runtimeNotProvided: (name: string) =>
+    new Error(`Runtime was not provided for module: ${name}`),
+
+  methodInputsNotProvided: () =>
+    new Error(
+      "Method execution inputs not provided, provide them via context.inputs"
+    ),
+};
+
+export function toStateTransitionsHash(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  stateTransitions: StateTransition<any>[]
+) {
+  const stateTransitionsHashList = new DefaultProvableHashList(
+    ProvableStateTransition
+  );
+
+  return stateTransitions
+    .map((stateTransition) => stateTransition.toProvable())
+    .reduce(
+      (allStateTransitionsHashList, stateTransition) =>
+        allStateTransitionsHashList.push(stateTransition),
+      stateTransitionsHashList
+    )
+    .toField();
+}
+
+// eslint-disable-next-line etc/prefer-interface
+export type WrappedMethod = (...args: unknown[]) => MethodPublicOutput;
+
+export function toWrappedMethod(
+  this: RuntimeModule<unknown>,
+  methodName: string,
+  moduleMethod: (...args: unknown[]) => unknown
+) {
+  const executionContext = container.resolve<RuntimeMethodExecutionContext>(
+    RuntimeMethodExecutionContext
+  );
+
+  const wrappedMethod: WrappedMethod = (...args): MethodPublicOutput => {
+    Reflect.apply(moduleMethod, this, args);
+    const {
+      result: { stateTransitions, status },
+      input,
+    } = executionContext.current();
+
+    const stateTransitionsHash = toStateTransitionsHash(stateTransitions);
+
+    if (input === undefined) {
+      throw errors.methodInputsNotProvided();
+    }
+
+    const transactionHash = input.transaction.hash();
+    const networkStateHash = input.networkState.hash();
+
+    return new MethodPublicOutput({
+      stateTransitionsHash,
+      status,
+      transactionHash,
+      networkStateHash,
+    });
+  };
+
+  Object.defineProperty(wrappedMethod, "name", {
+    value: `wrapped_${methodName}`,
+    writable: false,
+  });
+
+  return wrappedMethod;
+}
+
+export function combineMethodName(
+  runtimeModuleName: string,
+  methodName: string
+) {
+  return `${runtimeModuleName}.${methodName}`;
+}
+
+export const runtimeMethodMetadataKey = "yab-method";
+
+/**
+ * Checks the metadata of the provided runtime module and its method,
+ * to see if it has been decorated with @runtimeMethod()
+ *
+ * @param target - Runtime module to check
+ * @param propertyKey - Name of the method to check in the prior runtime module
+ * @returns - If the provided method name is a runtime method or not
+ */
+export function isRuntimeMethod(
+  target: RuntimeModule<unknown>,
+  propertyKey: string
+) {
+  return Boolean(
+    Reflect.getMetadata(runtimeMethodMetadataKey, target, propertyKey)
+  );
+}
+
+export function runtimeMethod() {
+  return (
+    target: RuntimeModule<unknown>,
+    methodName: string,
+    descriptor: PropertyDescriptor
+  ) => {
+    const executionContext = container.resolve<RuntimeMethodExecutionContext>(
+      RuntimeMethodExecutionContext
+    );
+
+    Reflect.defineMetadata(runtimeMethodMetadataKey, true, target, methodName);
+
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    const simulatedMethod = descriptor.value as DecoratedMethod;
+
+    descriptor.value = function value(
+      this: RuntimeModule<unknown>,
+      ...args: FlexibleProvable<unknown>[]
+    ) {
+      const constructorName = this.constructor.name;
+
+      /**
+       * If its a top level method call, wrap it into a wrapped method,
+       * since it'll be turned into a real/mock prover in provableMethod().
+       *
+       * Otherwise provableMethod() will just call the originalMethod provided
+       * if method is not called at the top level.
+       */
+      const simulatedWrappedMethod = Reflect.apply(toWrappedMethod, this, [
+        methodName,
+        simulatedMethod,
+      ]);
+
+      /**
+       * Before the prover runs, make sure it is operating on the correct
+       * RuntimeMethodExecutionContext state, meaning it enters and exits
+       * the context properly.
+       */
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      async function prover(this: ZkProgrammable<any, any>) {
+        executionContext.beforeMethod(constructorName, methodName, args);
+        const innerProver = toProver(
+          combineMethodName(constructorName, methodName),
+          simulatedWrappedMethod,
+          false,
+          ...args
+        ).bind(this);
+        // eslint-disable-next-line @typescript-eslint/init-declarations
+        let result: Awaited<ReturnType<typeof innerProver>>;
+        try {
+          // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+          result = await Reflect.apply(innerProver, this, args);
+        } finally {
+          executionContext.afterMethod();
+        }
+
+        return result;
+      }
+
+      executionContext.beforeMethod(constructorName, methodName, args);
+
+      if (executionContext.isTopLevel) {
+        if (!this.runtime) {
+          throw errors.runtimeNotProvided(constructorName);
+        }
+        executionContext.setProver(prover.bind(this.runtime.zkProgrammable));
+      }
+
+      // eslint-disable-next-line @typescript-eslint/init-declarations
+      let result: unknown;
+      try {
+        result = Reflect.apply(simulatedMethod, this, args);
+      } finally {
+        executionContext.afterMethod();
+      }
+
+      return result;
+    };
+  };
+}

package/src/module/decorator.ts

@@ -0,0 +1,21 @@
+import { injectable } from "tsyringe";
+import { StaticConfigurableModule, TypedClass } from "@proto-kit/common";
+
+import { RuntimeModule } from "../runtime/RuntimeModule.js";
+
+/**
+ * Marks the decorated class as a runtime module, while also
+ * making it injectable with our dependency injection solution.
+ */
+export function runtimeModule() {
+  return (
+    /**
+     * Check if the target class extends RuntimeModule, while
+     * also providing static config presets
+     */
+    target: StaticConfigurableModule<unknown> &
+      TypedClass<RuntimeModule<unknown>>
+  ) => {
+    injectable()(target);
+  };
+}

package/src/runtime/Runtime.ts

@@ -0,0 +1,304 @@
+// 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";
+import {
+  StringKeyOf,
+  ModuleContainer,
+  ModulesConfig,
+  ModulesRecord,
+  TypedClass,
+  ZkProgrammable,
+  PlainZkProgram,
+  WithZkProgrammable,
+  AreProofsEnabled,
+} from "@proto-kit/common";
+import {
+  fieldToString,
+  stringToField,
+  MethodPublicOutput,
+} from "@proto-kit/protocol";
+
+import {
+  combineMethodName,
+  isRuntimeMethod,
+  toWrappedMethod,
+  WrappedMethod,
+} from "../method/runtimeMethod.js";
+import { StateService } from "../state/InMemoryStateService.js";
+import { StateServiceProvider } from "../state/StateServiceProvider";
+
+import { RuntimeModule } from "./RuntimeModule.js";
+
+/**
+ * Record of modules accepted by the Runtime module container.
+ *
+ * We have to use TypedClass since RuntimeModule
+ * is an abstract class
+ */
+export type RuntimeModulesRecord = ModulesRecord<
+  TypedClass<RuntimeModule<unknown>>
+>;
+
+const errors = {
+  methodNotFound: (methodKey: string) =>
+    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
+> 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 zkProgramFactory(): PlainZkProgram<undefined, MethodPublicOutput> {
+    type Methods = Record<
+      string,
+      {
+        privateInputs: any;
+        method: WrappedMethod;
+      }
+    >;
+    // 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
+    const runtime: Runtime<Modules> = this.runtime;
+
+    const runtimeMethods = runtime.runtimeModuleNames.reduce<Methods>(
+      (allMethods, runtimeModuleName) => {
+        runtime.isValidModuleName(
+          runtime.definition.modules,
+          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
+        >;
+
+        const modulePrototypeMethods =
+          Object.getOwnPropertyNames(modulePrototype);
+
+        const moduleMethods = modulePrototypeMethods.reduce<Methods>(
+          (allModuleMethods, methodName) => {
+            if (isRuntimeMethod(runtimeModule, methodName)) {
+              const combinedMethodName = combineMethodName(
+                runtimeModuleName,
+                methodName
+              );
+              const method = modulePrototype[methodName];
+              const wrappedMethod = Reflect.apply(
+                toWrappedMethod,
+                runtimeModule,
+                [methodName, method]
+              );
+
+              // eslint-disable-next-line no-warning-comments
+              // TODO: find out how to import the Tuple type
+
+              const privateInputs = Reflect.getMetadata(
+                "design:paramtypes",
+                runtimeModule,
+                methodName
+              );
+
+              return {
+                ...allModuleMethods,
+
+                [combinedMethodName]: {
+                  privateInputs,
+                  method: wrappedMethod,
+                },
+              };
+            }
+
+            return allModuleMethods;
+          },
+          {}
+        );
+
+        return {
+          ...allMethods,
+          ...moduleMethods,
+        };
+      },
+      {}
+    );
+
+    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,
+    };
+  }
+}
+
+/**
+ * Wrapper for an application specific runtime, which helps orchestrate
+ * runtime modules into an interoperable runtime.
+ */
+@injectable()
+export class Runtime<Modules extends RuntimeModulesRecord>
+  extends ModuleContainer<Modules>
+  implements WithZkProgrammable<undefined, MethodPublicOutput>
+{
+  public static from<Modules extends RuntimeModulesRecord>(
+    definition: RuntimeDefinition<Modules>
+  ) {
+    return new Runtime(definition);
+  }
+
+  // runtime modules composed into a ZkProgram
+  public program?: ReturnType<typeof Experimental.ZkProgram>;
+
+  public definition: RuntimeDefinition<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>) {
+    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");
+  }
+
+  public get stateService(): StateService {
+    return this.stateServiceProviderInstance.stateService;
+  }
+
+  public get stateServiceProvider(): StateServiceProvider {
+    return this.stateServiceProviderInstance;
+  }
+
+  /**
+   * @returns The dependency injection container of this runtime
+   */
+  public get dependencyContainer(): DependencyContainer {
+    return this.container;
+  }
+
+  /**
+   * @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);
+
+    this.isValidModuleName(this.definition.modules, 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
+    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()
+    );
+  }
+
+  /**
+   * 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
+   */
+  public decorateModule(
+    moduleName: StringKeyOf<Modules>,
+    containedModule: InstanceType<Modules[StringKeyOf<Modules>]>
+  ) {
+    containedModule.name = moduleName;
+    containedModule.runtime = this;
+
+    super.decorateModule(moduleName, containedModule);
+  }
+
+  /**
+   * @returns A list of names of all the registered module names
+   */
+  public get runtimeModuleNames() {
+    return Object.keys(this.definition.modules);
+  }
+}

package/src/runtime/RuntimeModule.ts

@@ -0,0 +1,68 @@
+import { ConfigurableModule, Presets } from "@proto-kit/common";
+import { container, injectable } from "tsyringe";
+import { NetworkState, RuntimeTransaction } from "@proto-kit/protocol";
+
+import {
+  RuntimeMethodExecutionContext,
+  RuntimeMethodExecutionData,
+} from "../method/RuntimeMethodExecutionContext";
+import { StateService } from "../state/InMemoryStateService";
+
+import type {
+  Runtime,
+  RuntimeDefinition,
+  RuntimeModulesRecord,
+} from "./Runtime";
+
+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">;
+
+  get stateService(): StateService;
+}
+
+/**
+ * Base class for runtime modules providing the necessary utilities.
+ */
+@injectable()
+export class RuntimeModule<Config> extends ConfigurableModule<Config> {
+  public static presets: Presets<unknown> = {};
+
+  /**
+   * This property exists only to typecheck that the RuntimeModule
+   * was extended correctly in e.g. a decorator. We need at least
+   * one non-optional property in this class to make the typechecking work.
+   */
+  public isRuntimeModule = true;
+
+  public name?: string;
+
+  public runtime?: Runtime<RuntimeModulesRecord>;
+
+  private getInputs(): RuntimeMethodExecutionData {
+    const { input } = container.resolve<RuntimeMethodExecutionContext>(
+      RuntimeMethodExecutionContext
+    );
+
+    if (input === undefined) {
+      throw errors.inputDataNotSet();
+    }
+    return input;
+  }
+
+  public get transaction(): RuntimeTransaction {
+    return this.getInputs().transaction;
+  }
+
+  public get network(): NetworkState {
+    return this.getInputs().networkState;
+  }
+}

package/src/state/InMemoryStateService.ts

@@ -0,0 +1,29 @@
+import { Field } from "snarkyjs";
+
+export interface StateService {
+  get: (key: Field) => Field[] | undefined;
+  set: (key: Field, value: Field[] | undefined) => void;
+}
+
+/**
+ * Naive implementation of a StateService for testing purposes
+ */
+export class InMemoryStateService implements StateService {
+  public values: Record<string, Field[] | undefined> = {};
+
+  public get(key: Field): Field[] | undefined {
+    return this.values[key.toString()];
+  }
+
+  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()];
+    } else {
+      this.values[key.toString()] = value;
+    }
+  }
+}