Dynamically set proofs enabled

proto-kit · Jan 8, 2025 · efaf83e · efaf83e
1 parent f7df3c5
commit efaf83e
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 20 deletions.
diff --git a/packages/common/src/zkProgram/WithZkProgram.ts b/packages/common/src/zkProgram/WithZkProgram.ts
@@ -55,6 +55,7 @@ export interface PlainZkProgram<PublicInput = undefined, PublicOutput = void> {
     Record<string, Awaited<ReturnType<typeof Provable.constraintSystem>>>
   >;
   proofsEnabled: boolean;
+  setProofsEnabled(proofsEnabled: boolean): void;
 }
 
 export interface ZkProgramFactory<PublicInput, PublicOutput> {

diff --git a/packages/module/src/runtime/Runtime.ts b/packages/module/src/runtime/Runtime.ts
@@ -72,7 +72,10 @@ export interface RuntimeDefinition<Modules extends RuntimeModulesRecord> {
 export class RuntimeZkProgramFactory<Modules extends RuntimeModulesRecord>
   implements ZkProgramFactory<undefined, MethodPublicOutput>
 {
-  public constructor(public runtime: Runtime<Modules>) {}
+  public constructor(
+    public runtime: Runtime<Modules>,
+    private readonly areProofsEnabled: boolean
+  ) {}
 
   public get appChain() {
     return this.runtime.appChain;
@@ -231,7 +234,7 @@ export class RuntimeZkProgramFactory<Modules extends RuntimeModulesRecord>
         publicOutput: MethodPublicOutput,
         methods: bucket,
       });
-
+      program.setProofsEnabled(this.areProofsEnabled);
       const SelfProof = ZkProgram.Proof(program);
 
       const methods = Object.keys(bucket).reduce<Record<string, any>>(
@@ -246,8 +249,8 @@ export class RuntimeZkProgramFactory<Modules extends RuntimeModulesRecord>
         compile: program.compile.bind(program),
         verify: program.verify.bind(program),
         analyzeMethods: program.analyzeMethods.bind(program),
-        // TODO set this dynamically somehow
-        proofsEnabled: true,
+        proofsEnabled: program.proofsEnabled,
+        setProofsEnabled: program.setProofsEnabled.bind(program),
         Proof: SelfProof,
         methods,
       };
@@ -274,12 +277,25 @@ export class Runtime<Modules extends RuntimeModulesRecord>
     };
   }
 
-  // runtime modules composed into a ZkProgram
-  public zkProgram: PlainZkProgram<undefined, MethodPublicOutput>[];
+  private _zkProgram?: PlainZkProgram<undefined, MethodPublicOutput>[];
+
+  private _zkProgramFactory?: ZkProgramFactory<undefined, MethodPublicOutput>;
 
   public definition: RuntimeDefinition<Modules>;
 
-  public zkProgramFactory: ZkProgramFactory<undefined, MethodPublicOutput>;
+  // runtime modules composed into a ZkProgram
+  public get zkProgram(): PlainZkProgram<undefined, MethodPublicOutput>[] {
+    // eslint-disable-next-line no-underscore-dangle
+    return this._zkProgram!;
+  }
+
+  public get zkProgramFactory(): ZkProgramFactory<
+    undefined,
+    MethodPublicOutput
+  > {
+    // eslint-disable-next-line no-underscore-dangle
+    return this._zkProgramFactory!;
+  }
 
   /**
    * Creates a new Runtime from the provided config
@@ -289,15 +305,20 @@ export class Runtime<Modules extends RuntimeModulesRecord>
   public constructor(definition: RuntimeDefinition<Modules>) {
     super(definition);
     this.definition = definition;
-    this.zkProgramFactory = new RuntimeZkProgramFactory<Modules>(this);
-    this.zkProgram = this.zkProgramFactory.zkProgramFactory();
   }
 
   // TODO Remove after changing DFs to type-based approach
   public create(childContainerProvider: ChildContainerProvider) {
     super.create(childContainerProvider);
 
     this.useDependencyFactory(this.container.resolve(MethodIdFactory));
+    // eslint-disable-next-line no-underscore-dangle
+    this._zkProgramFactory = new RuntimeZkProgramFactory<Modules>(
+      this,
+      this.appChain?.areProofsEnabled!
+    );
+    // eslint-disable-next-line no-underscore-dangle
+    this._zkProgram = this.zkProgramFactory.zkProgramFactory();
   }
 
   public get appChain(): AreProofsEnabled | undefined {

diff --git a/packages/protocol/src/prover/block/BlockProver.ts b/packages/protocol/src/prover/block/BlockProver.ts
@@ -139,7 +139,8 @@ export class BlockProverFactory
     public readonly runtime: PlainZkProgram<undefined, MethodPublicOutput>[],
     private readonly transactionHooks: ProvableTransactionHook<unknown>[],
     private readonly blockHooks: ProvableBlockHook<unknown>[],
-    private readonly verificationKeyService: MinimalVKTreeService
+    private readonly verificationKeyService: MinimalVKTreeService,
+    private readonly areProofsEnabled: boolean
   ) {}
 
   public get appChain(): AreProofsEnabled | undefined {
@@ -864,6 +865,7 @@ export class BlockProverFactory
         },
       },
     });
+    program.setProofsEnabled(this.areProofsEnabled);
 
     const methods = {
       proveTransaction: program.proveTransaction,
@@ -878,8 +880,8 @@ export class BlockProverFactory
         verify: program.verify.bind(program),
         analyzeMethods: program.analyzeMethods.bind(program),
         Proof: SelfProofClass,
-        // TODO Set this dynamically
-        proofsEnabled: true,
+        proofsEnabled: program.proofsEnabled,
+        setProofsEnabled: program.setProofsEnabled.bind(program),
         methods,
       },
     ];
@@ -912,7 +914,9 @@ export class BlockProver extends ProtocolModule implements BlockProvable {
     transactionHooks: ProvableTransactionHook<unknown>[],
     @injectAll("ProvableBlockHook")
     blockHooks: ProvableBlockHook<unknown>[],
-    verificationKeyService: RuntimeVerificationKeyRootService
+    verificationKeyService: RuntimeVerificationKeyRootService,
+    @inject("AreProofsEnabled")
+    proofsEnabled: AreProofsEnabled
   ) {
     super();
     this.zkProgramFactory = new BlockProverFactory(
@@ -922,7 +926,8 @@ export class BlockProver extends ProtocolModule implements BlockProvable {
       runtime.zkProgram,
       transactionHooks,
       blockHooks,
-      verificationKeyService
+      verificationKeyService,
+      proofsEnabled.areProofsEnabled
     );
     this.zkProgram = this.zkProgramFactory.zkProgramFactory();
   }

diff --git a/packages/protocol/src/prover/statetransition/StateTransitionProver.ts b/packages/protocol/src/prover/statetransition/StateTransitionProver.ts
@@ -5,7 +5,7 @@ import {
   ZkProgramFactory,
 } from "@proto-kit/common";
 import { Field, Provable, SelfProof, ZkProgram } from "o1js";
-import { injectable } from "tsyringe";
+import { inject, injectable } from "tsyringe";
 
 import { constants } from "../../Constants";
 import { ProvableStateTransition } from "../../model/StateTransition";
@@ -64,7 +64,8 @@ export class StateTransitionProverFactory
     >
 {
   public constructor(
-    private readonly stateTransitionProver: StateTransitionProver
+    private readonly stateTransitionProver: StateTransitionProver,
+    private readonly areProofsEnabled: boolean
   ) {}
 
   public get appChain(): AreProofsEnabled | undefined {
@@ -113,6 +114,7 @@ export class StateTransitionProverFactory
         },
       },
     });
+    program.setProofsEnabled(this.areProofsEnabled);
 
     const methods = {
       proveBatch: program.proveBatch.bind(program),
@@ -127,8 +129,8 @@ export class StateTransitionProverFactory
         verify: program.verify.bind(program),
         analyzeMethods: program.analyzeMethods.bind(program),
         Proof: SelfProofClass,
-        // TODO Set this dynamically
-        proofsEnabled: true,
+        proofsEnabled: program.proofsEnabled,
+        setProofsEnabled: program.setProofsEnabled.bind(program),
         methods,
       },
     ];
@@ -344,9 +346,15 @@ export class StateTransitionProver
     StateTransitionProverPublicOutput
   >[];
 
-  public constructor() {
+  public constructor(
+    @inject("AreProofsEnabled")
+    proofsEnabled: AreProofsEnabled
+  ) {
     super();
-    this.zkProgramFactory = new StateTransitionProverFactory(this);
+    this.zkProgramFactory = new StateTransitionProverFactory(
+      this,
+      proofsEnabled.areProofsEnabled
+    );
     this.zkProgram = this.zkProgramFactory.zkProgramFactory();
   }