Smart Contracts compiling

Author: rpanic

packages/protocol/src/compiling/AtomicCompileHelper.ts renamed to packages/common/src/compiling/AtomicCompileHelper.ts

@@ -1,10 +1,11 @@
 import {
   AreProofsEnabled,
   CompileArtifact,
-  log,
   MOCK_VERIFICATION_KEY,
-} from "@proto-kit/common";
-import { SmartContract, VerificationKey } from "o1js";
+} from "../zkProgrammable/ZkProgrammable";
+import { isSubtypeOfName } from "../utils";
+import { TypedClass } from "../types";
+import { log } from "../log";
 
 export type ArtifactRecord = Record<string, CompileArtifact>;
 
@@ -29,12 +30,24 @@ export class AtomicCompileHelper {
     if (this.compilationPromises[name] === undefined) {
       const proofsEnabled =
         overrideProofsEnabled ?? this.areProofsEnabled.areProofsEnabled;
-      // This wierd any is necessary otherwise the compiler optimized that check away
-      if (proofsEnabled || !((contract as any) instanceof SmartContract)) {
+
+      // 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,
         });

packages/protocol/src/compiling/CompilableModule.ts renamed to packages/common/src/compiling/CompilableModule.ts

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

packages/common/src/compiling/CompileRegistry.ts

@@ -0,0 +1,118 @@
+import { inject, injectable, singleton } from "tsyringe";
+
+import { AreProofsEnabled } from "../zkProgrammable/ZkProgrammable";
+
+import {
+  ArtifactRecord,
+  AtomicCompileHelper,
+  CompileTarget,
+} from "./AtomicCompileHelper";
+import { CompilableModule } from "./CompilableModule";
+
+interface GenericCompilableModule<Artifact> {
+  compile(registry: CompileRegistry): Promise<Artifact>;
+}
+
+export type InferDependencyArtifacts<
+  Dependencies extends Record<string, CompilableModule>,
+> = {
+  [Key in keyof Dependencies]: Dependencies[Key] extends GenericCompilableModule<
+    infer Artifact
+  >
+    ? Artifact
+    : void;
+};
+/**
+ * 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 = {};
+
+  // private cachedModuleOutputs: Record<string, ArtifactRecord | void> = {};
+
+  // 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 async compileModule<
+  //   ReturnType extends ArtifactRecord,
+  //   Dependencies extends Record<string, CompilableModule>,
+  // >(
+  //   compile: (
+  //     compiler: AtomicCompileHelper,
+  //     args: InferDependencyArtifacts<Dependencies>
+  //   ) => Promise<ReturnType>,
+  //   dependencies?: Dependencies
+  // ): Promise<ReturnType | undefined> {
+  //   const collectedArtifacts = await mapSequential(
+  //     Object.entries(dependencies ?? {}),
+  //     async ([depName, dep]) => {
+  //       if (this.cachedModuleOutputs[depName] !== undefined) {
+  //         return [depName, this.cachedModuleOutputs[depName]];
+  //       }
+  //       const artifact = await dep.compile(this);
+  //       if (artifact !== undefined) {
+  //         this.artifacts = {
+  //           ...this.artifacts,
+  //           ...artifact,
+  //         };
+  //       }
+  //       this.cachedModuleOutputs[depName] = artifact;
+  //       return [depName, artifact];
+  //     }
+  //   );
+  //
+  //   const artifacts = await compile(
+  //     this.compile,
+  //     Object.fromEntries(collectedArtifacts)
+  //   );
+  //
+  //   this.artifacts = {
+  //     ...this.artifacts,
+  //     ...artifacts,
+  //   };
+  //
+  //   return artifacts;
+  // }
+
+  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/protocol/src/compiling/services/ChildVerificationKeyService.ts renamed to packages/common/src/compiling/services/ChildVerificationKeyService.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/index.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/types.ts

@@ -6,6 +6,8 @@ import {
   Proof,
 } from "o1js";
 
+import { TypedClass } from "./types";
+
 export function requireTrue(
   condition: boolean,
   errorOrFunction: Error | (() => Error)
@@ -164,3 +166,19 @@ type NonMethodKeys<Type> = {
   [Key in keyof Type]: Type[Key] extends Function ? never : Key;
 }[keyof Type];
 export type NonMethods<Type> = Pick<Type, NonMethodKeys<Type>>;
+
+/**
+ * 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);
+}

packages/common/src/utils.ts

@@ -4,6 +4,7 @@ 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";
 
@@ -74,8 +75,6 @@ export function verifyToMockable<PublicInput, PublicOutput>(
       return verified;
     }
 
-    console.log("VerifyMocked");
-
     return proof.proof === MOCK_PROOF;
   };
 }
@@ -128,11 +127,11 @@ export abstract class ZkProgrammable<
     });
   }
 
-  public async compile() {
+  public async compile(registry: CompileRegistry) {
     return await reduceSequential(
       this.zkProgram,
       async (acc, program) => {
-        const result = await program.compile();
+        const result = await registry.compile(program);
         return {
           ...acc,
           [program.name]: result,

packages/common/src/zkProgrammable/ZkProgrammable.ts

@@ -11,13 +11,13 @@ import {
   PlainZkProgram,
   AreProofsEnabled,
   ChildContainerProvider,
+  CompilableModule,
+  CompileRegistry,
 } from "@proto-kit/common";
 import {
   MethodPublicOutput,
   StateServiceProvider,
   SimpleAsyncStateService,
-  CompilableModule,
-  CompileRegistry,
   RuntimeMethodExecutionContext,
   RuntimeTransaction,
   NetworkState,
@@ -384,14 +384,12 @@ export class Runtime<Modules extends RuntimeModulesRecord>
   }
 
   public async compile(registry: CompileRegistry) {
-    return await registry.compileModule(async () => {
-      const context = container.resolve(RuntimeMethodExecutionContext);
-      context.setup({
-        transaction: RuntimeTransaction.dummyTransaction(),
-        networkState: NetworkState.empty(),
-      });
-      return await this.zkProgrammable.compile();
+    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 */