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

package/src/method/runtimeMethod.ts

@@ -1,32 +1,42 @@
-import { FlexibleProvable } from "snarkyjs";
+import { Bool, Field, FlexibleProvablePure, Poseidon } from "o1js";
 import { container } from "tsyringe";
 import {
-  StateTransition,
-  DefaultProvableHashList,
   ProvableStateTransition,
   MethodPublicOutput,
+  RuntimeMethodExecutionContext,
+  StateTransitionReductionList,
+  DefaultProvableHashList,
 } from "@proto-kit/protocol";
-import { DecoratedMethod, toProver, ZkProgrammable } from "@proto-kit/common";
+import {
+  DecoratedMethod,
+  toProver,
+  ZkProgrammable,
+  ArgumentTypes,
+} from "@proto-kit/common";
 
 import type { RuntimeModule } from "../runtime/RuntimeModule.js";
 
-import { RuntimeMethodExecutionContext } from "./RuntimeMethodExecutionContext.js";
+import {
+  MethodParameterEncoder,
+  checkArgsProvable,
+} from "./MethodParameterEncoder";
 
 const errors = {
   runtimeNotProvided: (name: string) =>
     new Error(`Runtime was not provided for module: ${name}`),
 
-  methodInputsNotProvided: () =>
+  runtimeNameNotSet: () => new Error("Runtime name was not set"),
+
+  fieldNotConstant: (name: string) =>
     new Error(
-      "Method execution inputs not provided, provide them via context.inputs"
+      `In-circuit field ${name} not a constant, this is likely a framework bug`
     ),
 };
 
 export function toStateTransitionsHash(
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  stateTransitions: StateTransition<any>[]
+  stateTransitions: { toProvable: () => ProvableStateTransition }[]
 ) {
-  const stateTransitionsHashList = new DefaultProvableHashList(
+  const stateTransitionsHashList = new StateTransitionReductionList(
     ProvableStateTransition
   );
 
@@ -40,39 +50,102 @@ export function toStateTransitionsHash(
     .toField();
 }
 
-// eslint-disable-next-line etc/prefer-interface
-export type WrappedMethod = (...args: unknown[]) => MethodPublicOutput;
+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: unknown[]) => unknown
-) {
+  moduleMethod: (...args: ArgumentTypes) => Promise<any>,
+  options: {
+    invocationType: RuntimeMethodInvocationType;
+  }
+): AsyncWrappedMethod {
   const executionContext = container.resolve<RuntimeMethodExecutionContext>(
     RuntimeMethodExecutionContext
   );
 
-  const wrappedMethod: WrappedMethod = (...args): MethodPublicOutput => {
-    Reflect.apply(moduleMethod, this, args);
+  const wrappedMethod: AsyncWrappedMethod = async (
+    ...args
+  ): Promise<MethodPublicOutput> => {
+    await Reflect.apply(moduleMethod, this, args);
     const {
-      result: { stateTransitions, status },
-      input,
+      result: { stateTransitions, status, events },
     } = executionContext.current();
 
     const stateTransitionsHash = toStateTransitionsHash(stateTransitions);
+    const eventsHash = toEventsHash(events);
 
-    if (input === undefined) {
-      throw errors.methodInputsNotProvided();
+    const { name, parent: 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");
     }
 
-    const transactionHash = input.transaction.hash();
-    const networkStateHash = input.networkState.hash();
+    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,
     });
   };
 
@@ -93,6 +166,7 @@ export function combineMethodName(
 
 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,
@@ -111,12 +185,17 @@ export function isRuntimeMethod(
   );
 }
 
-export function runtimeMethod() {
+export type RuntimeMethodInvocationType = "SIGNATURE" | "INCOMING_MESSAGE";
+
+function runtimeMethodInternal(options: {
+  invocationType: RuntimeMethodInvocationType;
+}) {
   return (
     target: RuntimeModule<unknown>,
     methodName: string,
-    descriptor: PropertyDescriptor
+    descriptor: TypedPropertyDescriptor<(...args: any[]) => Promise<any>>
   ) => {
+    checkArgsProvable(target, methodName);
     const executionContext = container.resolve<RuntimeMethodExecutionContext>(
       RuntimeMethodExecutionContext
     );
@@ -135,14 +214,21 @@ export function runtimeMethod() {
 
     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 = function value(
+    descriptor.value = async function value(
       this: RuntimeModule<unknown>,
-      ...args: FlexibleProvable<unknown>[]
+      ...args: ArgumentTypes
     ) {
-      const constructorName = this.constructor.name;
+      const constructorName = this.name!;
 
       /**
        * If its a top level method call, wrap it into a wrapped method,
@@ -154,6 +240,7 @@ export function runtimeMethod() {
       const simulatedWrappedMethod = Reflect.apply(toWrappedMethod, this, [
         methodName,
         simulatedMethod,
+        options,
       ]);
 
       /**
@@ -161,7 +248,7 @@ export function runtimeMethod() {
        * 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(
@@ -170,7 +257,6 @@ export function runtimeMethod() {
           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
@@ -185,16 +271,15 @@ export function runtimeMethod() {
       executionContext.beforeMethod(constructorName, methodName, args);
 
       if (executionContext.isTopLevel) {
-        if (!this.runtime) {
+        if (!this.parent) {
           throw errors.runtimeNotProvided(constructorName);
         }
-        executionContext.setProver(prover.bind(this.runtime.zkProgrammable));
+        executionContext.setProver(prover.bind(this.parent.zkProgrammable));
       }
 
-      // eslint-disable-next-line @typescript-eslint/init-declarations
       let result: unknown;
       try {
-        result = Reflect.apply(simulatedMethod, this, args);
+        result = await Reflect.apply(simulatedMethod, this, args);
       } finally {
         executionContext.afterMethod();
       }
@@ -203,3 +288,15 @@ export function runtimeMethod() {
     };
   };
 }
+
+export function runtimeMessage() {
+  return runtimeMethodInternal({
+    invocationType: "INCOMING_MESSAGE",
+  });
+}
+
+export function runtimeMethod() {
+  return runtimeMethodInternal({
+    invocationType: "SIGNATURE",
+  });
+}

package/src/runtime/MethodIdResolver.ts

@@ -1,71 +1,108 @@
-import { singleFieldToString, stringToField } from "@proto-kit/protocol";
+import { filterNonUndefined } from "@proto-kit/common";
+import { stringToField, RuntimeMethodIdMapping } from "@proto-kit/protocol";
+import { Poseidon } from "o1js";
+import { inject, injectable } from "tsyringe";
 
-import type { Runtime, RuntimeModulesRecord } from "./Runtime";
+import {
+  RuntimeMethodInvocationType,
+  runtimeMethodTypeMetadataKey,
+} from "../method/runtimeMethod";
 
-const offset = 128n;
-// eslint-disable-next-line @typescript-eslint/no-magic-numbers
-const modulus = 2n ** (offset - 1n);
+import type { Runtime, RuntimeModulesRecord } from "./Runtime";
 
 /**
- * How do we encode MethodIds
- * A MethodId is defined as the following in little-endian
- * [0
- *   ...hash(stringToField(moduleName))[0..128],
- *   ...hash(stringToField(methodName))[0..128]
- * ]
+ * Please see `getMethodId` to learn more about
+ * methodId encoding
  */
+@injectable()
 export class MethodIdResolver {
-  private readonly dictionary: { [key: string]: string } = {};
+  private readonly dictionary: {
+    [key: string]: { moduleName: string; methodName: string };
+  } = {};
 
   public constructor(
-    private readonly runtime: Runtime<RuntimeModulesRecord>,
-    private readonly modules: RuntimeModulesRecord
+    @inject("Runtime") private readonly runtime: Runtime<RuntimeModulesRecord>
   ) {
-    this.dictionary = runtime.runtimeModuleNames.reduce<Record<string, string>>(
-      (dict, moduleName) => {
-        this.runtime.assertIsValidModuleName(modules, moduleName);
-
-        dict[(stringToField(moduleName).toBigInt() % modulus).toString()] =
-          moduleName;
-
-        runtime.resolve(moduleName).runtimeMethodNames.forEach((methodName) => {
-          dict[(stringToField(methodName).toBigInt() % modulus).toString()] =
-            methodName;
-        });
-        return dict;
-      },
-      {}
-    );
+    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;
+    }, {});
   }
 
-  public getMethodNameFromId(methodId: bigint): [string, string] | undefined {
-    const moduleNameHash = singleFieldToString(methodId >> offset);
-    const methodNameHash = singleFieldToString(methodId % modulus);
+  /**
+   * 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 moduleName: string | undefined = this.dictionary[moduleNameHash];
+    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;
 
-    // eslint-disable-next-line no-warning-comments
-    // TODO Replace by throwing exception?
-    if (moduleName === undefined) {
-      return undefined;
-    }
-    this.runtime.assertIsValidModuleName(this.modules, moduleName);
+        if (type !== undefined) {
+          return {
+            name: `${moduleName}.${method}`,
+            methodId: methodIdResolver.getMethodId(moduleName, method),
+            type,
+          } as const;
+        }
 
-    const methodName: string | undefined = this.dictionary[methodNameHash];
+        return undefined;
+      });
+    });
 
-    if (methodName === 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(this.modules, moduleName);
+    this.runtime.assertIsValidModuleName(moduleName);
 
-    return (
-      (stringToField(moduleName).toBigInt() % modulus << offset) +
-      (stringToField(methodName).toBigInt() % modulus)
-    );
+    return Poseidon.hash([
+      stringToField(moduleName),
+      stringToField(methodName),
+    ]).toBigInt();
   }
 }