Added ChildVerificationKeyService to allow for sideloading in protocol circuits

Author: rpanic

packages/protocol/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/protocol/src/index.ts

@@ -1,6 +1,7 @@
 export * from "./compiling/AtomicCompileHelper";
 export * from "./compiling/CompileRegistry";
 export * from "./compiling/CompilableModule";
+export * from "./compiling/services/ChildVerificationKeyService";
 export * from "./hooks/AccountStateHook";
 export * from "./hooks/BlockHeightHook";
 export * from "./hooks/LastStateRootBlockHook";

packages/sequencer/src/protocol/production/tasks/CircuitCompilerTask.ts

@@ -9,6 +9,7 @@ import {
   CompileRegistry,
   CompilableModule,
   BridgeContractProtocolModule,
+  RuntimeVerificationKeyRootService,
 } from "@proto-kit/protocol";
 
 import { TaskSerializer } from "../../../worker/flow/Task";
@@ -18,6 +19,7 @@ import { VerificationKeySerializer } from "../helpers/VerificationKeySerializer"
 export type CompilerTaskParams = {
   existingArtifacts: ArtifactRecord;
   targets: string[];
+  runtimeVKRoot?: string;
 };
 
 type SerializedArtifactRecord = Record<
@@ -117,6 +119,14 @@ export class CircuitCompilerTask extends UnpreparingTask<
   public async compute(input: CompilerTaskParams): Promise<ArtifactRecord> {
     this.compileRegistry.addArtifactsRaw(input.existingArtifacts);
 
+    // We need to initialize the VK tree root if we have it, so that
+    // the BlockProver can bake in that root
+    if (input.runtimeVKRoot !== undefined) {
+      this.protocol.dependencyContainer
+        .resolve(RuntimeVerificationKeyRootService)
+        .setRoot(BigInt(input.runtimeVKRoot));
+    }
+
     log.info("Computing VKs");
 
     // TODO make adaptive
@@ -132,7 +142,10 @@ export class CircuitCompilerTask extends UnpreparingTask<
       if (target in targets) {
         await targets[target].compile(this.compileRegistry);
       } else {
-        throw new Error(`Compile target ${target} not found`);
+        log.info(
+          // TODO Is that right? Or should we check that the bridge exists on the sequencer side?
+          `Compile target ${target} not found, skipping`
+        );
       }
     });
     log.timeEnd.info(msg);

packages/sequencer/src/protocol/runtime/RuntimeVerificationKeyService.ts

@@ -93,7 +93,7 @@ export class VerificationKeyService extends ConfigurableModule<{}> {
     );
   }
 
-  public async initializeVKTreeFromMethodMappings(verificationKeys: VKRecord) {
+  private async initializeVKTreeFromMethodMappings(verificationKeys: VKRecord) {
     const tree = new VKTree(new InMemoryMerkleTreeStorage());
     const valuesVK: Record<string, { data: string; hash: Field }> = {};
     const indexes: VKIndexes = {};

packages/sequencer/src/sequencer/SequencerStartupModule.ts

@@ -1,16 +1,20 @@
 import { inject } from "tsyringe";
 import {
   ArtifactRecord,
+  CompileRegistry,
   MandatoryProtocolModulesRecord,
   Protocol,
   RuntimeVerificationKeyRootService,
   SettlementSmartContractBase,
 } from "@proto-kit/protocol";
-import { log } from "@proto-kit/common";
+import { CompileArtifact, log } from "@proto-kit/common";
 
-import { FlowCreator } from "../worker/flow/Flow";
+import { Flow, FlowCreator } from "../worker/flow/Flow";
 import { WorkerRegistrationFlow } from "../worker/worker/startup/WorkerRegistrationFlow";
-import { CircuitCompilerTask } from "../protocol/production/tasks/CircuitCompilerTask";
+import {
+  CircuitCompilerTask,
+  CompilerTaskParams,
+} from "../protocol/production/tasks/CircuitCompilerTask";
 import { VerificationKeyService } from "../protocol/runtime/RuntimeVerificationKeyService";
 
 import { SequencerModule, sequencerModule } from "./builder/SequencerModule";
@@ -23,47 +27,114 @@ export class SequencerStartupModule extends SequencerModule {
     private readonly protocol: Protocol<MandatoryProtocolModulesRecord>,
     private readonly compileTask: CircuitCompilerTask,
     private readonly verificationKeyService: VerificationKeyService,
-    private readonly registrationFlow: WorkerRegistrationFlow
+    private readonly registrationFlow: WorkerRegistrationFlow,
+    private readonly compileRegistry: CompileRegistry
   ) {
     super();
   }
 
-  public async start() {
-    const flow = this.flowCreator.createFlow("compile-circuits", {});
+  private async pushCompileTask(
+    flow: Flow<{}>,
+    payload: CompilerTaskParams
+  ): Promise<ArtifactRecord> {
+    return await flow.withFlow<ArtifactRecord>(async (res, rej) => {
+      await flow.pushTask(this.compileTask, payload, async (result) => {
+        res(result);
+      });
+    });
+  }
 
-    log.info("Compiling Protocol circuits, this can take a few minutes");
+  public async compileRuntime(flow: Flow<{}>) {
+    const artifacts = await this.pushCompileTask(flow, {
+      existingArtifacts: {},
+      targets: ["runtime"],
+      runtimeVKRoot: undefined,
+    });
+
+    // Init runtime VK tree
+    await this.verificationKeyService.initializeVKTree(artifacts);
+
+    const root = this.verificationKeyService.getRoot();
+
+    this.protocol.dependencyContainer
+      .resolve(RuntimeVerificationKeyRootService)
+      .setRoot(root);
+
+    this.compileRegistry.addArtifactsRaw(artifacts);
+
+    return root;
+  }
+
+  private async compileProtocolAndBridge(flow: Flow<{}>) {
+    // Can happen in parallel
+    type ParallelResult = {
+      protocol?: ArtifactRecord;
+      bridge?: ArtifactRecord;
+    };
+    const result = await flow.withFlow<ArtifactRecord>(async (res, rej) => {
+      const results: ParallelResult = {};
+
+      const resolveIfPossible = () => {
+        const { bridge, protocol } = results;
+        if (bridge !== undefined && protocol !== undefined) {
+          res({ ...protocol, ...bridge });
+        }
+      };
 
-    const artifacts = await flow.withFlow<ArtifactRecord>(async (res, rej) => {
       await flow.pushTask(
         this.compileTask,
-        { existingArtifacts: {}, targets: ["runtime"] },
+        {
+          existingArtifacts: {},
+          targets: ["protocol"],
+          runtimeVKRoot: undefined,
+        },
         async (result) => {
-          res(result);
+          results.protocol = result;
+          resolveIfPossible();
+        }
+      );
+
+      await flow.pushTask(
+        this.compileTask,
+        {
+          existingArtifacts: {},
+          targets: ["bridge"],
+          runtimeVKRoot: undefined,
+        },
+        async (result) => {
+          results.bridge = result;
+          resolveIfPossible();
         }
       );
     });
+    this.compileRegistry.addArtifactsRaw(result);
+    return result;
+  }
 
-    log.info("Protocol circuits compiled");
+  public async start() {
+    const flow = this.flowCreator.createFlow("compile-circuits", {});
 
-    // Init runtime VK tree
-    await this.verificationKeyService.initializeVKTree(artifacts);
+    log.info("Compiling Protocol circuits, this can take a few minutes");
 
-    const root = this.verificationKeyService.getRoot();
+    const root = await this.compileRuntime(flow);
 
-    this.protocol.dependencyContainer
-      .resolve(RuntimeVerificationKeyRootService)
-      .setRoot(root);
+    const protocolBridgeArtifacts = await this.compileProtocolAndBridge(flow);
+
+    log.info("Protocol circuits compiled");
 
     // Init BridgeContract vk for settlement contract
-    const bridgeVk = artifacts.BridgeContract;
+    const bridgeVk = protocolBridgeArtifacts.BridgeContract;
     if (bridgeVk !== undefined) {
       SettlementSmartContractBase.args.BridgeContractVerificationKey =
         bridgeVk.verificationKey;
     }
 
+    // TODO Add vk record to start params
+
     await this.registrationFlow.start({
       runtimeVerificationKeyRoot: root,
-      bridgeContractVerificationKey: bridgeVk.verificationKey,
+      bridgeContractVerificationKey: bridgeVk?.verificationKey,
+      compiledArtifacts: this.compileRegistry.getAllArtifacts(),
     });
 
     log.info("Protocol circuits compiled successfully, commencing startup");

packages/sequencer/src/worker/worker/startup/WorkerRegistrationTask.ts

@@ -1,6 +1,9 @@
 import { log, noop } from "@proto-kit/common";
 import { inject, injectable } from "tsyringe";
 import {
+  ArtifactRecord,
+  ChildVerificationKeyService,
+  CompileRegistry,
   Protocol,
   RuntimeVerificationKeyRootService,
   SettlementSmartContractBase,
@@ -12,11 +15,13 @@ import { AbstractStartupTask } from "../../flow/AbstractStartupTask";
 import { VerificationKeySerializer } from "../../../protocol/production/helpers/VerificationKeySerializer";
 
 import { CloseWorkerError } from "./CloseWorkerError";
+import { ArtifactRecordSerializer } from "../../../protocol/production/tasks/CircuitCompilerTask";
 
 export type WorkerStartupPayload = {
   runtimeVerificationKeyRoot: bigint;
   // This has to be nullable, since
   bridgeContractVerificationKey?: VerificationKey;
+  compiledArtifacts: ArtifactRecord;
 };
 
 @injectable()
@@ -28,7 +33,8 @@ export class WorkerRegistrationTask
   private done = false;
 
   public constructor(
-    @inject("Protocol") private readonly protocol: Protocol<any>
+    @inject("Protocol") private readonly protocol: Protocol<any>,
+    private readonly compileRegistry: CompileRegistry
   ) {
     super();
   }
@@ -55,13 +61,19 @@ export class WorkerRegistrationTask
         input.bridgeContractVerificationKey;
     }
 
+    this.compileRegistry.addArtifactsRaw(input.compiledArtifacts);
+    this.protocol.dependencyContainer
+      .resolve(ChildVerificationKeyService)
+      .setCompileRegistry(this.compileRegistry);
+
     this.events.emit("startup-task-finished");
 
     this.done = true;
     return true;
   }
 
   public inputSerializer() {
+    const artifactSerializer = new ArtifactRecordSerializer();
     return {
       toJSON: (payload: WorkerStartupPayload) => {
         return JSON.stringify({
@@ -73,6 +85,9 @@ export class WorkerRegistrationTask
                   payload.bridgeContractVerificationKey
                 )
               : undefined,
+          compiledArtifacts: artifactSerializer.toJSON(
+            payload.compiledArtifacts
+          ),
         });
       },
       fromJSON: (payload: string) => {
@@ -91,6 +106,9 @@ export class WorkerRegistrationTask
                   jsonObject.bridgeContractVerificationKey
                 )
               : undefined,
+          compiledArtifacts: artifactSerializer.fromJSON(
+            jsonObject.compiledArtifacts
+          ),
         };
       },
     };