Fix state carrying issue with block prover

packages/common/src/trees/sparse/RollupMerkleTree.ts

@@ -189,6 +189,8 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
       return hash;
     }
 
+    // TODO Make sure this implementation is as efficient as it gets.
+    //  Especially compared to doing calculateRoot + witness new witness + check index
     public calculateRootIncrement(
       leafIndex: Field,
       leaf: Field

packages/protocol/src/prover/accumulators/WitnessedRootHashList.ts

@@ -1,4 +1,4 @@
-import { Bool, Field, Struct } from "o1js";
+import { Bool, Field, Provable, Struct } from "o1js";
 
 import { DefaultProvableHashList } from "../../utils/ProvableHashList";
 
@@ -13,15 +13,17 @@ export class WitnessedRoot extends Struct({
 
 export class WitnessedRootWitness extends Struct({
   witnessedRoot: Field,
-  preimage: Field,
 }) {}
 
 /**
  * Accumulator as of section "Intermediary state roots" of the STProver v2 spec
  */
 
 export class WitnessedRootHashList extends DefaultProvableHashList<WitnessedRoot> {
-  public constructor(commitment: Field = Field(0)) {
+  public constructor(
+    commitment: Field = Field(0),
+    public preimage: Field = Field(0)
+  ) {
     super(WitnessedRoot, commitment);
   }
 
@@ -34,20 +36,14 @@ export class WitnessedRootHashList extends DefaultProvableHashList<WitnessedRoot
    * points fall back to the same ST (because any batches in between were empty),
    * this has to be detected and compensated for.
    * This function does this using the preimage of the current list state.
-   *
-   * @param preimage The preimage to the **current** state of the list.
    */
-  public witnessRoot(
-    witnessedRoot: WitnessedRoot,
-    preimage: Field,
-    condition: Bool
-  ) {
+  public witnessRoot(witnessedRoot: WitnessedRoot, condition: Bool) {
     // Note, we don't have to validate the preimage here because of the following
     // 1. If the sequencer doesn't provide the correct witness, the BlockProver's
     //    equality check will fail
     // 2. If the list is empty, no preimage exists, therefore condition (2) doesn't
     //    apply, which is the same outcome when the sequencer provides an arbitrary witness
-    const preimageCheckList = new WitnessedRootHashList(preimage).push(
+    const preimageCheckList = new WitnessedRootHashList(this.preimage).push(
       witnessedRoot
     );
 
@@ -56,6 +52,11 @@ export class WitnessedRootHashList extends DefaultProvableHashList<WitnessedRoot
     // (2) don't append if preimage.push({ finalizedRoot, pendingSTBatchesHash }) == this.commitment
     const skipPush = preimageCheckList.commitment.equals(this.commitment);
 
-    return this.pushIf(witnessedRoot, condition.and(skipPush.not()));
+    const fromCommitment = this.commitment;
+
+    const pushCondition = condition.and(skipPush.not());
+    this.pushIf(witnessedRoot, pushCondition);
+
+    this.preimage = Provable.if(pushCondition, fromCommitment, this.preimage);
   }
 }

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

@@ -1,4 +1,4 @@
-import { Bool, Field, Proof, Provable, Struct } from "o1js";
+import { Bool, Field, Poseidon, Proof, Provable, Struct } from "o1js";
 import { CompilableModule, WithZkProgrammable } from "@proto-kit/common";
 
 import { StateTransitionProof } from "../statetransition/StateTransitionProvable";
@@ -32,7 +32,6 @@ export class BlockArguments extends Struct({
     return new BlockArguments({
       afterBlockRootWitness: {
         witnessedRoot: stateRoot,
-        preimage: Field(0),
       },
       transactionsHash: Field(0),
       pendingSTBatchesHash: {
@@ -52,6 +51,92 @@ export class BlockArgumentsBatch extends Struct({
   batch: Provable.Array(BlockArguments, BLOCK_ARGUMENT_BATCH_SIZE),
 }) {}
 
+const BlockProverStateBaseFields = {
+  eternalTransactionsHash: Field,
+  incomingMessagesHash: Field,
+  stateRoot: Field,
+  blockHashRoot: Field,
+  blockNumber: Field,
+  networkStateHash: Field,
+};
+
+export class BlockProverPublicInput extends Struct({
+  // Tracker of the current block prover state
+  proverStateRemainder: Field,
+  ...BlockProverStateBaseFields,
+}) {
+  public equals(input: BlockProverPublicInput): Bool {
+    const output2 = BlockProverPublicInput.toFields(input);
+    const output1 = BlockProverPublicInput.toFields(this);
+    return output1
+      .map((value1, index) => value1.equals(output2[index]))
+      .reduce((a, b) => a.and(b));
+  }
+
+  public clone() {
+    return new BlockProverPublicInput(
+      BlockProverPublicInput.fromFields(BlockProverPublicInput.toFields(this))
+    );
+  }
+}
+
+export const BlockProverStateCommitments = {
+  remainders: {
+    // Commitment to the list of unprocessed (pending) batches of STs that need to be proven
+    pendingSTBatchesHash: Field,
+    witnessedRootsHash: Field,
+    bundlesHash: Field,
+    witnessedRootsPreimage: Field,
+  },
+  ...BlockProverStateBaseFields,
+};
+
+export class BlockProverStateInput extends Struct(BlockProverStateCommitments) {
+  public hash() {
+    return Poseidon.hash(BlockProverStateInput.toFields(this));
+  }
+
+  public static fromPublicInput(input: BlockProverPublicInput) {
+    return new BlockProverStateInput({
+      remainders: {
+        bundlesHash: Field(0),
+        pendingSTBatchesHash: Field(0),
+        witnessedRootsHash: Field(0),
+        witnessedRootsPreimage: Field(0),
+      },
+      eternalTransactionsHash: input.eternalTransactionsHash,
+      incomingMessagesHash: input.incomingMessagesHash,
+      stateRoot: input.stateRoot,
+      blockHashRoot: input.blockHashRoot,
+      blockNumber: input.blockNumber,
+      networkStateHash: input.networkStateHash,
+    });
+  }
+
+  public finalize(condition: Bool) {
+    condition
+      .implies(
+        this.remainders.bundlesHash
+          .equals(0)
+          .and(this.remainders.pendingSTBatchesHash.equals(0))
+          .and(this.remainders.witnessedRootsHash.equals(0))
+      )
+      .assertTrue("Remainers not fully removed");
+
+    return new BlockProverPublicInput({
+      proverStateRemainder: Field(0),
+      eternalTransactionsHash: this.eternalTransactionsHash,
+      incomingMessagesHash: this.incomingMessagesHash,
+      stateRoot: this.stateRoot,
+      blockHashRoot: this.blockHashRoot,
+      blockNumber: this.blockNumber,
+      networkStateHash: this.networkStateHash,
+    });
+  }
+}
+
+export class BlockProverPublicOutput extends BlockProverPublicInput {}
+
 export class BlockProverState {
   /**
    * The network state which gives access to values such as blockHeight
@@ -113,44 +198,46 @@ export class BlockProverState {
     this.incomingMessages = args.incomingMessages;
   }
 
-  public toCommitments(): BlockProverPublicInput {
-    return {
+  public toCommitments(): BlockProverStateInput {
+    return new BlockProverStateInput({
       remainders: {
         bundlesHash: this.bundleList.commitment,
         pendingSTBatchesHash: this.pendingSTBatches.commitment,
         witnessedRootsHash: this.witnessedRoots.commitment,
+        witnessedRootsPreimage: this.witnessedRoots.preimage,
       },
       eternalTransactionsHash: this.eternalTransactionsList.commitment,
       incomingMessagesHash: this.incomingMessages.commitment,
       stateRoot: this.stateRoot,
       blockHashRoot: this.blockHashRoot,
       blockNumber: this.blockNumber,
       networkStateHash: this.networkState.hash(),
-    };
+    });
   }
 
   public static blockProverFromCommitments(
-    publicInput: BlockProverPublicInput,
+    stateInput: NonMethods<BlockProverStateInput>,
     networkState: NetworkState,
     blockWitness: BlockHashMerkleTreeWitness
   ): BlockProverState {
     return new BlockProverState({
-      bundleList: new BundleHashList(publicInput.remainders.bundlesHash),
+      bundleList: new BundleHashList(stateInput.remainders.bundlesHash),
       eternalTransactionsList: new TransactionHashList(
-        publicInput.eternalTransactionsHash
+        stateInput.eternalTransactionsHash
       ),
       incomingMessages: new MinaActionsHashList(
-        publicInput.incomingMessagesHash
+        stateInput.incomingMessagesHash
       ),
       pendingSTBatches: new AppliedBatchHashList(
-        publicInput.remainders.pendingSTBatchesHash
+        stateInput.remainders.pendingSTBatchesHash
       ),
       witnessedRoots: new WitnessedRootHashList(
-        publicInput.remainders.witnessedRootsHash
+        stateInput.remainders.witnessedRootsHash,
+        stateInput.remainders.witnessedRootsPreimage
       ),
-      stateRoot: publicInput.stateRoot,
-      blockHashRoot: publicInput.blockHashRoot,
-      blockNumber: publicInput.blockNumber,
+      stateRoot: stateInput.stateRoot,
+      blockHashRoot: stateInput.blockHashRoot,
+      blockNumber: stateInput.blockNumber,
       networkState,
       blockWitness,
     });
@@ -167,6 +254,7 @@ export class BlockProverState {
       this.pendingSTBatches.commitment,
       this.incomingMessages.commitment,
       this.witnessedRoots.commitment,
+      this.witnessedRoots.preimage,
       this.stateRoot,
       this.blockHashRoot,
       this.blockNumber,
@@ -182,14 +270,14 @@ export class BlockProverState {
       eternalTransactionsList: new TransactionHashList(fields[1]),
       pendingSTBatches: new AppliedBatchHashList(fields[2]),
       incomingMessages: new MinaActionsHashList(fields[3]),
-      witnessedRoots: new WitnessedRootHashList(fields[4]),
-      stateRoot: fields[5],
-      blockHashRoot: fields[6],
-      blockNumber: fields[7],
-      networkState: new NetworkState(NetworkState.fromFields(fields.slice(8))),
+      witnessedRoots: new WitnessedRootHashList(fields[4], fields[5]),
+      stateRoot: fields[6],
+      blockHashRoot: fields[7],
+      blockNumber: fields[8],
+      networkState: new NetworkState(NetworkState.fromFields(fields.slice(9))),
       blockWitness: new BlockHashMerkleTreeWitness(
         BlockHashMerkleTreeWitness.fromFields(
-          fields.slice(8 + NetworkState.sizeInFields())
+          fields.slice(9 + NetworkState.sizeInFields())
         )
       ),
     });
@@ -230,6 +318,11 @@ export class BlockProverState {
           condition,
           a.witnessedRoots.commitment,
           b.witnessedRoots.commitment
+        ),
+        Provable.if(
+          condition,
+          a.witnessedRoots.preimage,
+          b.witnessedRoots.preimage
         )
       ),
       stateRoot: Provable.if(condition, a.stateRoot, b.stateRoot),
@@ -250,51 +343,30 @@ export class BlockProverState {
   }
 }
 
-export const BlockProverStateCommitments = {
-  remainders: {
-    // Commitment to the list of unprocessed (pending) batches of STs that need to be proven
-    pendingSTBatchesHash: Field,
-    witnessedRootsHash: Field,
-    bundlesHash: Field,
-  },
-  eternalTransactionsHash: Field,
-  incomingMessagesHash: Field,
-  stateRoot: Field,
-  blockHashRoot: Field,
-  blockNumber: Field,
-  networkStateHash: Field,
-};
-
-export class BlockProverPublicInput extends Struct(
-  BlockProverStateCommitments
-) {}
-
-export class BlockProverPublicOutput extends Struct({
-  ...BlockProverStateCommitments,
-}) {
-  public equals(input: BlockProverPublicInput): Bool {
-    const output2 = BlockProverPublicOutput.toFields(input);
-    const output1 = BlockProverPublicOutput.toFields(this);
-    return output1
-      .map((value1, index) => value1.equals(output2[index]))
-      .reduce((a, b) => a.and(b));
-  }
-}
-
 export type BlockProof = Proof<BlockProverPublicInput, BlockProverPublicOutput>;
 
 export interface BlockProvable
   extends WithZkProgrammable<BlockProverPublicInput, BlockProverPublicOutput>,
     CompilableModule {
-  proveBlockBatch: (
+  proveBlockBatchNoProofs: (
     publicInput: BlockProverPublicInput,
+    stateWitness: BlockProverStateInput,
     networkState: NetworkState,
     blockWitness: BlockHashMerkleTreeWitness,
-    stateTransitionProof: StateTransitionProof,
+    batch: BlockArgumentsBatch,
+    finalize: Bool
+  ) => Promise<BlockProverPublicOutput>;
+
+  proveBlockBatchWithProofs: (
+    publicInput: BlockProverPublicInput,
+    stateWitness: BlockProverStateInput,
+    networkState: NetworkState,
+    blockWitness: BlockHashMerkleTreeWitness,
+    batch: BlockArgumentsBatch,
     deferSTProof: Bool,
-    transactionProof: TransactionProof,
     deferTransactionProof: Bool,
-    batch: BlockArgumentsBatch
+    stateTransitionProof: StateTransitionProof,
+    transactionProof: TransactionProof
   ) => Promise<BlockProverPublicOutput>;
 
   merge: (