Merge pull request #263 from proto-kit/feature/worker-readiness

Feature/worker readiness

Author: rpanic

packages/common/src/compiling/AtomicCompileHelper.ts

@@ -0,0 +1,62 @@
+import {
+  AreProofsEnabled,
+  CompileArtifact,
+  MOCK_VERIFICATION_KEY,
+} from "../zkProgrammable/ZkProgrammable";
+import { isSubtypeOfName } from "../utils";
+import { TypedClass } from "../types";
+import { log } from "../log";
+
+export type ArtifactRecord = Record<string, CompileArtifact>;
+
+export type CompileTarget = {
+  name: string;
+  compile: () => Promise<CompileArtifact>;
+};
+
+export class AtomicCompileHelper {
+  public constructor(private readonly areProofsEnabled: AreProofsEnabled) {}
+
+  private compilationPromises: {
+    [key: string]: Promise<CompileArtifact>;
+  } = {};
+
+  public async compileContract(
+    contract: CompileTarget,
+    overrideProofsEnabled?: boolean
+  ): Promise<CompileArtifact> {
+    let newPromise = false;
+    const { name } = contract;
+    if (this.compilationPromises[name] === undefined) {
+      const proofsEnabled =
+        overrideProofsEnabled ?? this.areProofsEnabled.areProofsEnabled;
+
+      // We only care about proofs enabled here if it's a contract, because
+      // in all other cases, ZkProgrammable already handles this switch, and we
+      // want to preserve the artifact layout (which might be more than one
+      // entry for ZkProgrammables)
+      if (
+        proofsEnabled ||
+        !isSubtypeOfName(
+          // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+          contract as unknown as TypedClass<any>,
+          "SmartContract"
+        )
+      ) {
+        log.time(`Compiling ${name}`);
+        this.compilationPromises[name] = contract.compile();
+        newPromise = true;
+      } else {
+        log.trace(`Compiling ${name} - mock`);
+        this.compilationPromises[name] = Promise.resolve({
+          verificationKey: MOCK_VERIFICATION_KEY,
+        });
+      }
+    }
+    const result = await this.compilationPromises[name];
+    if (newPromise) {
+      log.timeEnd.info(`Compiling ${name}`);
+    }
+    return result;
+  }
+}

packages/common/src/compiling/CompilableModule.ts

@@ -0,0 +1,6 @@
+import type { CompileRegistry } from "./CompileRegistry";
+import type { ArtifactRecord } from "./AtomicCompileHelper";
+
+export interface CompilableModule {
+  compile(registry: CompileRegistry): Promise<ArtifactRecord | void>;
+}

packages/common/src/compiling/CompileRegistry.ts

@@ -0,0 +1,61 @@
+import { inject, injectable, singleton } from "tsyringe";
+
+import { AreProofsEnabled } from "../zkProgrammable/ZkProgrammable";
+
+import {
+  ArtifactRecord,
+  AtomicCompileHelper,
+  CompileTarget,
+} from "./AtomicCompileHelper";
+
+/**
+ * The CompileRegistry compiles "compilable modules"
+ * (i.e. zkprograms, contracts or contractmodules)
+ * while making sure they don't get compiled twice in the same process in parallel.
+ */
+@injectable()
+@singleton()
+export class CompileRegistry {
+  public constructor(
+    @inject("AreProofsEnabled")
+    private readonly areProofsEnabled: AreProofsEnabled
+  ) {
+    this.compiler = new AtomicCompileHelper(this.areProofsEnabled);
+  }
+
+  private compiler: AtomicCompileHelper;
+
+  private artifacts: ArtifactRecord = {};
+
+  // TODO Add possibility to force recompilation for non-sideloaded dependencies
+
+  public async compile(target: CompileTarget) {
+    if (this.artifacts[target.name] === undefined) {
+      const artifact = await this.compiler.compileContract(target);
+      this.artifacts[target.name] = artifact;
+      return artifact;
+    }
+    return this.artifacts[target.name];
+  }
+
+  public getArtifact(name: string) {
+    if (this.artifacts[name] === undefined) {
+      throw new Error(
+        `Artifact for ${name} not available, did you compile it via the CompileRegistry?`
+      );
+    }
+
+    return this.artifacts[name];
+  }
+
+  public addArtifactsRaw(artifacts: ArtifactRecord) {
+    this.artifacts = {
+      ...this.artifacts,
+      ...artifacts,
+    };
+  }
+
+  public getAllArtifacts() {
+    return this.artifacts;
+  }
+}

packages/common/src/compiling/services/ChildVerificationKeyService.ts

@@ -0,0 +1,26 @@
+import { injectable, Lifecycle, scoped } from "tsyringe";
+
+import { CompileRegistry } from "../CompileRegistry";
+
+@injectable()
+@scoped(Lifecycle.ContainerScoped)
+export class ChildVerificationKeyService {
+  private compileRegistry?: CompileRegistry;
+
+  public setCompileRegistry(registry: CompileRegistry) {
+    this.compileRegistry = registry;
+  }
+
+  public getVerificationKey(name: string) {
+    if (this.compileRegistry === undefined) {
+      throw new Error("CompileRegistry hasn't been set yet");
+    }
+    const artifact = this.compileRegistry.getArtifact(name);
+    if (artifact === undefined) {
+      throw new Error(
+        `Verification Key for child program ${name} not found in registry`
+      );
+    }
+    return artifact.verificationKey;
+  }
+}

packages/common/src/index.ts

@@ -18,3 +18,7 @@ export * from "./trees/RollupMerkleTree";
 export * from "./events/EventEmitterProxy";
 export * from "./events/ReplayingSingleUseEventEmitter";
 export * from "./trees/MockAsyncMerkleStore";
+export * from "./compiling/AtomicCompileHelper";
+export * from "./compiling/CompileRegistry";
+export * from "./compiling/CompilableModule";
+export * from "./compiling/services/ChildVerificationKeyService";

packages/common/src/types.ts

@@ -1,5 +1,5 @@
 // allows to reference interfaces as 'classes' rather than instances
-import { Bool, Field, PublicKey } from "o1js";
+import { Bool, DynamicProof, Field, Proof, ProofBase, PublicKey } from "o1js";
 
 export type TypedClass<Class> = new (...args: any[]) => Class;
 
@@ -47,3 +47,12 @@ export const EMPTY_PUBLICKEY = PublicKey.fromObject({
 export type OverwriteObjectType<Base, New> = {
   [Key in keyof Base]: Key extends keyof New ? New[Key] : Base[Key];
 } & New;
+
+export type InferProofBase<
+  ProofType extends Proof<any, any> | DynamicProof<any, any>,
+> =
+  ProofType extends Proof<infer PI, infer PO>
+    ? ProofBase<PI, PO>
+    : ProofType extends DynamicProof<infer PI, infer PO>
+      ? ProofBase<PI, PO>
+      : undefined;

packages/common/src/utils.ts

@@ -6,6 +6,8 @@ import {
   Proof,
 } from "o1js";
 
+import { TypedClass } from "./types";
+
 export function requireTrue(
   condition: boolean,
   errorOrFunction: Error | (() => Error)
@@ -60,7 +62,7 @@ export function reduceSequential<T, U>(
     array: T[]
   ) => Promise<U>,
   initialValue: U
-) {
+): Promise<U> {
   return array.reduce<Promise<U>>(
     async (previousPromise, current, index, arr) => {
       const previous = await previousPromise;
@@ -170,3 +172,25 @@ type NonMethodKeys<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.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;
+}

packages/common/src/zkProgrammable/ZkProgrammable.ts

@@ -3,6 +3,8 @@ 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";
 
@@ -32,6 +34,7 @@ export interface Compile {
 }
 
 export interface PlainZkProgram<PublicInput = undefined, PublicOutput = void> {
+  name: string;
   compile: Compile;
   verify: Verify<PublicInput, PublicOutput>;
   Proof: ReturnType<
@@ -72,8 +75,6 @@ export function verifyToMockable<PublicInput, PublicOutput>(
       return verified;
     }
 
-    console.log("VerifyMocked");
-
     return proof.proof === MOCK_PROOF;
   };
 }
@@ -125,6 +126,21 @@ export abstract class ZkProgrammable<
       };
     });
   }
+
+  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<

packages/module/src/runtime/Runtime.ts

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-unsafe-argument */
 import { ZkProgram } from "o1js";
-import { DependencyContainer, injectable } from "tsyringe";
+import { container, DependencyContainer, injectable } from "tsyringe";
 import {
   StringKeyOf,
   ModuleContainer,
@@ -11,11 +11,16 @@ import {
   PlainZkProgram,
   AreProofsEnabled,
   ChildContainerProvider,
+  CompilableModule,
+  CompileRegistry,
 } from "@proto-kit/common";
 import {
   MethodPublicOutput,
   StateServiceProvider,
   SimpleAsyncStateService,
+  RuntimeMethodExecutionContext,
+  RuntimeTransaction,
+  NetworkState,
 } from "@proto-kit/protocol";
 
 import {
@@ -227,9 +232,10 @@ export class RuntimeZkProgrammable<
       return buckets;
     };
 
-    return splitRuntimeMethods().map((bucket) => {
+    return splitRuntimeMethods().map((bucket, index) => {
+      const name = `RuntimeProgram-${index}`;
       const program = ZkProgram({
-        name: "RuntimeProgram",
+        name,
         publicOutput: MethodPublicOutput,
         methods: bucket,
       });
@@ -245,6 +251,7 @@ export class RuntimeZkProgrammable<
       );
 
       return {
+        name,
         compile: program.compile.bind(program),
         verify: program.verify.bind(program),
         analyzeMethods: program.analyzeMethods.bind(program),
@@ -262,7 +269,7 @@ export class RuntimeZkProgrammable<
 @injectable()
 export class Runtime<Modules extends RuntimeModulesRecord>
   extends ModuleContainer<Modules>
-  implements RuntimeEnvironment
+  implements RuntimeEnvironment, CompilableModule
 {
   public static from<Modules extends RuntimeModulesRecord>(
     definition: RuntimeDefinition<Modules>
@@ -375,5 +382,14 @@ export class Runtime<Modules extends RuntimeModulesRecord>
   public get runtimeModuleNames() {
     return Object.keys(this.definition.modules);
   }
+
+  public async compile(registry: CompileRegistry) {
+    const context = container.resolve(RuntimeMethodExecutionContext);
+    context.setup({
+      transaction: RuntimeTransaction.dummyTransaction(),
+      networkState: NetworkState.empty(),
+    });
+    return await this.zkProgrammable.compile(registry);
+  }
 }
 /* eslint-enable @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-unsafe-argument */

packages/persistance/test-integration/SequencerRestart.test.ts

@@ -32,7 +32,7 @@ describe("sequencer restart", () => {
       },
     });
 
-    await appChain.start(container.createChildContainer());
+    await appChain.start(false, container.createChildContainer());
   };
 
   const teardown = async () => {