@proto-kit/common 0.1.1-develop.0

package/src/zkProgrammable/ProvableMethodExecutionContext.ts

@@ -0,0 +1,122 @@
+import type { Proof } from "o1js";
+import { singleton } from "tsyringe";
+import uniqueId from "lodash/uniqueId";
+
+import type { ArgumentTypes } from "./provableMethod";
+
+const errors = {
+  moduleOrMethodNameNotSet: () => new Error("Module or method name not set"),
+
+  proverNotSet: (moduleName: string, methodName: string) =>
+    new Error(
+      `Prover not set for '${moduleName}.${methodName}', did you forget to decorate your method?`
+    ),
+};
+
+export class ProvableMethodExecutionResult {
+  public moduleName?: string;
+
+  public methodName?: string;
+
+  public args?: ArgumentTypes;
+
+  public prover?: () => Promise<Proof<unknown, unknown>>;
+
+  public async prove<
+    ProofType extends Proof<unknown, unknown>,
+  >(): Promise<ProofType> {
+    if (!this.prover) {
+      // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
+      if (!this.moduleName || !this.methodName) {
+        throw errors.moduleOrMethodNameNotSet();
+      }
+      throw errors.proverNotSet(this.moduleName, this.methodName);
+    }
+
+    // turn the prover result into the desired proof type
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    return (await this.prover()) as ProofType;
+  }
+}
+
+/**
+ * Execution context used to wrap runtime module methods,
+ * allowing them to post relevant information (such as execution status)
+ * into the context without any unnecessary 'prop drilling'.
+ */
+@singleton()
+export class ProvableMethodExecutionContext {
+  public id = uniqueId();
+
+  public methods: string[] = [];
+
+  public result: ProvableMethodExecutionResult =
+    new ProvableMethodExecutionResult();
+
+  // TODO See if we should make this class generic, bc I think we can persist the type
+  /**
+   * Adds a method prover to the current execution context,
+   * which can be collected and ran asynchronously at a later point in time.
+   *
+   * @param prove - Prover function to be ran later,
+   * when the method execution needs to be proven
+   */
+  public setProver(prover: () => Promise<Proof<unknown, unknown>>) {
+    this.result.prover = prover;
+  }
+
+  /**
+   * Adds a method to the method execution stack, reseting the execution context
+   * in a case a new top-level (non nested) method call is made.
+   *
+   * @param methodName - Name of the method being captured in the context
+   */
+  public beforeMethod(
+    moduleName: string,
+    methodName: string,
+    args: ArgumentTypes
+  ) {
+    if (this.isFinished) {
+      this.clear();
+      this.result.moduleName = moduleName;
+      this.result.methodName = methodName;
+      this.result.args = args;
+    }
+
+    this.methods.push(methodName);
+  }
+
+  /**
+   * Removes the latest method from the execution context stack,
+   * keeping track of the amount of 'unfinished' methods. Allowing
+   * for the context to distinguish between top-level and nested method calls.
+   */
+  public afterMethod() {
+    this.methods.pop();
+  }
+
+  public get isTopLevel() {
+    return this.methods.length === 1;
+  }
+
+  public get isFinished() {
+    return this.methods.length === 0;
+  }
+
+  /**
+   * @returns - Current execution context state
+   */
+  public current() {
+    return {
+      isFinished: this.isFinished,
+      result: this.result,
+    };
+  }
+
+  /**
+   * Manually clears/resets the execution context
+   */
+  public clear() {
+    this.result = new ProvableMethodExecutionResult();
+  }
+}

package/src/zkProgrammable/ZkProgrammable.ts

@@ -0,0 +1,151 @@
+import { ZkProgram, FlexibleProvablePure, Proof, Field, Provable } from "o1js";
+import { Memoize } from "typescript-memoize";
+
+import { log } from "../log";
+import { dummyVerificationKey } from "../dummyVerificationKey";
+import { reduceSequential } from "../utils";
+import type { CompileRegistry } from "../compiling/CompileRegistry";
+
+import { MOCK_PROOF } from "./provableMethod";
+
+const errors = {
+  areProofsEnabledNotSet: (name: string) =>
+    new Error(`AreProofsEnabled was not injected for: ${name}`),
+};
+
+export interface CompileArtifact {
+  verificationKey: {
+    data: string;
+    hash: Field;
+  };
+}
+
+export interface AreProofsEnabled {
+  areProofsEnabled: boolean;
+  setProofsEnabled: (areProofsEnabled: boolean) => void;
+}
+
+export interface Verify<PublicInput, PublicOutput> {
+  (proof: Proof<PublicInput, PublicOutput>): Promise<boolean>;
+}
+
+export interface Compile {
+  (): Promise<CompileArtifact>;
+}
+
+export interface PlainZkProgram<PublicInput = undefined, PublicOutput = void> {
+  name: string;
+  compile: Compile;
+  verify: Verify<PublicInput, PublicOutput>;
+  Proof: ReturnType<
+    typeof ZkProgram.Proof<
+      FlexibleProvablePure<PublicInput>,
+      FlexibleProvablePure<PublicOutput>
+    >
+  >;
+  methods: Record<
+    string,
+    | ((...args: any) => Promise<Proof<PublicInput, PublicOutput>>)
+    | ((
+        publicInput: PublicInput,
+        ...args: any
+      ) => Promise<Proof<PublicInput, PublicOutput>>)
+  >;
+  analyzeMethods: () => Promise<
+    Record<string, Awaited<ReturnType<typeof Provable.constraintSystem>>>
+  >;
+}
+
+export function verifyToMockable<PublicInput, PublicOutput>(
+  verify: Verify<PublicInput, PublicOutput>,
+  { areProofsEnabled }: AreProofsEnabled
+) {
+  return async (proof: Proof<PublicInput, PublicOutput>) => {
+    if (areProofsEnabled) {
+      let verified = false;
+
+      try {
+        verified = await verify(proof);
+      } catch (error: unknown) {
+        // silently fail verification
+        log.error(error);
+        verified = false;
+      }
+
+      return verified;
+    }
+
+    return proof.proof === MOCK_PROOF;
+  };
+}
+
+export const MOCK_VERIFICATION_KEY = dummyVerificationKey();
+
+export function compileToMockable(
+  compile: Compile,
+  { areProofsEnabled }: AreProofsEnabled
+): () => Promise<CompileArtifact> {
+  return async () => {
+    if (areProofsEnabled) {
+      return await compile();
+    }
+
+    return {
+      verificationKey: MOCK_VERIFICATION_KEY,
+    };
+  };
+}
+
+export abstract class ZkProgrammable<
+  PublicInput = undefined,
+  PublicOutput = void,
+> {
+  public abstract get areProofsEnabled(): AreProofsEnabled | undefined;
+
+  public abstract zkProgramFactory(): PlainZkProgram<
+    PublicInput,
+    PublicOutput
+  >[];
+
+  private zkProgramSingleton?: PlainZkProgram<PublicInput, PublicOutput>[];
+
+  @Memoize()
+  public get zkProgram(): PlainZkProgram<PublicInput, PublicOutput>[] {
+    if (this.zkProgramSingleton === undefined) {
+      this.zkProgramSingleton = this.zkProgramFactory();
+    }
+
+    return this.zkProgramSingleton.map((bucket) => {
+      if (!this.areProofsEnabled) {
+        throw errors.areProofsEnabledNotSet(this.constructor.name);
+      }
+      return {
+        ...bucket,
+        verify: verifyToMockable(bucket.verify, this.areProofsEnabled),
+        compile: compileToMockable(bucket.compile, this.areProofsEnabled),
+      };
+    });
+  }
+
+  public async compile(registry: CompileRegistry) {
+    return await reduceSequential(
+      this.zkProgram,
+      async (acc, program) => {
+        const result = await registry.compile(program);
+        return {
+          ...acc,
+          [program.name]: result,
+        };
+      },
+      // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+      {} as Record<string, CompileArtifact>
+    );
+  }
+}
+
+export interface WithZkProgrammable<
+  PublicInput = undefined,
+  PublicOutput = void,
+> {
+  zkProgrammable: ZkProgrammable<PublicInput, PublicOutput>;
+}

package/src/zkProgrammable/provableMethod.ts

@@ -0,0 +1,124 @@
+import { Proof, DynamicProof } from "o1js";
+import { container } from "tsyringe";
+
+import { ProvableMethodExecutionContext } from "./ProvableMethodExecutionContext";
+import type { WithZkProgrammable, ZkProgrammable } from "./ZkProgrammable";
+
+// Now, in o1js provable types are normal js objects, therefore any
+export type O1JSPrimitive = object | string | boolean | number;
+export type ArgumentTypes = (
+  | O1JSPrimitive
+  | Proof<unknown, unknown>
+  | DynamicProof<unknown, unknown>
+)[];
+
+export type DecoratedMethod = (...args: ArgumentTypes) => Promise<unknown>;
+
+export const MOCK_PROOF = "mock-proof";
+// (await Proof.dummy(Field(0), Field(0), 2)).proof as string;
+
+export function toProver(
+  methodName: string,
+  simulatedMethod: DecoratedMethod,
+  isFirstParameterPublicInput: boolean,
+  ...args: ArgumentTypes
+) {
+  return async function prover(this: ZkProgrammable<any, any>) {
+    const { areProofsEnabled } = this.areProofsEnabled!;
+
+    const zkProgram = this.zkProgram.find((prog) =>
+      Object.keys(prog.methods).includes(methodName)
+    );
+
+    if (zkProgram === undefined) {
+      throw new Error("Correct ZkProgram not found");
+    }
+
+    if (areProofsEnabled) {
+      const programProvableMethod = zkProgram.methods[methodName];
+      return await Reflect.apply(programProvableMethod, this, args);
+    }
+
+    // create a mock proof by simulating method> execution in JS
+    const publicOutput = await Reflect.apply(simulatedMethod, this, args);
+
+    return new zkProgram.Proof({
+      proof: MOCK_PROOF,
+
+      // TODO: provide undefined if public input is not used
+      publicInput: isFirstParameterPublicInput ? args[0] : undefined,
+      publicOutput,
+
+      /**
+       * We set this to the max possible number, to avoid having
+       * to manually count in-circuit proof verifications
+       */
+      maxProofsVerified: 2,
+    });
+  };
+}
+
+/**
+ * Decorates a provable method on a 'prover class', depending on
+ * if proofs are enabled or not, either runs the respective zkProgram prover,
+ * or simulates the method execution and issues a mock proof.
+ *
+ * @param isFirstParameterPublicInput
+ * @param executionContext
+ * @returns
+ */
+export function provableMethod(
+  isFirstParameterPublicInput = true,
+  executionContext: ProvableMethodExecutionContext = container.resolve<ProvableMethodExecutionContext>(
+    ProvableMethodExecutionContext
+  )
+) {
+  return <
+    Target extends WithZkProgrammable<any, any> | ZkProgrammable<any, any>,
+  >(
+    target: Target,
+    methodName: string,
+    descriptor: TypedPropertyDescriptor<(...args: any[]) => Promise<any> | any>
+  ) => {
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    const simulatedMethod = descriptor.value as DecoratedMethod;
+
+    descriptor.value = async function value(
+      this: ZkProgrammable<unknown, unknown>,
+      ...args: ArgumentTypes
+    ) {
+      const prover = toProver(
+        methodName,
+        simulatedMethod,
+        isFirstParameterPublicInput,
+        ...args
+      );
+
+      executionContext.beforeMethod(this.constructor.name, methodName, args);
+
+      /**
+       * Check if the method is called at the top level,
+       * if yes then create a prover.
+       */
+      if (executionContext.isTopLevel) {
+        executionContext.setProver(prover.bind(this));
+      }
+
+      /**
+       * Regardless of if the method is called from the top level
+       * or not, execute its simulated (Javascript) version and
+       * return the result.
+       */
+      let result: unknown;
+      try {
+        result = Reflect.apply(simulatedMethod, this, args);
+      } finally {
+        executionContext.afterMethod();
+      }
+
+      return result;
+    };
+
+    return descriptor;
+  };
+}

package/test/config/ContainerEvents.test.ts

@@ -0,0 +1,67 @@
+import "reflect-metadata";
+import { injectable, container as tsyringeContainer } from "tsyringe";
+
+import {
+  BaseModuleInstanceType,
+  ConfigurableModule,
+  EventEmitter,
+  EventEmittingComponent,
+  ModuleContainer,
+  ModulesRecord,
+} from "../../src";
+
+class TestContainer<
+  Modules extends ModulesRecord,
+> extends ModuleContainer<Modules> {}
+
+type TestEvents = {
+  test: [string];
+};
+
+@injectable()
+class TestModule
+  extends ConfigurableModule<{}>
+  implements BaseModuleInstanceType, EventEmittingComponent<TestEvents>
+{
+  events = new EventEmitter<TestEvents>();
+}
+
+type TestEvents2 = {
+  test2: [number];
+};
+
+class TestModule2
+  extends ConfigurableModule<{}>
+  implements BaseModuleInstanceType, EventEmittingComponent<TestEvents2>
+{
+  events = new EventEmitter<TestEvents2>();
+}
+
+describe("test event propagation", () => {
+  it("should propagate events up", () => {
+    expect.assertions(1);
+
+    const container = new TestContainer({
+      modules: {
+        test: TestModule,
+        test2: TestModule2,
+      },
+    });
+
+    container.configure({
+      test: {},
+      test2: {},
+    });
+
+    container.create(() => tsyringeContainer.createChildContainer());
+
+    const testvalue = "testString123";
+
+    container.events.on("test", (value: string) => {
+      expect(value).toStrictEqual(testvalue);
+    });
+
+    const module = container.resolve("test");
+    module.events.emit("test", testvalue);
+  });
+});

package/test/config/ModuleContainer.test.ts

@@ -0,0 +1,215 @@
+import "reflect-metadata";
+import { container as tsyringeContainer, inject, injectable } from "tsyringe";
+
+import {
+  ConfigurableModule,
+  NoConfig,
+} from "../../src/config/ConfigurableModule";
+import {
+  ModuleContainer,
+  ModulesRecord,
+} from "../../src/config/ModuleContainer";
+import { TypedClass } from "../../src/types";
+import { DependencyFactory, expectDefined } from "../../src";
+import { injectAlias } from "../../src/config/injectAlias";
+
+// module container will accept modules that extend this type
+class BaseTestModule<Config> extends ConfigurableModule<Config> {}
+
+type TestModulesRecord = ModulesRecord<TypedClass<BaseTestModule<unknown>>>;
+
+interface TestModuleConfig {
+  testConfigProperty: number;
+  testConfigProperty2?: number;
+  testConfigProperty3?: number;
+}
+
+@injectable()
+@injectAlias(["child-alias", "multi-alias"])
+class ChildModule extends BaseTestModule<NoConfig> {
+  public constructor(@inject("TestModule") public readonly testModule: any) {
+    super();
+  }
+
+  x() {
+    return "dependency factory works";
+  }
+}
+
+@injectAlias(["base-alias", "multi-alias"])
+class TestModule
+  extends BaseTestModule<TestModuleConfig>
+  implements DependencyFactory
+{
+  public dependencies() {
+    return {
+      dependencyModule1: {
+        useClass: ChildModule,
+      },
+    };
+  }
+}
+
+interface OtherTestModuleConfig {
+  otherTestConfigProperty: number;
+}
+
+class OtherTestModule extends BaseTestModule<OtherTestModuleConfig> {
+  public x() {
+    return "";
+  }
+}
+
+/**
+ * Showcases a wrongly typed/defined module as
+ * per the TestModuleContainer requirements
+ */
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+class WrongTestModule {}
+
+class TestModuleContainer<
+  Modules extends TestModulesRecord,
+> extends ModuleContainer<Modules> {
+  public get dependencyContainer() {
+    return this.container;
+  }
+}
+
+describe("moduleContainer", () => {
+  let container: TestModuleContainer<{
+    TestModule: typeof TestModule;
+    OtherTestModule: typeof OtherTestModule;
+  }>;
+  const testConfigProperty = 0;
+
+  beforeEach(() => {
+    container = new TestModuleContainer({
+      modules: {
+        TestModule,
+        OtherTestModule,
+        // this module would not be assignable to TestModuleContainer
+        // WrongTestModule,
+      },
+    });
+  });
+
+  it("should resolve dependency factory dependencies correctly", () => {
+    container.configure({
+      TestModule: {
+        testConfigProperty,
+      },
+
+      OtherTestModule: {
+        otherTestConfigProperty: testConfigProperty,
+      },
+    });
+
+    container.create(() => tsyringeContainer.createChildContainer());
+
+    // Unfortunately we still need this so that the dependencies are registered
+    container.resolve("TestModule");
+    const dm = container.resolve("DependencyModule1");
+
+    expect(dm.x()).toBe("dependency factory works");
+    expect(dm.testModule).toBeDefined();
+  });
+
+  it("should throw on resolution, if config was not provided", () => {
+    expect.assertions(1);
+
+    container.create(() => tsyringeContainer.createChildContainer());
+
+    expect(() => {
+      container.resolve("TestModule");
+    }).toThrow();
+  });
+
+  it("should resolve the registered module with the provided config", () => {
+    expect.assertions(1);
+
+    container.configure({
+      TestModule: {
+        testConfigProperty,
+      },
+
+      OtherTestModule: {
+        otherTestConfigProperty: testConfigProperty,
+      },
+    });
+    container.create(() => tsyringeContainer.createChildContainer());
+
+    const testModule = container.resolve("TestModule");
+
+    expect(testModule.config.testConfigProperty).toBe(testConfigProperty);
+
+    const dependency = container.resolve("DependencyModule1");
+    dependency.x();
+  });
+
+  it("should stack configurations correctly", () => {
+    container.configure({
+      TestModule: {
+        testConfigProperty: 1,
+      },
+      OtherTestModule: {
+        otherTestConfigProperty: 4,
+      },
+    });
+
+    container.configurePartial({
+      TestModule: {
+        testConfigProperty2: 2,
+      },
+    });
+
+    container.configurePartial({
+      TestModule: {
+        testConfigProperty: 3,
+      },
+    });
+
+    container.create(() => tsyringeContainer.createChildContainer());
+
+    const { config } = container.resolve("TestModule");
+
+    expect(config.testConfigProperty).toBe(3);
+    expect(config.testConfigProperty2).toBe(2);
+    expect(config.testConfigProperty3).toBe(undefined);
+  });
+
+  it("should resolve dependencies correctly via alias", () => {
+    container.configure({
+      TestModule: {
+        testConfigProperty,
+      },
+
+      OtherTestModule: {
+        otherTestConfigProperty: testConfigProperty,
+      },
+    });
+
+    container.create(() => tsyringeContainer.createChildContainer());
+
+    // Unfortunately we still need this so that the dependencies are registered
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    const m1 = container.resolve("base-alias" as any);
+    const m2 = container.resolve("TestModule");
+
+    expectDefined(m1);
+    // Check if its the same reference
+    expect(m1).toBe(m2);
+
+    const dm1 = container.resolve("child-alias" as any) as ChildModule;
+    const dm2 = container.resolve("DependencyModule1");
+
+    expect(dm1.x()).toBe("dependency factory works");
+    expect(dm1.testModule).toBeDefined();
+    expect(dm1).toBe(dm2);
+
+    const multi =
+      container.dependencyContainer.resolveAll<BaseTestModule<unknown>>(
+        "multi-alias"
+      );
+    expect(multi).toHaveLength(2);
+  });
+});

package/test/config/injectAlias.test.ts

@@ -0,0 +1,28 @@
+import "reflect-metadata";
+import { getInjectAliases, injectAlias } from "../../src/config/injectAlias";
+
+@injectAlias(["foo", "bar"])
+class TestClass {}
+
+@injectAlias(["ayy"])
+class TestClass2 extends TestClass {}
+
+describe("injectAlias metadata", () => {
+  it("set and retrieve", () => {
+    expect.assertions(2);
+
+    const aliases = getInjectAliases(TestClass);
+
+    expect(aliases).toHaveLength(2);
+    expect(aliases).toStrictEqual(["foo", "bar"]);
+  });
+
+  it("recursive", () => {
+    expect.assertions(2);
+
+    const aliases = getInjectAliases(TestClass2);
+
+    expect(aliases).toHaveLength(3);
+    expect(aliases).toStrictEqual(["ayy", "foo", "bar"]);
+  });
+});