@proto-kit/common 0.1.1-develop.211 → 0.1.1-develop.2137

package/src/utils.ts

@@ -1,4 +1,13 @@
-import { Field, FlexibleProvable, FlexibleProvablePure } from "snarkyjs";
+import {
+  Field,
+  FlexibleProvablePure,
+  Poseidon,
+  DynamicProof,
+  Proof,
+} from "o1js";
+import _ from "lodash";
+
+import { TypedClass } from "./types";
 
 export function requireTrue(
   condition: boolean,
@@ -11,14 +20,106 @@ export function requireTrue(
   }
 }
 
-export function range(startOrEnd: number, end: number | undefined): number[] {
+/**
+ * Utility function to split an array of type T into a record <K, T[]> based on a
+ * function T => K that determines the key of each record
+ */
+export function splitArray<T, K extends string | number>(
+  arr: T[],
+  split: (t: T) => K
+): Record<K, T[] | undefined> {
+  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+  const record = {} as { [Key in K]: T[] };
+  arr.forEach((element) => {
+    const k = split(element);
+    if (record[k] !== undefined) {
+      record[k].push(element);
+    } else {
+      record[k] = [element];
+    }
+  });
+  return record;
+}
+
+export function range(
+  startOrEnd: number,
+  endOrNothing?: number | undefined
+): number[] {
+  let end = endOrNothing;
+  let start = startOrEnd;
   if (end === undefined) {
     end = startOrEnd;
-    startOrEnd = 0;
+    start = 0;
   }
-  return Array.from({ length: end - startOrEnd }, (ignored, index) => index);
+  return Array.from({ length: end - start }, (ignored, index) => index + start);
+}
+
+export function reduceSequential<T, U>(
+  array: T[],
+  callbackfn: (
+    previousValue: U,
+    currentValue: T,
+    currentIndex: number,
+    array: T[]
+  ) => Promise<U>,
+  initialValue: U
+): Promise<U> {
+  return array.reduce<Promise<U>>(
+    async (previousPromise, current, index, arr) => {
+      const previous = await previousPromise;
+      return await callbackfn(previous, current, index, arr);
+    },
+    Promise.resolve(initialValue)
+  );
+}
+
+export function yieldSequential<Source, State, Target>(
+  array: Source[],
+  callbackfn: (
+    previousValue: State,
+    currentValue: Source,
+    currentIndex: number,
+    array: Source[]
+  ) => Promise<[State, Target]>,
+  initialValue: State
+): Promise<[State, Target[]]> {
+  return reduceSequential<Source, [State, Target[]]>(
+    array,
+    async ([state, collectedTargets], curr, index, arr) => {
+      const [newState, addition] = await callbackfn(state, curr, index, arr);
+      // The reason we wrap this in an array here is for a special case where Target is a tuple
+      // or array itself. In this case, js interprets by flattening the Value in the array
+      // (which it does when a function (like concat) uses a spread operator and the
+      // input is an array)
+      return [newState, collectedTargets.concat([addition])];
+    },
+    [initialValue, []]
+  );
 }
 
+export function mapSequential<T, R>(
+  array: T[],
+  f: (element: T, index: number, array: T[]) => Promise<R>
+): Promise<R[]> {
+  return array.reduce<Promise<R[]>>(async (r, element, index, a) => {
+    const ret = await r;
+    const next = await f(element, index, a);
+    ret.push(next);
+    return ret;
+  }, Promise.resolve([]));
+}
+
+export function unzip<A, B>(array: [A, B][]): [A[], B[]] {
+  const as = array.map(([a]) => a);
+  const bs = array.map(([, b]) => b);
+  return [as, bs];
+}
+
+export function assertSizeOneOrTwo<T>(arr: T[]): asserts arr is [T] | [T, T] {
+  if (!(arr.length === 1 || arr.length === 2)) {
+    throw new Error("Given array not size 1 or 2");
+  }
+}
 
 /**
  * Computes a dummy value for the given value type.
@@ -34,4 +135,179 @@ export function dummyValue<Value>(
 
   // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
   return valueType.fromFields(fields) as Value;
-}
+}
+
+export function noop(): void {}
+
+export interface ToFieldable {
+  toFields: () => Field[];
+}
+
+export interface ToFieldableStatic {
+  toFields: (value: unknown) => Field[];
+}
+
+export interface ToJSONableStatic {
+  toJSON: (value: unknown) => any;
+}
+
+// export interface ProofTypes {
+//   publicOutputType?: ToFieldableStatic;
+//   publicInputType?: ToFieldableStatic;
+// }
+
+export type ProofTypes =
+  | typeof Proof<unknown, unknown>
+  | typeof DynamicProof<unknown, unknown>;
+
+export async function sleep(ms: number) {
+  await new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}
+
+export function filterNonNull<Type>(value: Type | null): value is Type {
+  return value !== null;
+}
+
+export function filterNonUndefined<Type>(
+  value: Type | undefined
+): value is Type {
+  return value !== undefined;
+}
+
+const encoder = new TextEncoder();
+
+// Copied from o1js binable.ts:317
+export function prefixToField(prefix: string): Field {
+  const fieldSize = Field.sizeInBytes;
+  if (prefix.length >= fieldSize) throw Error("prefix too long");
+  const stringBytes = [...encoder.encode(prefix)];
+  return Field.fromBytes(
+    stringBytes.concat(Array(fieldSize - stringBytes.length).fill(0))
+  );
+}
+
+export function hashWithPrefix(prefix: string, input: Field[]) {
+  const salt = Poseidon.update(
+    [Field(0), Field(0), Field(0)],
+    [prefixToField(prefix)]
+  );
+  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+  return Poseidon.update(salt as [Field, Field, Field], input)[0];
+}
+
+// end copy
+
+export function expectDefined<T>(value: T | undefined): asserts value is T {
+  expect(value).toBeDefined();
+}
+
+type NonMethodKeys<Type> = {
+  [Key in keyof Type]: Type[Key] extends Function ? never : Key;
+}[keyof Type];
+export type NonMethods<Type> = Pick<Type, NonMethodKeys<Type>>;
+
+export const MAX_FIELD = Field(Field.ORDER - 1n);
+
+/**
+ * Returns a boolean indicating whether a given class is a subclass of another class,
+ * indicated by the name parameter.
+ */
+// TODO Change to class reference based comparisons
+export function isSubtypeOfName(
+  clas: TypedClass<unknown>,
+  name: string
+): boolean {
+  if (clas === undefined || clas === null) {
+    return false;
+  }
+
+  if (clas.name === name) {
+    return true;
+  }
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+  return isSubtypeOfName(Object.getPrototypeOf(clas), name);
+}
+
+// TODO Eventually, replace this by a schema validation library
+export function safeParseJson<T>(json: string) {
+  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+  return JSON.parse(json) as T;
+}
+
+export type Nullable<T> = {
+  [Key in keyof T]: T[Key] | undefined;
+};
+
+export function isFull<T>(t: Nullable<T>): t is T {
+  return Object.values(t).findIndex((v) => v === undefined) === -1;
+}
+
+// TODO Restructure utils into separate package and multiple files
+
+export function padArray<T>(
+  array: T[],
+  batchSize: number,
+  generator: (index: number) => T
+): T[] {
+  const slice = array.slice();
+  const dummies = range(0, batchSize - (array.length % batchSize)).map((i) =>
+    generator(i + array.length)
+  );
+  slice.push(...dummies);
+  return slice;
+}
+
+export function batch<T>(
+  arr: T[],
+  batchSize: number,
+  dummy: (index: number) => T
+): T[][] {
+  const padded = padArray(arr, batchSize, dummy);
+
+  const partitioned = _.groupBy(
+    padded.map((v, i) => [v, i] as const),
+    ([v, i]) => Math.floor(i / batchSize)
+  );
+
+  const numBatches = Math.ceil(arr.length / batchSize);
+
+  return range(0, numBatches).map((i) => partitioned[i].map((x) => x[0]));
+}
+
+export type Reference<T> = {
+  set value(value: T);
+  get value(): T;
+};
+
+class ReferenceObject<T> {
+  public constructor(private internalValue: T) {}
+
+  get value() {
+    return this.internalValue;
+  }
+
+  set value(t: T) {
+    this.internalValue = t;
+  }
+}
+
+export function createReference<T>(initial: T): Reference<T> {
+  return new ReferenceObject(initial);
+}
+
+export namespace BigIntMath {
+  export function max(...args: bigint[]) {
+    return args.reduce((m, e) => (e > m ? e : m));
+  }
+}
+
+export function assertDefined<T>(
+  t: T | undefined,
+  msg?: string
+): asserts t is T {
+  if (t === undefined) {
+    throw new Error(msg ?? "Value is undefined");
+  }
+}

package/src/zkProgrammable/ProvableMethodExecutionContext.ts

@@ -1,7 +1,9 @@
-import type { FlexibleProvable, Proof } from "snarkyjs";
+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"),
 
@@ -16,13 +18,12 @@ export class ProvableMethodExecutionResult {
 
   public methodName?: string;
 
-  public args?: FlexibleProvable<unknown>[];
+  public args?: ArgumentTypes;
 
   public prover?: () => Promise<Proof<unknown, unknown>>;
 
   public async prove<
-    // eslint-disable-next-line etc/no-misused-generics
-    ProofType extends Proof<unknown, unknown>
+    ProofType extends Proof<unknown, unknown>,
   >(): Promise<ProofType> {
     if (!this.prover) {
       // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
@@ -52,7 +53,6 @@ export class ProvableMethodExecutionContext {
   public result: ProvableMethodExecutionResult =
     new ProvableMethodExecutionResult();
 
-  // eslint-disable-next-line no-warning-comments,max-len
   // 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,
@@ -74,7 +74,7 @@ export class ProvableMethodExecutionContext {
   public beforeMethod(
     moduleName: string,
     methodName: string,
-    args: FlexibleProvable<unknown>[]
+    args: ArgumentTypes
   ) {
     if (this.isFinished) {
       this.clear();

package/src/zkProgrammable/ZkProgrammable.ts

@@ -1,15 +1,30 @@
-import { Experimental, FlexibleProvablePure, Proof } from "snarkyjs";
+import {
+  ZkProgram,
+  FlexibleProvablePure,
+  Proof,
+  Field,
+  Provable,
+  Cache as O1Cache,
+} from "o1js";
 import { Memoize } from "typescript-memoize";
 
-import { mockProof } from "./provableMethod";
+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 = {
-  appChainNotSet: (name: string) =>
-    new Error(`Appchain was not injected for: ${name}`),
+  areProofsEnabledNotSet: (name: string) =>
+    new Error(`AreProofsEnabled was not injected for: ${name}`),
 };
 
 export interface CompileArtifact {
-  verificationKey: string;
+  verificationKey: {
+    data: string;
+    hash: Field;
+  };
 }
 
 export interface AreProofsEnabled {
@@ -22,29 +37,45 @@ export interface Verify<PublicInput, PublicOutput> {
 }
 
 export interface Compile {
-  (): Promise<CompileArtifact>;
+  (options?: {
+    cache?: O1Cache;
+    forceRecompile?: boolean;
+    proofsEnabled?: boolean;
+    withRuntimeTables?: boolean;
+    numChunks?: number;
+    lazyMode?: boolean;
+  }): Promise<CompileArtifact>;
 }
 
-export interface PlainZkProgram<PublicInput = undefined, PublicOutput = void> {
+export interface PlainZkProgram<
+  PublicInput = undefined,
+  PublicOutput = undefined,
+> {
+  name: string;
   compile: Compile;
   verify: Verify<PublicInput, PublicOutput>;
   Proof: ReturnType<
-    typeof Experimental.ZkProgram.Proof<
+    typeof ZkProgram.Proof<
       FlexibleProvablePure<PublicInput>,
       FlexibleProvablePure<PublicOutput>
     >
   >;
   methods: Record<
     string,
-    | ((
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        ...args: any
-      ) => Promise<Proof<PublicInput, PublicOutput>>)
+    | ((...args: any) => Promise<{
+        proof: Proof<PublicInput, PublicOutput>;
+        auxiliaryOutput: any;
+      }>)
     | ((
         publicInput: PublicInput,
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
         ...args: any
-      ) => Promise<Proof<PublicInput, PublicOutput>>)
+      ) => Promise<{
+        proof: Proof<PublicInput, PublicOutput>;
+        auxiliaryOutput: any;
+      }>)
+  >;
+  analyzeMethods: () => Promise<
+    Record<string, Awaited<ReturnType<typeof Provable.constraintSystem>>>
   >;
 }
 
@@ -60,60 +91,84 @@ export function verifyToMockable<PublicInput, PublicOutput>(
         verified = await verify(proof);
       } catch (error: unknown) {
         // silently fail verification
-        console.error(error);
+        log.error(error);
         verified = false;
       }
 
       return verified;
     }
 
-    return proof.proof === mockProof;
+    return proof.proof === MOCK_PROOF;
   };
 }
 
-export const mockVerificationKey = "mock-verification-key";
+export const MOCK_VERIFICATION_KEY = dummyVerificationKey();
+
 export function compileToMockable(
   compile: Compile,
   { areProofsEnabled }: AreProofsEnabled
-): () => Promise<CompileArtifact> {
-  return async () => {
+): Compile {
+  return async (...args) => {
     if (areProofsEnabled) {
-      return await compile();
+      return await compile(...args);
     }
 
     return {
-      verificationKey: mockVerificationKey,
+      verificationKey: MOCK_VERIFICATION_KEY,
     };
   };
 }
 
 export abstract class ZkProgrammable<
   PublicInput = undefined,
-  PublicOutput = void
+  PublicOutput = void,
 > {
-  public abstract get appChain(): AreProofsEnabled | undefined;
+  public abstract get areProofsEnabled(): AreProofsEnabled | undefined;
 
-  public abstract zkProgramFactory(): PlainZkProgram<PublicInput, PublicOutput>;
+  public abstract zkProgramFactory(): PlainZkProgram<
+    PublicInput,
+    PublicOutput
+  >[];
 
-  @Memoize()
-  public get zkProgram(): PlainZkProgram<PublicInput, PublicOutput> {
-    const zkProgram = this.zkProgramFactory();
+  private zkProgramSingleton?: PlainZkProgram<PublicInput, PublicOutput>[];
 
-    if (!this.appChain) {
-      throw errors.appChainNotSet(this.constructor.name);
+  @Memoize()
+  public get zkProgram(): PlainZkProgram<PublicInput, PublicOutput>[] {
+    if (this.zkProgramSingleton === undefined) {
+      this.zkProgramSingleton = this.zkProgramFactory();
     }
 
-    return {
-      ...zkProgram,
-      verify: verifyToMockable(zkProgram.verify, this.appChain),
-      compile: compileToMockable(zkProgram.compile, this.appChain),
-    };
+    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
+  PublicOutput = void,
 > {
   zkProgrammable: ZkProgrammable<PublicInput, PublicOutput>;
 }

package/src/zkProgrammable/provableMethod.ts

@@ -1,39 +1,68 @@
-import { FlexibleProvable } from "snarkyjs";
+import { Proof, DynamicProof } from "o1js";
 import { container } from "tsyringe";
 
 import { ProvableMethodExecutionContext } from "./ProvableMethodExecutionContext";
 import type { WithZkProgrammable, ZkProgrammable } from "./ZkProgrammable";
 
-// eslint-disable-next-line etc/prefer-interface
-export type DecoratedMethod = (...args: unknown[]) => unknown;
+// 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 const mockProof = "mock-proof";
+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: unknown[]
+  ...args: ArgumentTypes
 ) {
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
   return async function prover(this: ZkProgrammable<any, any>) {
-    const areProofsEnabled = this.appChain?.areProofsEnabled;
-    if (areProofsEnabled ?? false) {
-      const programProvableMethod = this.zkProgram.methods[methodName];
-      return await Reflect.apply(programProvableMethod, this, args);
+    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 (searching for method ${methodName})`
+      );
     }
 
-    // create a mock proof by simulating method execution in JS
-    const publicOutput = Reflect.apply(simulatedMethod, this, args);
+    if (areProofsEnabled) {
+      const programProvableMethod = zkProgram.methods[methodName];
+      const result = await Reflect.apply(programProvableMethod, this, args);
+      // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+      return result.proof as Proof<any, any>;
+    }
+
+    // create a mock proof by simulating method> execution in JS
+    const publicOutput = await Reflect.apply(simulatedMethod, this, args);
 
-    return new this.zkProgram.Proof({
-      proof: mockProof,
+    return new zkProgram.Proof({
+      proof: MOCK_PROOF,
 
-      // eslint-disable-next-line no-warning-comments
       // TODO: provide undefined if public input is not used
       publicInput: isFirstParameterPublicInput ? args[0] : undefined,
-      publicOutput,
-
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+      publicOutput: (() => {
+        // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
+        if (!publicOutput || typeof publicOutput !== "object") {
+          return undefined;
+        }
+        if ("publicOutput" in publicOutput) {
+          // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+          return (publicOutput as any).publicOutput;
+        }
+        return publicOutput;
+      })(),
       /**
        * We set this to the max possible number, to avoid having
        * to manually count in-circuit proof verifications
@@ -59,19 +88,18 @@ export function provableMethod(
   )
 ) {
   return <
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    Target extends WithZkProgrammable<any, any> | ZkProgrammable<any, any>
+    Target extends WithZkProgrammable<any, any> | ZkProgrammable<any, any>,
   >(
     target: Target,
     methodName: string,
-    descriptor: PropertyDescriptor
+    descriptor: TypedPropertyDescriptor<(...args: any[]) => Promise<any> | any>
   ) => {
     // 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: ZkProgrammable<unknown, unknown>,
-      ...args: FlexibleProvable<unknown>[]
+      ...args: ArgumentTypes
     ) {
       const prover = toProver(
         methodName,
@@ -95,7 +123,6 @@ export function provableMethod(
        * or not, execute its simulated (Javascript) version and
        * return the result.
        */
-      // eslint-disable-next-line @typescript-eslint/init-declarations
       let result: unknown;
       try {
         result = Reflect.apply(simulatedMethod, this, args);