@proto-kit/module 0.1.1-develop.0

package/src/method/runtimeMethod.ts

@@ -0,0 +1,308 @@
+import { Bool, Field, FlexibleProvablePure, Poseidon } from "o1js";
+import { container } from "tsyringe";
+import {
+  StateTransition,
+  ProvableStateTransition,
+  MethodPublicOutput,
+  RuntimeMethodExecutionContext,
+  StateTransitionReductionList,
+  DefaultProvableHashList,
+} from "@proto-kit/protocol";
+import {
+  DecoratedMethod,
+  toProver,
+  ZkProgrammable,
+  ArgumentTypes,
+} from "@proto-kit/common";
+
+import type { RuntimeModule } from "../runtime/RuntimeModule.js";
+
+import {
+  MethodParameterEncoder,
+  checkArgsProvable,
+} from "./MethodParameterEncoder";
+
+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"
+    ),
+
+  runtimeNameNotSet: () => new Error("Runtime name was not set"),
+
+  fieldNotConstant: (name: string) =>
+    new Error(
+      `In-circuit field ${name} not a constant, this is likely a framework bug`
+    ),
+};
+
+export function toStateTransitionsHash(
+  stateTransitions: StateTransition<any>[]
+) {
+  const stateTransitionsHashList = new StateTransitionReductionList(
+    ProvableStateTransition
+  );
+
+  return stateTransitions
+    .map((stateTransition) => stateTransition.toProvable())
+    .reduce(
+      (allStateTransitionsHashList, stateTransition) =>
+        allStateTransitionsHashList.push(stateTransition),
+      stateTransitionsHashList
+    )
+    .toField();
+}
+
+export function toEventsHash(
+  events: {
+    eventType: FlexibleProvablePure<any>;
+    event: any;
+    eventName: string;
+    condition: Bool;
+  }[]
+) {
+  return events.reduce((acc, event) => {
+    const hashList = new DefaultProvableHashList(event.eventType, acc);
+    hashList.pushIf(event.event, event.condition);
+    return hashList.commitment;
+  }, Field(0));
+}
+
+export type WrappedMethod = (...args: ArgumentTypes) => MethodPublicOutput;
+export type AsyncWrappedMethod = (
+  ...args: ArgumentTypes
+) => Promise<MethodPublicOutput>;
+
+export function toWrappedMethod(
+  this: RuntimeModule<unknown>,
+  methodName: string,
+  moduleMethod: (...args: ArgumentTypes) => Promise<any>,
+  options: {
+    invocationType: RuntimeMethodInvocationType;
+  }
+): AsyncWrappedMethod {
+  const executionContext = container.resolve<RuntimeMethodExecutionContext>(
+    RuntimeMethodExecutionContext
+  );
+
+  const wrappedMethod: AsyncWrappedMethod = async (
+    ...args
+  ): Promise<MethodPublicOutput> => {
+    await Reflect.apply(moduleMethod, this, args);
+    const {
+      result: { stateTransitions, status, events },
+    } = executionContext.current();
+
+    const stateTransitionsHash = toStateTransitionsHash(stateTransitions);
+    const eventsHash = toEventsHash(events);
+
+    const { name, runtime } = this;
+
+    if (name === undefined) {
+      throw errors.runtimeNameNotSet();
+    }
+    if (runtime === undefined) {
+      throw errors.runtimeNotProvided(name);
+    }
+
+    const { transaction, networkState } = executionContext.witnessInput();
+    const { methodIdResolver } = runtime;
+
+    // Assert that the given transaction has the correct methodId
+    const thisMethodId = Field(methodIdResolver.getMethodId(name, methodName));
+    if (!thisMethodId.isConstant()) {
+      throw errors.fieldNotConstant("methodId");
+    }
+
+    transaction.methodId.assertEquals(
+      thisMethodId,
+      "Runtimemethod called with wrong methodId on the transaction object"
+    );
+
+    /**
+     * Use the type info obtained previously to convert
+     * the args passed to fields
+     */
+    const { fields } = MethodParameterEncoder.fromMethod(
+      this,
+      methodName
+    ).encode(args);
+
+    // Assert that the argsHash that has been signed matches the given arguments
+    const argsHash = Poseidon.hash(fields);
+
+    transaction.argsHash.assertEquals(
+      argsHash,
+      "argsHash and therefore arguments of transaction and runtime call does not match"
+    );
+
+    const isMessage = Bool(options.invocationType === "INCOMING_MESSAGE");
+    transaction.assertTransactionType(Bool(isMessage));
+
+    const transactionHash = transaction.hash();
+    const networkStateHash = networkState.hash();
+
+    return new MethodPublicOutput({
+      stateTransitionsHash,
+      status,
+      transactionHash,
+      networkStateHash,
+      isMessage,
+      eventsHash,
+    });
+  };
+
+  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";
+export const runtimeMethodNamesMetadataKey = "proto-kit-runtime-methods";
+export const runtimeMethodTypeMetadataKey = "proto-kit-runtime-method-type";
+
+/**
+ * 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 type RuntimeMethodInvocationType = "SIGNATURE" | "INCOMING_MESSAGE";
+
+function runtimeMethodInternal(options: {
+  invocationType: RuntimeMethodInvocationType;
+}) {
+  return (
+    target: RuntimeModule<unknown>,
+    methodName: string,
+    descriptor: TypedPropertyDescriptor<(...args: any[]) => Promise<any>>
+  ) => {
+    checkArgsProvable(target, methodName);
+    const executionContext = container.resolve<RuntimeMethodExecutionContext>(
+      RuntimeMethodExecutionContext
+    );
+
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    let data: string[] | undefined = Reflect.getMetadata(
+      runtimeMethodNamesMetadataKey,
+      target
+    );
+    if (data !== undefined) {
+      data.push(methodName);
+    } else {
+      data = [methodName];
+    }
+    Reflect.defineMetadata(runtimeMethodNamesMetadataKey, data, target);
+
+    Reflect.defineMetadata(runtimeMethodMetadataKey, true, target, methodName);
+
+    Reflect.defineMetadata(
+      runtimeMethodTypeMetadataKey,
+      options.invocationType,
+      target,
+      methodName
+    );
+
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    const simulatedMethod = descriptor.value as DecoratedMethod;
+
+    descriptor.value = async function value(
+      this: RuntimeModule<unknown>,
+      ...args: ArgumentTypes
+    ) {
+      const constructorName = this.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,
+        options,
+      ]);
+
+      /**
+       * Before the prover runs, make sure it is operating on the correct
+       * RuntimeMethodExecutionContext state, meaning it enters and exits
+       * the context properly.
+       */
+
+      async function prover(this: ZkProgrammable<any, any>) {
+        executionContext.beforeMethod(constructorName, methodName, args);
+        const innerProver = toProver(
+          combineMethodName(constructorName, methodName),
+          simulatedWrappedMethod,
+          false,
+          ...args
+        ).bind(this);
+        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));
+      }
+
+      let result: unknown;
+      try {
+        result = await Reflect.apply(simulatedMethod, this, args);
+      } finally {
+        executionContext.afterMethod();
+      }
+
+      return result;
+    };
+  };
+}
+
+export function runtimeMessage() {
+  return runtimeMethodInternal({
+    invocationType: "INCOMING_MESSAGE",
+  });
+}
+
+export function runtimeMethod() {
+  return runtimeMethodInternal({
+    invocationType: "SIGNATURE",
+  });
+}

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/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();
+  }
+}