Added new CompileRegistry with adaptive compiling and caching

Author: rpanic

packages/common/src/utils.ts

@@ -56,7 +56,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;

packages/common/src/zkProgrammable/ZkProgrammable.ts

@@ -3,6 +3,7 @@ import { Memoize } from "typescript-memoize";
 
 import { log } from "../log";
 import { dummyVerificationKey } from "../dummyVerificationKey";
+import { reduceSequential } from "../utils";
 
 import { MOCK_PROOF } from "./provableMethod";
 
@@ -126,6 +127,21 @@ export abstract class ZkProgrammable<
       };
     });
   }
+
+  public async compile() {
+    return await reduceSequential(
+      this.zkProgram,
+      async (acc, program) => {
+        const result = await program.compile();
+        return {
+          ...acc,
+          [program.name]: result,
+        };
+      },
+      // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+      {} as Record<string, CompileArtifact>
+    );
+  }
 }
 
 export interface WithZkProgrammable<

packages/protocol/src/compiling/AtomicCompileHelper.ts

@@ -0,0 +1,61 @@
+import {
+  AreProofsEnabled,
+  log,
+  MOCK_VERIFICATION_KEY,
+} from "@proto-kit/common";
+
+export type Artifact = string | object | undefined;
+
+export type GenericCompileTarget<T extends Artifact> = {
+  compile: () => Promise<T>;
+};
+
+export class AtomicCompileHelper {
+  public constructor(private readonly areProofsEnabled: AreProofsEnabled) {}
+
+  private compilationPromises: {
+    [key: string]: Promise<Artifact | undefined>;
+  } = {};
+
+  // Generic params for zkProgrammable should be unknown, but verify makes those types invariant
+  // public async zkProgrammable(zkProgrammable: ZkProgrammable<any, any>) {
+  //   await reduceSequential(
+  //     zkProgrammable.zkProgram,
+  //     async (acc, program) => {
+  //       const res = await this.program(program);
+  //       return {
+  //         ...acc,
+  //         [program.name]: res,
+  //       };
+  //     },
+  //     {}
+  //   );
+  // }
+
+  public async program<ReturnArtifact extends Artifact>(
+    name: string,
+    contract: GenericCompileTarget<ReturnArtifact>,
+    overrideProofsEnabled?: boolean
+  ): Promise<ReturnArtifact> {
+    let newPromise = false;
+    if (this.compilationPromises[name] === undefined) {
+      const proofsEnabled =
+        overrideProofsEnabled ?? this.areProofsEnabled.areProofsEnabled;
+      if (proofsEnabled) {
+        log.time(`Compiling ${name}`);
+        this.compilationPromises[name] = contract.compile();
+        newPromise = true;
+      } else {
+        // TODO Mock VK here is not at all generalized and safe
+        //  Better way would be to package the Smart contracts into a mock-compile as well
+        this.compilationPromises[name] = Promise.resolve(MOCK_VERIFICATION_KEY);
+      }
+    }
+    const result = await this.compilationPromises[name];
+    if (newPromise) {
+      log.timeEnd.info(`Compiling ${name}`);
+    }
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    return result as ReturnArtifact;
+  }
+}

packages/protocol/src/compiling/CompilableModule.ts

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

packages/protocol/src/compiling/CompileRegistry.ts

@@ -0,0 +1,89 @@
+import { inject, injectable, singleton } from "tsyringe";
+import {
+  AreProofsEnabled,
+  mapSequential,
+  MOCK_VERIFICATION_KEY,
+} from "@proto-kit/common";
+
+import {
+  Artifact,
+  AtomicCompileHelper,
+  GenericCompileTarget,
+} from "./AtomicCompileHelper";
+import { CompilableModule } from "./CompilableModule";
+
+/**
+ * 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.compile = new AtomicCompileHelper(this.areProofsEnabled);
+  }
+
+  public compile: AtomicCompileHelper;
+
+  private artifacts: Record<string, Artifact | "compiled-but-no-artifact"> = {};
+
+  public async compileModule<ReturnArtifact extends Artifact>(
+    // TODO Make name inferred by the module token
+    name: string,
+    compile: (
+      registry: CompileRegistry,
+      ...args: unknown[]
+    ) => GenericCompileTarget<ReturnArtifact>,
+    dependencies: Record<string, CompilableModule> = {}
+  ): Promise<ReturnArtifact | undefined> {
+    const collectedArtifacts = await mapSequential(
+      Object.entries(dependencies),
+      async ([depName, dep]) => {
+        if (this.artifacts[depName] !== undefined) {
+          return this.artifacts[depName];
+        }
+        const artifact =
+          (await dep.compile(this)) ?? "compiled-but-no-artifact";
+        this.artifacts[depName] = artifact;
+        return artifact;
+      }
+    );
+
+    const target = compile(this, ...collectedArtifacts);
+    const artifact = await this.compile.program<ReturnArtifact>(name, target);
+
+    this.artifacts[name] = artifact ?? "compiled-but-no-artifact";
+
+    return artifact;
+  }
+
+  public getArtifact<ArtifactType extends Artifact>(name: string) {
+    if (this.artifacts[name] === undefined) {
+      throw new Error(
+        `Artifact for ${name} not available, did you compile it via the CompileRegistry?`
+      );
+    }
+    if (!this.areProofsEnabled.areProofsEnabled) {
+      return MOCK_VERIFICATION_KEY;
+    }
+
+    const artifact = this.artifacts[name];
+    if (artifact === "compiled-but-no-artifact") {
+      throw new Error(
+        `Module ${name} didn't return the requested artifact even though proofs are enabled`
+      );
+    }
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    return artifact as ArtifactType;
+  }
+
+  public addArtifactsRaw(artifacts: Record<string, Artifact>) {
+    Object.entries(artifacts).forEach(([key, value]) => {
+      this.artifacts[key] = value;
+    });
+  }
+}

packages/protocol/src/index.ts

@@ -1,3 +1,6 @@
+export * from "./compiling/AtomicCompileHelper";
+export * from "./compiling/CompileRegistry";
+export * from "./compiling/CompilableModule";
 export * from "./hooks/AccountStateHook";
 export * from "./hooks/BlockHeightHook";
 export * from "./hooks/LastStateRootBlockHook";
@@ -17,7 +20,6 @@ export * from "./utils/MinaPrefixedProvableHashList";
 export * from "./utils/ProvableReductionHashList";
 export * from "./utils/StateTransitionReductionList";
 export * from "./utils/utils";
-export * from "./utils/CompileRegistry";
 export * from "./prover/block/BlockProver";
 export * from "./prover/block/BlockProvable";
 export * from "./prover/block/accummulators/RuntimeVerificationKeyTree";

packages/protocol/src/prover/block/BlockProvable.ts

@@ -16,6 +16,7 @@ import { NetworkState } from "../../model/network/NetworkState";
 
 import { BlockHashMerkleTreeWitness } from "./accummulators/BlockHashMerkleTree";
 import { RuntimeVerificationKeyAttestation } from "./accummulators/RuntimeVerificationKeyTree";
+import { CompilableModule } from "../../compiling/CompilableModule";
 
 export class BlockProverPublicInput extends Struct({
   transactionsHash: Field,
@@ -77,7 +78,8 @@ export class DynamicRuntimeProof extends DynamicProof<
 }
 
 export interface BlockProvable
-  extends WithZkProgrammable<BlockProverPublicInput, BlockProverPublicOutput> {
+  extends WithZkProgrammable<BlockProverPublicInput, BlockProverPublicOutput>,
+    CompilableModule {
   proveTransaction: (
     publicInput: BlockProverPublicInput,
     stateProof: StateTransitionProof,

packages/protocol/src/prover/block/BlockProver.ts

@@ -12,6 +12,7 @@ import {
 import { container, inject, injectable, injectAll } from "tsyringe";
 import {
   AreProofsEnabled,
+  CompileArtifact,
   PlainZkProgram,
   provableMethod,
   WithZkProgrammable,
@@ -41,6 +42,7 @@ import {
   MinaActionsHashList,
 } from "../../utils/MinaPrefixedProvableHashList";
 import { StateTransitionReductionList } from "../../utils/StateTransitionReductionList";
+import { CompileRegistry } from "../../compiling/CompileRegistry";
 
 import {
   BlockProvable,
@@ -60,6 +62,7 @@ import {
   RuntimeVerificationKeyAttestation,
 } from "./accummulators/RuntimeVerificationKeyTree";
 import { RuntimeVerificationKeyRootService } from "./services/RuntimeVerificationKeyRootService";
+import { CompilableModule } from "../../compiling/CompilableModule";
 
 const errors = {
   stateProofNotStartingAtZero: () =>
@@ -149,6 +152,8 @@ export class BlockProverProgrammable extends ZkProgrammable<
     super();
   }
 
+  name = "BlockProver";
+
   public get areProofsEnabled(): AreProofsEnabled | undefined {
     return this.prover.areProofsEnabled;
   }
@@ -889,7 +894,10 @@ export class BlockProverProgrammable extends ZkProgrammable<
  * then be merged to be committed to the base-layer contract
  */
 @injectable()
-export class BlockProver extends ProtocolModule implements BlockProvable {
+export class BlockProver
+  extends ProtocolModule
+  implements BlockProvable, CompilableModule
+{
   public zkProgrammable: BlockProverProgrammable;
 
   public constructor(
@@ -917,6 +925,15 @@ export class BlockProver extends ProtocolModule implements BlockProvable {
     );
   }
 
+  public async compile(
+    registry: CompileRegistry
+  ): Promise<Record<string, CompileArtifact> | undefined> {
+    return await registry.compileModule(
+      "BlockProver",
+      () => this.zkProgrammable
+    );
+  }
+
   public proveTransaction(
     publicInput: BlockProverPublicInput,
     stateProof: StateTransitionProof,

packages/protocol/src/prover/statetransition/StateTransitionProvable.ts

@@ -4,6 +4,7 @@ import { WithZkProgrammable } from "@proto-kit/common";
 import { StateTransitionProvableBatch } from "../../model/StateTransitionProvableBatch";
 
 import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
+import { CompilableModule } from "../../compiling/CompilableModule";
 
 export class StateTransitionProverPublicInput extends Struct({
   stateTransitionsHash: Field,
@@ -26,9 +27,10 @@ export type StateTransitionProof = Proof<
 
 export interface StateTransitionProvable
   extends WithZkProgrammable<
-    StateTransitionProverPublicInput,
-    StateTransitionProverPublicOutput
-  > {
+      StateTransitionProverPublicInput,
+      StateTransitionProverPublicOutput
+    >,
+    CompilableModule {
   witnessProviderReference: StateTransitionWitnessProviderReference;
 
   runBatch: (

packages/protocol/src/prover/statetransition/StateTransitionProver.ts

@@ -29,6 +29,9 @@ import {
 } from "./StateTransitionProvable";
 import { StateTransitionWitnessProvider } from "./StateTransitionWitnessProvider";
 import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
+import { CompilableModule } from "../../compiling/CompilableModule";
+import { Artifact } from "../../compiling/AtomicCompileHelper";
+import { CompileRegistry } from "../../compiling/CompileRegistry";
 
 const errors = {
   propertyNotMatching: (property: string, step: string) =>
@@ -346,7 +349,10 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
 @injectable()
 export class StateTransitionProver
   extends ProtocolModule
-  implements StateTransitionProvable, StateTransitionProverType
+  implements
+    StateTransitionProvable,
+    StateTransitionProverType,
+    CompilableModule
 {
   public zkProgrammable: StateTransitionProverProgrammable;
 
@@ -361,6 +367,13 @@ export class StateTransitionProver
     );
   }
 
+  public compile(registry: CompileRegistry): Promise<void | Artifact> {
+    return registry.compileModule(
+      "StateTransitionProver",
+      () => this.zkProgrammable
+    );
+  }
+
   public runBatch(
     publicInput: StateTransitionProverPublicInput,
     batch: StateTransitionProvableBatch