@proto-kit/protocol 0.1.1-develop.191 → 0.1.1-develop.2024

package/src/prover/statetransition/StateTransitionProver.ts

@@ -1,49 +1,51 @@
-import { Experimental, Field, Provable, SelfProof } from "snarkyjs";
+import {
+  AreProofsEnabled,
+  PlainZkProgram,
+  provableMethod,
+  ZkProgrammable,
+  CompilableModule,
+  type ArtifactRecord,
+  type CompileRegistry,
+  TreeWrite,
+  LinkedMerkleTreeCircuitOps,
+  LinkedMerkleTreeWitness,
+} from "@proto-kit/common";
+import { Field, Provable, SelfProof, ZkProgram } from "o1js";
 import { injectable } from "tsyringe";
-import { PlainZkProgram, provableMethod } from "@proto-kit/common";
 
-import {
-  MerkleTreeUtils,
-  RollupMerkleWitness,
-} from "../../utils/merkletree/RollupMerkleTree.js";
-import {
-  DefaultProvableHashList,
-  ProvableHashList,
-} from "../../utils/ProvableHashList";
-import { ProvableStateTransition } from "../../model/StateTransition";
-import { StateTransitionProvableBatch } from "../../model/StateTransitionProvableBatch";
 import { constants } from "../../Constants";
+import { ProvableStateTransition } from "../../model/StateTransition";
+import {
+  MerkleWitnessBatch,
+  StateTransitionProvableBatch,
+  StateTransitionType,
+} from "../../model/StateTransitionProvableBatch";
 import { ProtocolModule } from "../../protocol/ProtocolModule";
+import { DefaultProvableHashList } from "../../utils/ProvableHashList";
+import { WitnessedRootHashList } from "../accumulators/WitnessedRootHashList";
+import { AppliedBatchHashList } from "../accumulators/AppliedBatchHashList";
+import { AppliedStateTransitionBatchState } from "../../model/AppliedStateTransitionBatch";
 
-import { StateTransitionWitnessProvider } from "./StateTransitionWitnessProvider.js";
 import {
+  StateTransitionProof,
   StateTransitionProvable,
   StateTransitionProverPublicInput,
-  StateTransitionProof,
   StateTransitionProverPublicOutput,
 } from "./StateTransitionProvable";
-import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
 
 const errors = {
-  stateRootNotMatching: (step: string) => `StateRoots not matching ${step}`,
-
-  stateTransitionsHashNotMatching: (step: string) =>
-    `State transitions hash not matching ${step}`,
+  propertyNotMatching: (property: string, step: string) =>
+    `${property} not matching ${step}`,
 
   merkleWitnessNotCorrect: (index: number) =>
     `MerkleWitness not valid for StateTransition (${index})`,
-
-  noWitnessProviderSet: () =>
-    new Error(
-      "WitnessProvider not set, set it before you use StateTransitionProvider"
-    ),
-
-  propertyNotMatching: (propertyName: string) => `${propertyName} not matching`,
 };
 
 interface StateTransitionProverExecutionState {
-  stateRoot: Field;
-  stateTransitionList: ProvableHashList<ProvableStateTransition>;
+  currentBatch: AppliedStateTransitionBatchState;
+  batchList: AppliedBatchHashList;
+  finalizedRoot: Field;
+  witnessedRoots: WitnessedRootHashList;
 }
 
 const StateTransitionSelfProofClass = SelfProof<
@@ -55,42 +57,56 @@ const StateTransitionSelfProofClass = SelfProof<
  * StateTransitionProver is the prover that proves the application of some state
  * transitions and checks and updates their merkle-tree entries
  */
-@injectable()
-export class StateTransitionProver
-  extends ProtocolModule<
-    StateTransitionProverPublicInput,
-    StateTransitionProverPublicOutput
-  >
-  implements StateTransitionProvable
-{
+export class StateTransitionProverProgrammable extends ZkProgrammable<
+  StateTransitionProverPublicInput,
+  StateTransitionProverPublicOutput
+> {
   public constructor(
-    // Injected
-    public readonly witnessProviderReference: StateTransitionWitnessProviderReference
+    private readonly stateTransitionProver: Pick<
+      StateTransitionProver,
+      "areProofsEnabled"
+    >
   ) {
     super();
   }
 
+  public get areProofsEnabled(): AreProofsEnabled | undefined {
+    return this.stateTransitionProver.areProofsEnabled;
+  }
+
   public zkProgramFactory(): PlainZkProgram<
     StateTransitionProverPublicInput,
     StateTransitionProverPublicOutput
-  > {
-    // eslint-disable-next-line max-len
-    // eslint-disable-next-line @typescript-eslint/no-this-alias,consistent-this,unicorn/no-this-assignment
+  >[] {
     const instance = this;
 
-    const program = Experimental.ZkProgram({
+    const program = ZkProgram({
+      name: "StateTransitionProver",
       publicInput: StateTransitionProverPublicInput,
       publicOutput: StateTransitionProverPublicOutput,
 
       methods: {
         proveBatch: {
-          privateInputs: [StateTransitionProvableBatch],
+          privateInputs: [
+            StateTransitionProvableBatch,
+            MerkleWitnessBatch,
+            AppliedStateTransitionBatchState,
+          ],
 
-          method(
+          async method(
             publicInput: StateTransitionProverPublicInput,
-            batch: StateTransitionProvableBatch
+            batch: StateTransitionProvableBatch,
+            witnesses: MerkleWitnessBatch,
+            currentAppliedBatch: AppliedStateTransitionBatchState
           ) {
-            return instance.runBatch(publicInput, batch);
+            return {
+              publicOutput: await instance.proveBatch(
+                publicInput,
+                batch,
+                witnesses,
+                currentAppliedBatch
+              ),
+            };
           },
         },
 
@@ -100,12 +116,14 @@ export class StateTransitionProver
             StateTransitionSelfProofClass,
           ],
 
-          method(
+          async method(
             publicInput: StateTransitionProverPublicInput,
             proof1: StateTransitionProof,
             proof2: StateTransitionProof
           ) {
-            return instance.merge(publicInput, proof1, proof2);
+            return {
+              publicOutput: await instance.merge(publicInput, proof1, proof2),
+            };
           },
         },
       },
@@ -116,22 +134,39 @@ export class StateTransitionProver
       merge: program.merge.bind(program),
     };
 
-    const SelfProofClass = Experimental.ZkProgram.Proof(program);
+    const SelfProofClass = ZkProgram.Proof(program);
 
-    return {
-      compile: program.compile.bind(program),
-      verify: program.verify.bind(program),
-      Proof: SelfProofClass,
-      methods,
-    };
+    return [
+      {
+        name: program.name,
+        compile: program.compile.bind(program),
+        verify: program.verify.bind(program),
+        analyzeMethods: program.analyzeMethods.bind(program),
+        Proof: SelfProofClass,
+        methods,
+      },
+    ];
   }
 
-  private get witnessProvider(): StateTransitionWitnessProvider {
-    const provider = this.witnessProviderReference.getWitnessProvider();
-    if (provider === undefined) {
-      throw errors.noWitnessProviderSet();
-    }
-    return provider;
+  private transitionToTreeWrite(
+    st: ProvableStateTransition,
+    witness: LinkedMerkleTreeWitness
+  ): TreeWrite {
+    // If from isSome isn't set, the user "ignored the previous value", i.e. we
+    // can assume the value in the witness is correct
+    const from = Provable.if(
+      st.from.isSome,
+      st.from.value,
+      witness.leafCurrent.leaf.value
+    );
+
+    // If the user doesn't want to write, we just carry over the from-value
+    const to = Provable.if(st.to.isSome, st.to.value, from);
+    return {
+      path: st.path,
+      from,
+      to,
+    };
   }
 
   /**
@@ -139,26 +174,94 @@ export class StateTransitionProver
    * and returns the new prover state
    */
   public applyTransitions(
-    stateRoot: Field,
-    stateTransitionCommitmentFrom: Field,
-    transitionBatch: StateTransitionProvableBatch
-  ): StateTransitionProverExecutionState {
-    const state: StateTransitionProverExecutionState = {
-      stateRoot,
-
-      stateTransitionList: new DefaultProvableHashList(
-        ProvableStateTransition,
-        stateTransitionCommitmentFrom
-      ),
-    };
+    state: StateTransitionProverExecutionState,
+    batch: StateTransitionProvableBatch,
+    witnesses: MerkleWitnessBatch
+  ) {
+    const transitions = batch.batch;
 
-    const transitions = transitionBatch.batch;
     for (
       let index = 0;
       index < constants.stateTransitionProverBatchSize;
       index++
     ) {
-      this.applyTransition(state, transitions[index], index);
+      const updatedBatchState = this.applyTransition(
+        state.currentBatch,
+        transitions[index].stateTransition,
+        witnesses.witnesses[index],
+        index
+      );
+
+      // If the current batch is finished, we push it to the list
+      // and initialize the next
+      const { type, witnessRoot } = transitions[index];
+      const closing = type.isClosing();
+      const closingAndApply = type.type.equals(
+        StateTransitionType.closeAndApply
+      );
+
+      // Create the newBatch
+      // The root is based on if the previous batch will be applied or not
+      const base = Provable.if(
+        closingAndApply,
+        updatedBatchState.root,
+        state.finalizedRoot
+      );
+      const newBatchState = new AppliedStateTransitionBatchState({
+        batchHash: Field(0),
+        root: base,
+      });
+
+      const updatedBatch = {
+        applied: closingAndApply,
+        batchHash: updatedBatchState.batchHash,
+      };
+      state.batchList.pushIf(updatedBatch, closing);
+      state.finalizedRoot = Provable.if(
+        closingAndApply,
+        updatedBatchState.root,
+        state.finalizedRoot
+      );
+
+      // Add computed root to the witnessed root list if needed
+      witnessRoot
+        .implies(closing)
+        .assertTrue("Can only witness roots at closing batches");
+      state.witnessedRoots.pushIf(
+        {
+          root: state.finalizedRoot,
+          appliedBatchListState: state.batchList.commitment,
+        },
+        witnessRoot
+      );
+
+      const isDummy = ProvableStateTransition.isDummy(
+        transitions[index].stateTransition
+      );
+
+      // Dummy STs cannot change any state, as to prevent any
+      // dummy-in-the-middle attacks. This is given if the type is nothing.
+      isDummy
+        .implies(type.isNothing())
+        .assertTrue("Dummies have to be of type 'nothing'");
+
+      isDummy
+        .implies(state.currentBatch.batchHash.equals(0))
+        .assertTrue("Dummies can only be placed on closed batchLists");
+
+      // Dummies don't close the batch, but we still want to ignore any
+      // updated batch, since we need to result to stay.
+      // This will break the pipeline if there is a dummy in the middle,
+      // but will only end up to invalid proofs (i.e. mismatched batches)
+
+      state.currentBatch = new AppliedStateTransitionBatchState(
+        Provable.if(
+          closing.or(isDummy),
+          AppliedStateTransitionBatchState,
+          newBatchState,
+          updatedBatchState
+        )
+      );
     }
 
     return state;
@@ -169,37 +272,45 @@ export class StateTransitionProver
    * and mutates it in place
    */
   public applyTransition(
-    state: StateTransitionProverExecutionState,
+    currentBatch: AppliedStateTransitionBatchState,
     transition: ProvableStateTransition,
+    witness: LinkedMerkleTreeWitness,
     index = 0
   ) {
-    const treeWitness = Provable.witness(RollupMerkleWitness, () =>
-      this.witnessProvider.getWitness(transition.path)
+    const impliedRoot = this.applyTransitionToRoot(
+      transition,
+      currentBatch.root,
+      witness,
+      index
     );
 
-    const membershipValid = MerkleTreeUtils.checkMembership(
-      treeWitness,
-      state.stateRoot,
-      transition.path,
-      transition.from.value
+    // Append ST to the current batch's ST-list
+    const stList = new DefaultProvableHashList(
+      ProvableStateTransition,
+      currentBatch.batchHash
     );
-    membershipValid
-      .or(transition.from.isSome.not())
-      .assertTrue(errors.merkleWitnessNotCorrect(index));
+    stList.push(transition);
 
-    const newRoot = MerkleTreeUtils.computeRoot(
-      treeWitness,
-      transition.to.value
-    );
-    state.stateRoot = Provable.if(
-      transition.to.isSome,
-      newRoot,
-      state.stateRoot
-    );
+    // Update batch
+    return new AppliedStateTransitionBatchState({
+      batchHash: stList.commitment,
+      root: impliedRoot,
+    });
+  }
 
-    state.stateTransitionList.pushIf(
-      transition,
-      transition.path.equals(Field(0)).not()
+  private applyTransitionToRoot(
+    transition: ProvableStateTransition,
+    root: Field,
+    merkleWitness: LinkedMerkleTreeWitness,
+    index: number
+  ): Field {
+    const treeWrite = this.transitionToTreeWrite(transition, merkleWitness);
+
+    return LinkedMerkleTreeCircuitOps.applyTreeWrite(
+      root,
+      merkleWitness,
+      treeWrite,
+      index
     );
   }
 
@@ -207,54 +318,154 @@ export class StateTransitionProver
    * Applies a whole batch of StateTransitions at once
    */
   @provableMethod()
-  public runBatch(
+  public async proveBatch(
     publicInput: StateTransitionProverPublicInput,
-    batch: StateTransitionProvableBatch
-  ): StateTransitionProverPublicOutput {
-    const result = this.applyTransitions(
-      publicInput.stateRoot,
-      publicInput.stateTransitionsHash,
-      batch
-    );
+    batch: StateTransitionProvableBatch,
+    witnesses: MerkleWitnessBatch,
+    currentAppliedBatch: AppliedStateTransitionBatchState
+  ): Promise<StateTransitionProverPublicOutput> {
+    currentAppliedBatch
+      .hashOrZero()
+      .assertEquals(
+        publicInput.currentBatchStateHash,
+        "Provided startingAppliedBatch not matching PI hash"
+      );
+
+    // Assert that either the currentAppliedBatch is somewhere intermediary
+    // or the root is the current "finalized" root
+    currentAppliedBatch.root
+      .equals(publicInput.root)
+      .or(publicInput.currentBatchStateHash.equals(0).not())
+      .assertTrue();
+
+    const state: StateTransitionProverExecutionState = {
+      batchList: new AppliedBatchHashList(publicInput.batchesHash),
+      currentBatch: currentAppliedBatch,
+      finalizedRoot: publicInput.root,
+      witnessedRoots: new WitnessedRootHashList(publicInput.witnessedRootsHash),
+    };
+
+    const result = this.applyTransitions(state, batch, witnesses);
 
     return new StateTransitionProverPublicOutput({
-      stateRoot: result.stateRoot,
-      stateTransitionsHash: result.stateTransitionList.commitment,
+      batchesHash: result.batchList.commitment,
+      currentBatchStateHash: result.currentBatch.hashOrZero(),
+      root: result.finalizedRoot,
+      witnessedRootsHash: result.witnessedRoots.commitment,
     });
   }
 
   @provableMethod()
-  public merge(
+  public async merge(
     publicInput: StateTransitionProverPublicInput,
     proof1: StateTransitionProof,
     proof2: StateTransitionProof
-  ): StateTransitionProverPublicOutput {
+  ): Promise<StateTransitionProverPublicOutput> {
     proof1.verify();
     proof2.verify();
 
-    // Check state
-    publicInput.stateRoot.assertEquals(
-      proof1.publicInput.stateRoot,
-      errors.stateRootNotMatching("publicInput.from -> proof1.from")
+    // Check current batch hash
+    publicInput.currentBatchStateHash.assertEquals(
+      proof1.publicInput.currentBatchStateHash,
+      errors.propertyNotMatching(
+        "currentBatchStateHash",
+        "publicInput.from -> proof1.from"
+      )
     );
-    proof1.publicOutput.stateRoot.assertEquals(
-      proof2.publicInput.stateRoot,
-      errors.stateRootNotMatching("proof1.to -> proof2.from")
+    proof1.publicOutput.currentBatchStateHash.assertEquals(
+      proof2.publicInput.currentBatchStateHash,
+      errors.propertyNotMatching(
+        "currentBatchStateHash",
+        "proof1.to -> proof2.from"
+      )
     );
 
-    // Check ST list
-    publicInput.stateTransitionsHash.assertEquals(
-      proof1.publicInput.stateTransitionsHash,
-      errors.stateTransitionsHashNotMatching("publicInput.from -> proof1.from")
+    // Check batches hash
+    publicInput.batchesHash.assertEquals(
+      proof1.publicInput.batchesHash,
+      errors.propertyNotMatching(
+        "batchesHash",
+        "publicInput.from -> proof1.from"
+      )
+    );
+    proof1.publicOutput.batchesHash.assertEquals(
+      proof2.publicInput.batchesHash,
+      errors.propertyNotMatching("batchesHash", "proof1.to -> proof2.from")
     );
-    proof1.publicOutput.stateTransitionsHash.assertEquals(
-      proof2.publicInput.stateTransitionsHash,
-      errors.stateTransitionsHashNotMatching("proof1.to -> proof2.from")
+
+    // Check root
+    publicInput.root.assertEquals(
+      proof1.publicInput.root,
+      errors.propertyNotMatching("root", "publicInput.from -> proof1.from")
+    );
+
+    proof1.publicOutput.root.assertEquals(
+      proof2.publicInput.root,
+      errors.propertyNotMatching("root", "proof1.to -> proof2.from")
+    );
+
+    // Check root accumulator
+    publicInput.witnessedRootsHash.assertEquals(
+      proof1.publicInput.witnessedRootsHash,
+      errors.propertyNotMatching(
+        "witnessedRootsHash",
+        "publicInput.from -> proof1.from"
+      )
+    );
+    proof1.publicOutput.witnessedRootsHash.assertEquals(
+      proof2.publicInput.witnessedRootsHash,
+      errors.propertyNotMatching(
+        "witnessedRootsHash",
+        "proof1.to -> proof2.from"
+      )
     );
 
     return new StateTransitionProverPublicInput({
-      stateRoot: proof2.publicOutput.stateRoot,
-      stateTransitionsHash: proof2.publicOutput.stateTransitionsHash,
+      currentBatchStateHash: proof2.publicOutput.currentBatchStateHash,
+      batchesHash: proof2.publicOutput.batchesHash,
+      root: proof2.publicOutput.root,
+      witnessedRootsHash: proof2.publicOutput.witnessedRootsHash,
     });
   }
 }
+
+@injectable()
+export class StateTransitionProver
+  extends ProtocolModule
+  implements StateTransitionProvable, CompilableModule
+{
+  public zkProgrammable: StateTransitionProverProgrammable;
+
+  public constructor() {
+    super();
+    this.zkProgrammable = new StateTransitionProverProgrammable(this);
+  }
+
+  public async compile(
+    registry: CompileRegistry
+  ): Promise<void | ArtifactRecord> {
+    return await this.zkProgrammable.compile(registry);
+  }
+
+  public proveBatch(
+    publicInput: StateTransitionProverPublicInput,
+    batch: StateTransitionProvableBatch,
+    witnesses: MerkleWitnessBatch,
+    startingAppliedBatch: AppliedStateTransitionBatchState
+  ): Promise<StateTransitionProverPublicOutput> {
+    return this.zkProgrammable.proveBatch(
+      publicInput,
+      batch,
+      witnesses,
+      startingAppliedBatch
+    );
+  }
+
+  public merge(
+    publicInput: StateTransitionProverPublicInput,
+    proof1: StateTransitionProof,
+    proof2: StateTransitionProof
+  ): Promise<StateTransitionProverPublicOutput> {
+    return this.zkProgrammable.merge(publicInput, proof1, proof2);
+  }
+}