Merge pull request #222 from proto-kit/feature/refactor-stprover-merkle-witness

Refactor STProver to Accept Merkle Witness as Argument

Author: ejMina226

packages/protocol/src/index.ts

@@ -24,8 +24,6 @@ export * from "./prover/block/accummulators/BlockHashMerkleTree";
 export * from "./prover/block/services/RuntimeVerificationKeyRootService";
 export * from "./prover/statetransition/StateTransitionProver";
 export * from "./prover/statetransition/StateTransitionProvable";
-export * from "./prover/statetransition/StateTransitionWitnessProvider";
-export * from "./prover/statetransition/StateTransitionWitnessProviderReference";
 export * from "./protocol/Protocol";
 export * from "./protocol/ProtocolModule";
 export * from "./protocol/ProtocolEnvironment";

packages/protocol/src/model/StateTransitionProvableBatch.ts

@@ -1,5 +1,10 @@
 import { Bool, Provable, Struct } from "o1js";
-import { range } from "@proto-kit/common";
+import {
+  InMemoryMerkleTreeStorage,
+  range,
+  RollupMerkleTree,
+  RollupMerkleTreeWitness,
+} from "@proto-kit/common";
 
 import { constants } from "../Constants";
 
@@ -60,16 +65,22 @@ export class StateTransitionProvableBatch extends Struct({
     ProvableStateTransitionType,
     constants.stateTransitionProverBatchSize
   ),
+
+  merkleWitnesses: Provable.Array(
+    RollupMerkleTreeWitness,
+    constants.stateTransitionProverBatchSize
+  ),
 }) {
   public static fromMappings(
     transitions: {
       transition: ProvableStateTransition;
       type: ProvableStateTransitionType;
-    }[]
+    }[],
+    merkleWitnesses: RollupMerkleTreeWitness[]
   ): StateTransitionProvableBatch {
     const batch = transitions.map((entry) => entry.transition);
     const transitionTypes = transitions.map((entry) => entry.type);
-
+    const witnesses = merkleWitnesses.slice();
     // Check that order is correct
     let normalSTsStarted = false;
     transitionTypes.forEach((x) => {
@@ -84,16 +95,23 @@ export class StateTransitionProvableBatch extends Struct({
     while (batch.length < constants.stateTransitionProverBatchSize) {
       batch.push(ProvableStateTransition.dummy());
       transitionTypes.push(ProvableStateTransitionType.normal);
+      witnesses.push(
+        new RollupMerkleTree(new InMemoryMerkleTreeStorage()).getWitness(
+          BigInt(0)
+        )
+      );
     }
     return new StateTransitionProvableBatch({
       batch,
       transitionTypes,
+      merkleWitnesses: witnesses,
     });
   }
 
   public static fromTransitions(
     transitions: ProvableStateTransition[],
-    protocolTransitions: ProvableStateTransition[]
+    protocolTransitions: ProvableStateTransition[],
+    merkleWitnesses: RollupMerkleTreeWitness[]
   ): StateTransitionProvableBatch {
     const array = transitions.slice().concat(protocolTransitions);
 
@@ -113,12 +131,14 @@ export class StateTransitionProvableBatch extends Struct({
     return new StateTransitionProvableBatch({
       batch: array,
       transitionTypes,
+      merkleWitnesses,
     });
   }
 
   private constructor(object: {
     batch: ProvableStateTransition[];
     transitionTypes: ProvableStateTransitionType[];
+    merkleWitnesses: RollupMerkleTreeWitness[];
   }) {
     super(object);
   }

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

@@ -3,8 +3,6 @@ import { WithZkProgrammable } from "@proto-kit/common";
 
 import { StateTransitionProvableBatch } from "../../model/StateTransitionProvableBatch";
 
-import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
-
 export class StateTransitionProverPublicInput extends Struct({
   stateTransitionsHash: Field,
   protocolTransitionsHash: Field,
@@ -29,8 +27,6 @@ export interface StateTransitionProvable
     StateTransitionProverPublicInput,
     StateTransitionProverPublicOutput
   > {
-  witnessProviderReference: StateTransitionWitnessProviderReference;
-
   runBatch: (
     publicInput: StateTransitionProverPublicInput,
     batch: StateTransitionProvableBatch

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

@@ -27,8 +27,6 @@ import {
   StateTransitionProverPublicInput,
   StateTransitionProverPublicOutput,
 } from "./StateTransitionProvable";
-import { StateTransitionWitnessProvider } from "./StateTransitionWitnessProvider";
-import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
 
 const errors = {
   propertyNotMatching: (property: string, step: string) =>
@@ -64,8 +62,7 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
   StateTransitionProverPublicOutput
 > {
   public constructor(
-    private readonly stateTransitionProver: StateTransitionProver,
-    public readonly witnessProviderReference: StateTransitionWitnessProviderReference
+    private readonly stateTransitionProver: StateTransitionProver
   ) {
     super();
   }
@@ -132,14 +129,6 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
     ];
   }
 
-  private get witnessProvider(): StateTransitionWitnessProvider {
-    const provider = this.witnessProviderReference.getWitnessProvider();
-    if (provider === undefined) {
-      throw errors.noWitnessProviderSet();
-    }
-    return provider;
-  }
-
   /**
    * Applies the state transitions to the current stateRoot
    * and returns the new prover state
@@ -168,12 +157,19 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
 
     const transitions = transitionBatch.batch;
     const types = transitionBatch.transitionTypes;
+    const merkleWitness = transitionBatch.merkleWitnesses;
     for (
       let index = 0;
       index < constants.stateTransitionProverBatchSize;
       index++
     ) {
-      this.applyTransition(state, transitions[index], types[index], index);
+      this.applyTransition(
+        state,
+        transitions[index],
+        types[index],
+        merkleWitness[index],
+        index
+      );
     }
 
     return state;
@@ -187,13 +183,10 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
     state: StateTransitionProverExecutionState,
     transition: ProvableStateTransition,
     type: ProvableStateTransitionType,
+    merkleWitness: RollupMerkleTreeWitness,
     index = 0
   ) {
-    const witness = Provable.witness(RollupMerkleTreeWitness, () =>
-      this.witnessProvider.getWitness(transition.path)
-    );
-
-    const membershipValid = witness.checkMembership(
+    const membershipValid = merkleWitness.checkMembership(
       state.stateRoot,
       transition.path,
       transition.from.value
@@ -208,7 +201,7 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
         )
       );
 
-    const newRoot = witness.calculateRoot(transition.to.value);
+    const newRoot = merkleWitness.calculateRoot(transition.to.value);
 
     state.stateRoot = Provable.if(
       transition.to.isSome,
@@ -343,15 +336,9 @@ export class StateTransitionProver
 {
   public zkProgrammable: StateTransitionProverProgrammable;
 
-  public constructor(
-    // Injected
-    public readonly witnessProviderReference: StateTransitionWitnessProviderReference
-  ) {
+  public constructor() {
     super();
-    this.zkProgrammable = new StateTransitionProverProgrammable(
-      this,
-      witnessProviderReference
-    );
+    this.zkProgrammable = new StateTransitionProverProgrammable(this);
   }
 
   public runBatch(

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

@@ -7,7 +7,6 @@ import {
   BlockProver,
   LastStateRootBlockHook,
   MethodPublicOutput,
-  NoOpStateTransitionWitnessProvider,
   Protocol,
   StateTransitionProver,
 } from "../src";
@@ -39,7 +38,6 @@ export function createAndInitTestingProtocol() {
   protocol.create(() => container.createChildContainer());
 
   protocol.registerValue({
-    StateTransitionWitnessProvider: new NoOpStateTransitionWitnessProvider(),
     Runtime: new RuntimeMock(),
   });
 

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

@@ -60,14 +60,12 @@ export * from "./helpers/query/QueryBuilderFactory";
 export * from "./helpers/query/NetworkStateQuery";
 export * from "./helpers/query/NetworkStateTransportModule";
 export * from "./state/prefilled/PreFilledStateService";
-export * from "./state/prefilled/PreFilledWitnessProvider";
 export * from "./state/async/AsyncMerkleTreeStore";
 export * from "./state/async/AsyncStateService";
 export * from "./state/merkle/CachedMerkleTreeStore";
 export * from "./state/merkle/SyncCachedMerkleTreeStore";
 export * from "./state/state/DummyStateService";
 export * from "./state/state/CachedStateService";
-export * from "./state/MerkleStoreWitnessProvider";
 export * from "./settlement/SettlementModule";
 export * from "./settlement/messages/WithdrawalQueue";
 export * from "./settlement/messages/IncomingMessageAdapter";

packages/protocol/test/TestingProtocol.ts

@@ -18,7 +18,6 @@ import {
   ProofTaskSerializer,
 } from "../../../helpers/utils";
 import { TaskWorkerModule } from "../../../worker/worker/TaskWorkerModule";
-import { PreFilledWitnessProvider } from "../../../state/prefilled/PreFilledWitnessProvider";
 
 import {
   StateTransitionParametersSerializer,
@@ -61,13 +60,8 @@ export class StateTransitionTask
   public async compute(
     input: StateTransitionProofParameters
   ): Promise<StateTransitionProof> {
-    const witnessProvider = new PreFilledWitnessProvider(input.merkleWitnesses);
-
-    const { witnessProviderReference } = this.stateTransitionProver;
-    const previousProvider = witnessProviderReference.getWitnessProvider();
-    witnessProviderReference.setWitnessProvider(witnessProvider);
-
     const stBatch = input.stateTransitions.slice();
+    const merkleWitnesses = input.merkleWitnesses.slice();
     // Array.from({
     //   length: ProtocolConstants.stateTransitionProverBatchSize - stBatch.length,
     // }).forEach(() => {
@@ -78,21 +72,16 @@ export class StateTransitionTask
 
     const output = await this.stateTransitionProver.runBatch(
       input.publicInput,
-      StateTransitionProvableBatch.fromMappings(stBatch)
+      StateTransitionProvableBatch.fromMappings(stBatch, merkleWitnesses)
     );
     log.debug("STTask public io:", {
       input: StateTransitionProverPublicInput.toJSON(input.publicInput),
       output: StateTransitionProverPublicOutput.toJSON(output),
     });
 
-    const proof = await this.executionContext
+    return await this.executionContext
       .current()
       .result.prove<StateTransitionProof>();
-
-    if (previousProvider !== undefined) {
-      witnessProviderReference.setWitnessProvider(previousProvider);
-    }
-    return proof;
   }
 
   public async prepare(): Promise<void> {