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

package/src/prover/block/BlockProver.ts

@@ -1,233 +1,615 @@
-/* eslint-disable max-lines */
-import { Experimental, Field, type Proof, Provable, SelfProof } from "snarkyjs";
-import { inject, injectable } from "tsyringe";
+import { Bool, Field, Provable, SelfProof, ZkProgram } from "o1js";
+import { container, inject, injectable, injectAll } from "tsyringe";
 import {
+  AreProofsEnabled,
+  CompilableModule,
+  CompileArtifact,
+  CompileRegistry,
+  log,
+  NonMethods,
   PlainZkProgram,
   provableMethod,
+  reduceSequential,
   WithZkProgrammable,
   ZkProgrammable,
 } from "@proto-kit/common";
 
-import { DefaultProvableHashList } from "../../utils/ProvableHashList";
-import { MethodPublicOutput } from "../../model/MethodPublicOutput";
 import { ProtocolModule } from "../../protocol/ProtocolModule";
 import {
   StateTransitionProof,
+  StateTransitionProvable,
   StateTransitionProverPublicInput,
   StateTransitionProverPublicOutput,
 } from "../statetransition/StateTransitionProvable";
 import { RuntimeTransaction } from "../../model/transaction/RuntimeTransaction";
+import { RuntimeMethodExecutionContext } from "../../state/context/RuntimeMethodExecutionContext";
+import {
+  AfterBlockHookArguments,
+  BeforeBlockHookArguments,
+  ProvableBlockHook,
+  toAfterBlockHookArgument,
+  toBeforeBlockHookArgument,
+} from "../../protocol/ProvableBlockHook";
+import { NetworkState } from "../../model/network/NetworkState";
+import { assertEqualsIf } from "../../utils/utils";
+import { StateServiceProvider } from "../../state/StateServiceProvider";
+import { executeHooks } from "../utils";
+import {
+  TransactionProof,
+  TransactionProvable,
+  TransactionProverPublicInput,
+  TransactionProverPublicOutput,
+} from "../transaction/TransactionProvable";
+import { Bundle } from "../accumulators/BlockHashList";
 
 import {
+  BlockArguments,
+  BlockArgumentsBatch,
+  BlockProof,
   BlockProvable,
-  BlockProverExecutionData,
-  BlockProverProof,
   BlockProverPublicInput,
   BlockProverPublicOutput,
+  BlockProverState,
+  BlockProverStateInput,
 } from "./BlockProvable";
+import {
+  BlockHashMerkleTreeWitness,
+  BlockHashTreeEntry,
+} from "./accummulators/BlockHashMerkleTree";
 
 const errors = {
-  stateProofNotStartingAtZero: () =>
-    "StateProof not starting ST-commitment at zero",
-
-  stateTransitionsHashNotEqual: () =>
-    "StateTransition list commitments are not equal",
+  propertyNotMatchingStep: (propertyName: string, step: string) =>
+    `${propertyName} not matching: ${step}`,
 
   propertyNotMatching: (propertyName: string) => `${propertyName} not matching`,
-
-  stateRootNotMatching: (step: string) => `StateRoots not matching ${step}`,
-
-  transactionsHashNotMatching: (step: string) =>
-    `transactions hash not matching ${step}`,
 };
 
-export interface BlockProverState {
-  // The current state root of the block prover
-  stateRoot: Field;
+type BlockHookArgument<T extends "before" | "after"> = T extends "before"
+  ? BeforeBlockHookArguments
+  : AfterBlockHookArguments;
 
-  /**
-   * The current commitment of the transaction-list which
-   * will at the end equal the bundle hash
-   */
-  transactionsHash: Field;
-
-  /**
-   * The network state which gives access to values such as blockHeight
-   * This value is the same for the whole batch (L2 block)
-   */
-  networkStateHash: Field;
-}
-
-/**
- * BlockProver class, which aggregates a AppChainProof and
- * a StateTransitionProof into a single BlockProof, that can
- * then be merged to be committed to the base-layer contract
- */
-@injectable()
-export class BlockProver
-  extends ProtocolModule<BlockProverPublicInput, BlockProverPublicOutput>
-  implements BlockProvable
-{
+export class BlockProverProgrammable extends ZkProgrammable<
+  BlockProverPublicInput,
+  BlockProverPublicOutput
+> {
   public constructor(
-    @inject("StateTransitionProver")
-    private readonly stateTransitionProver: ZkProgrammable<
+    private readonly prover: BlockProver,
+    public readonly stateTransitionProver: ZkProgrammable<
       StateTransitionProverPublicInput,
       StateTransitionProverPublicOutput
     >,
-    @inject("Runtime")
-    private readonly runtime: WithZkProgrammable<void, MethodPublicOutput>
+    public readonly transactionProver: ZkProgrammable<
+      TransactionProverPublicInput,
+      TransactionProverPublicOutput
+    >,
+    private readonly blockHooks: ProvableBlockHook<unknown>[],
+    private readonly stateServiceProvider: StateServiceProvider
   ) {
     super();
   }
 
-  /**
-   * Applies and checks the two proofs and applies the corresponding state
-   * changes to the given state
-   *
-   * @param state The from-state of the BlockProver
-   * @param stateTransitionProof
-   * @param appProof
-   * @returns The new BlockProver-state to be used as public output
-   */
-  public applyTransaction(
-    state: BlockProverState,
-    stateTransitionProof: Proof<
-      StateTransitionProverPublicInput,
-      StateTransitionProverPublicOutput
-    >,
-    appProof: Proof<void, MethodPublicOutput>,
-    { transaction, networkState }: BlockProverExecutionData
-  ): BlockProverState {
-    appProof.verify();
-    stateTransitionProof.verify();
+  name = "BlockProver";
+
+  public get areProofsEnabled(): AreProofsEnabled | undefined {
+    return this.prover.areProofsEnabled;
+  }
 
-    const stateTo = { ...state };
+  public async executeBlockHooks<T extends "before" | "after">(
+    type: T,
+    hook: (
+      module: ProvableBlockHook<unknown>,
+      networkState: NetworkState,
+      args: BlockHookArgument<T>
+    ) => Promise<NetworkState>,
+    hookArguments: BlockHookArgument<T>,
+    inputNetworkState: NetworkState,
+    isDummy: Bool
+  ) {
+    const transaction = RuntimeTransaction.dummyTransaction();
+    const startingInputs = {
+      transaction,
+      networkState: inputNetworkState,
+    };
 
-    // eslint-disable-next-line no-warning-comments
-    // TODO Check methodId?
+    return await executeHooks(
+      startingInputs,
+      `${type}Block`,
+      async () => {
+        const executionContext = container.resolve(
+          RuntimeMethodExecutionContext
+        );
+
+        return await this.blockHooks.reduce<Promise<NetworkState>>(
+          async (networkStatePromise, blockHook) => {
+            const networkState = await networkStatePromise;
+
+            // Setup context for potential calls to runtime methods.
+            // With the special case that we set the new networkstate for every hook
+            // We also have to put in a dummy transaction for network.transaction
+            executionContext.setup({
+              transaction: RuntimeTransaction.dummyTransaction(),
+              networkState,
+            });
+
+            return await hook(blockHook, networkState, hookArguments);
+          },
+          Promise.resolve(inputNetworkState)
+        );
+      },
+      isDummy
+    );
+  }
 
-    // Checks for the stateTransitionProof and appProof matching
-    stateTransitionProof.publicInput.stateTransitionsHash.assertEquals(
+  public includeSTProof(
+    stateTransitionProof: StateTransitionProof,
+    apply: Bool,
+    stateRoot: Field,
+    pendingSTBatchesHash: Field,
+    witnessedRootsHash: Field
+  ): {
+    stateRoot: Field;
+    pendingSTBatchesHash: Field;
+    witnessedRootsHash: Field;
+  } {
+    assertEqualsIf(
+      stateTransitionProof.publicInput.currentBatchStateHash,
+      Field(0),
+      apply,
+      "State for STProof has to be empty at the start"
+    );
+    assertEqualsIf(
+      stateTransitionProof.publicOutput.currentBatchStateHash,
       Field(0),
-      errors.stateProofNotStartingAtZero()
+      apply,
+      "State for STProof has to be empty at the end"
     );
 
-    appProof.publicOutput.stateTransitionsHash.assertEquals(
-      stateTransitionProof.publicOutput.stateTransitionsHash,
-      errors.stateTransitionsHashNotEqual()
+    assertEqualsIf(
+      stateTransitionProof.publicInput.batchesHash,
+      Field(0),
+      apply,
+      "Batcheshash doesn't start at 0"
     );
 
-    // Apply state if status success
-    state.stateRoot.assertEquals(
-      stateTransitionProof.publicInput.stateRoot,
+    // Assert from state root
+    assertEqualsIf(
+      stateRoot,
+      stateTransitionProof.publicInput.root,
+      apply,
       errors.propertyNotMatching("from state root")
     );
-    stateTo.stateRoot = Provable.if(
-      appProof.publicOutput.status,
-      stateTransitionProof.publicOutput.stateRoot,
-      stateTransitionProof.publicInput.stateRoot
+
+    // Assert the stBatchesHash executed is the same
+    assertEqualsIf(
+      pendingSTBatchesHash,
+      stateTransitionProof.publicOutput.batchesHash,
+      apply,
+      "Pending STBatches are not the same that have been executed by the ST proof"
     );
 
-    // Check transaction signature
-    transaction
-      .validateSignature()
-      .assertTrue("Transaction signature not valid");
+    // Assert root Accumulator
+    assertEqualsIf(
+      Field(0),
+      stateTransitionProof.publicInput.witnessedRootsHash,
+      apply,
+      errors.propertyNotMatching("from state root")
+    );
+    // Assert the witnessedRootsHash created is the same
+    assertEqualsIf(
+      witnessedRootsHash,
+      stateTransitionProof.publicOutput.witnessedRootsHash,
+      apply,
+      "Root accumulator Commitment is not the same that have been executed by the ST proof"
+    );
+
+    // update root only if we didn't defer
+    const newRoot = Provable.if(
+      apply,
+      stateTransitionProof.publicOutput.root,
+      stateRoot
+    );
+    // Reset only if we didn't defer
+    const newBatchesHash = Provable.if(apply, Field(0), pendingSTBatchesHash);
+    const newWitnessedRootsHash = Provable.if(
+      apply,
+      Field(0),
+      witnessedRootsHash
+    );
+    return {
+      stateRoot: newRoot,
+      pendingSTBatchesHash: newBatchesHash,
+      witnessedRootsHash: newWitnessedRootsHash,
+    };
+  }
+
+  private verifySTProof(
+    state: BlockProverState,
+    stateTransitionProof: StateTransitionProof,
+    deferSTProof: Bool
+  ) {
+    // Verify ST Proof only if STs have been emitted,
+    // and we don't defer the verification of the STs
+    // otherwise we can input a dummy proof
+    const batchesEmpty = state.pendingSTBatches.commitment.equals(Field(0));
+    const verifyStProof = deferSTProof.not().and(batchesEmpty.not());
+    log.provable.debug("Verify STProof", verifyStProof);
+    stateTransitionProof.verifyIf(verifyStProof);
+
+    // Apply STProof if not deferred
+    const stateProofResult = this.includeSTProof(
+      stateTransitionProof,
+      verifyStProof,
+      state.stateRoot,
+      state.pendingSTBatches.commitment,
+      state.witnessedRoots.commitment
+    );
+    state.stateRoot = stateProofResult.stateRoot;
+    state.pendingSTBatches.commitment = stateProofResult.pendingSTBatchesHash;
+    state.witnessedRoots.commitment = stateProofResult.witnessedRootsHash;
+  }
+
+  private verifyTransactionProof(
+    state: BlockProverState,
+    transactionProof: TransactionProof,
+    deferTransactionProof: Bool
+  ) {
+    // Verify Transaction proof if it has at least 1 tx and it isn't deferred
+    const finalizeBlockProof = deferTransactionProof.not();
+    const verifyTransactionProof = finalizeBlockProof.and(
+      state.bundleList.isEmpty().not()
+    );
 
-    // Check if the methodId is correct
-    // to do
+    transactionProof.verifyIf(verifyTransactionProof);
+
+    // Fast-forward transaction trackers by the results of the aggregated transaction proof
+    // Implicitly, the 'from' values here are asserted against the publicInput, since the hashlists
+    // are created out of the public input
+    state.eternalTransactionsList.fastForwardIf(
+      {
+        from: transactionProof.publicInput.eternalTransactionsHash,
+        to: transactionProof.publicOutput.eternalTransactionsHash,
+      },
+      verifyTransactionProof,
+      "eternalTransactionsList"
+    );
+
+    state.incomingMessages.fastForwardIf(
+      {
+        from: transactionProof.publicInput.incomingMessagesHash,
+        to: transactionProof.publicOutput.incomingMessagesHash,
+      },
+      verifyTransactionProof,
+      "incomingMessages"
+    );
 
-    // Check transaction integrity against appProof
-    const blockTransactionHash =
-      RuntimeTransaction.fromProtocolTransaction(transaction).hash();
+    // Cancel out remainders for transaction proof
+    assertEqualsIf(
+      transactionProof.publicInput.bundlesHash,
+      Field(0),
+      verifyTransactionProof,
+      "TransactionProof has to start bundles at 0"
+    );
 
-    blockTransactionHash.assertEquals(
-      appProof.publicOutput.transactionHash,
-      "Transactions provided in AppProof and BlockProof do not match"
+    // Fast Backwards actually, but logic holds
+    state.bundleList.fastForwardIf(
+      {
+        from: transactionProof.publicOutput.bundlesHash,
+        to: state.bundleList.empty(),
+      },
+      verifyTransactionProof,
+      "bundles hash"
     );
+  }
 
-    // Check network state integrity against appProof
-    state.networkStateHash.assertEquals(
-      appProof.publicOutput.networkStateHash,
-      "Network state does not match state used in AppProof"
+  private parseState(
+    publicInput: BlockProverPublicInput,
+    stateWitness: BlockProverStateInput,
+    networkState: NetworkState,
+    blockWitness: BlockHashMerkleTreeWitness
+  ) {
+    const hasNoStateRemained = publicInput.proverStateRemainder.equals(0);
+
+    // If the state is supplied as a witness, we check that it is equals the PI's stateHash
+    stateWitness
+      .hash()
+      .equals(publicInput.proverStateRemainder)
+      .or(hasNoStateRemained)
+      .assertTrue("Input state witness is invalid");
+
+    const stateInputs = Provable.if(
+      hasNoStateRemained,
+      BlockProverStateInput,
+      BlockProverStateInput.fromPublicInput(publicInput),
+      stateWitness
     );
-    state.networkStateHash.assertEquals(
+
+    stateInputs.networkStateHash.assertEquals(
       networkState.hash(),
-      "Network state provided to BlockProver does not match the publicInput"
+      "Network state not valid"
     );
 
-    // Append tx to transaction list
-    const transactionList = new DefaultProvableHashList(
-      Field,
-      state.transactionsHash
+    const state = BlockProverState.blockProverFromCommitments(
+      stateInputs,
+      networkState,
+      blockWitness
     );
 
-    const { transactionHash } = appProof.publicOutput;
-    transactionList.push(transactionHash);
+    // Verify block witness validity
+    const blockIndex = blockWitness.calculateIndex();
+
+    blockIndex.assertEquals(stateInputs.blockNumber);
 
-    stateTo.transactionsHash = transactionList.commitment;
+    blockWitness
+      .calculateRoot(Field(0))
+      .assertEquals(
+        stateInputs.blockHashRoot,
+        "Supplied block hash witness not matching state root"
+      );
 
-    return stateTo;
+    return state;
   }
 
-  @provableMethod()
-  public proveTransaction(
+  private computeOutput(
     publicInput: BlockProverPublicInput,
-    stateProof: StateTransitionProof,
-    appProof: Proof<void, MethodPublicOutput>,
-    executionData: BlockProverExecutionData
-  ): BlockProverPublicOutput {
-    const state: BlockProverState = {
-      transactionsHash: publicInput.transactionsHash,
-      stateRoot: publicInput.stateRoot,
-      networkStateHash: publicInput.networkStateHash,
-    };
+    state: BlockProverState,
+    finalizeBlockProof: Bool
+  ) {
+    const finalizedOutput = state.toCommitments();
 
-    const stateTo = this.applyTransaction(
-      state,
-      stateProof,
-      appProof,
-      executionData
+    const deferredOutput = {
+      ...publicInput,
+    };
+    deferredOutput.proverStateRemainder = finalizedOutput.hash();
+
+    return new BlockProverPublicOutput(
+      Provable.if(
+        finalizeBlockProof,
+        BlockProverPublicOutput,
+        finalizedOutput.finalize(finalizeBlockProof),
+        deferredOutput
+      )
     );
+  }
 
-    return new BlockProverPublicOutput({
-      stateRoot: stateTo.stateRoot,
-      transactionsHash: stateTo.transactionsHash,
-    });
+  @provableMethod()
+  public async proveBlockBatchNoProofs(
+    publicInput: BlockProverPublicInput,
+    stateWitness: BlockProverStateInput,
+    networkState: NetworkState,
+    blockWitness: BlockHashMerkleTreeWitness,
+    batch: BlockArgumentsBatch,
+    finalize: Bool
+  ) {
+    return await this.proveBlockBatch(
+      false,
+      publicInput,
+      stateWitness,
+      networkState,
+      blockWitness,
+      batch,
+      Bool(true),
+      Bool(true),
+      finalize
+    );
   }
 
   @provableMethod()
-  public merge(
+  public async proveBlockBatchWithProofs(
     publicInput: BlockProverPublicInput,
-    proof1: BlockProverProof,
-    proof2: BlockProverProof
-  ): BlockProverPublicOutput {
-    proof1.verify();
-    proof2.verify();
+    stateWitness: BlockProverStateInput,
+    networkState: NetworkState,
+    blockWitness: BlockHashMerkleTreeWitness,
+    batch: BlockArgumentsBatch,
+    deferSTProof: Bool,
+    deferTransactionProof: Bool,
+    stateTransitionProof: StateTransitionProof,
+    transactionProof: TransactionProof
+  ) {
+    const finalize = deferTransactionProof.or(deferSTProof).not();
+
+    return await this.proveBlockBatch(
+      true,
+      publicInput,
+      stateWitness,
+      networkState,
+      blockWitness,
+      batch,
+      deferSTProof,
+      deferTransactionProof,
+      finalize,
+      stateTransitionProof,
+      transactionProof
+    );
+  }
 
-    // Check state
-    publicInput.stateRoot.assertEquals(
-      proof1.publicInput.stateRoot,
-      errors.stateRootNotMatching("publicInput.from -> proof1.from")
+  public async proveBlockBatch(
+    doProofVerification: boolean,
+    publicInput: BlockProverPublicInput,
+    stateWitness: BlockProverStateInput,
+    networkState: NetworkState,
+    blockWitness: BlockHashMerkleTreeWitness,
+    batch: BlockArgumentsBatch,
+    deferSTProof: Bool,
+    deferTransactionProof: Bool,
+    finalize: Bool,
+    stateTransitionProof?: StateTransitionProof,
+    transactionProof?: TransactionProof
+  ): Promise<BlockProverPublicOutput> {
+    let state = this.parseState(
+      publicInput,
+      stateWitness,
+      networkState,
+      blockWitness
+    );
+
+    // Prove blocks iteratively
+    state = await reduceSequential(
+      batch.batch,
+      async (current, block) => {
+        const result = await this.proveBlock(current.copy(), block);
+
+        this.stateServiceProvider.popCurrentStateService();
+
+        return BlockProverState.choose(block.isDummy, current, result);
+      },
+      state
     );
-    proof1.publicOutput.stateRoot.assertEquals(
-      proof2.publicInput.stateRoot,
-      errors.stateRootNotMatching("proof1.to -> proof2.from")
+
+    if (doProofVerification) {
+      this.verifyTransactionProof(
+        state,
+        transactionProof!,
+        deferTransactionProof
+      );
+      this.verifySTProof(state, stateTransitionProof!, deferSTProof);
+    }
+
+    return this.computeOutput(publicInput, state, finalize);
+  }
+
+  private async proveBlock(
+    state: BlockProverState,
+    args: BlockArguments
+  ): Promise<BlockProverState> {
+    const { networkState, blockWitness } = state;
+    const { afterBlockRootWitness, transactionsHash, isDummy } = args;
+
+    const startingPendingStBatches = state.pendingSTBatches.commitment;
+
+    // 1. Execute beforeBlock hooks
+    const beforeBlockArgs = toBeforeBlockHookArgument(state);
+    const beforeBlockResult = await this.executeBlockHooks(
+      "before",
+      async (module, networkStateArg, hookArgs) =>
+        await module.beforeBlock(networkStateArg, hookArgs),
+      beforeBlockArgs,
+      networkState,
+      isDummy
     );
 
-    // Check transaction list
-    publicInput.transactionsHash.assertEquals(
-      proof1.publicInput.transactionsHash,
-      errors.transactionsHashNotMatching("publicInput.from -> proof1.from")
+    state.pendingSTBatches.push(beforeBlockResult.batch);
+
+    // 2. "Apply" TX-type BlockProof
+    args.pendingSTBatchesHash.from.assertEquals(
+      state.pendingSTBatches.commitment
     );
-    proof1.publicOutput.transactionsHash.assertEquals(
-      proof2.publicInput.transactionsHash,
-      errors.transactionsHashNotMatching("proof1.to -> proof2.from")
+    args.witnessedRootsHash.from.assertEquals(state.witnessedRoots.commitment);
+    const isEmptyBlock = transactionsHash.equals(Field(0));
+    const isNotEmptyBlock = isEmptyBlock.not();
+
+    // Check & fast-forward the stBatchHashList to after all transactions appended
+    state.pendingSTBatches.fastForward(
+      args.pendingSTBatchesHash,
+      "Transaction proof doesn't start their STs after the beforeBlockHook"
+    );
+    // Same for witnessedRootsHash
+    state.witnessedRoots.fastForward(
+      args.witnessedRootsHash,
+      "Transaction proof doesn't start with correct witnessed roots hash"
     );
 
-    return new BlockProverPublicOutput({
-      stateRoot: proof2.publicOutput.stateRoot,
-      transactionsHash: proof2.publicOutput.transactionsHash,
+    // Add block to bundles list
+    const bundle = new Bundle({
+      transactionsHash: transactionsHash,
+      networkStateHash: beforeBlockResult.result.hash(),
+      pendingSTBatchesHash: args.pendingSTBatchesHash,
+      witnessedRootsHash: args.witnessedRootsHash,
     });
+    state.bundleList.pushIf(bundle, isNotEmptyBlock);
+
+    // 3.
+    // Calculate new block tree root and increment witness
+    // Blocknumber as the index here is already authenticated previously
+    const [root, newWitness] = blockWitness.calculateRootIncrement(
+      state.blockNumber,
+      new BlockHashTreeEntry({
+        block: {
+          index: state.blockNumber,
+          transactionListHash: transactionsHash,
+        },
+        closed: Bool(true),
+      }).hash()
+    );
+
+    state.blockHashRoot = root;
+    state.blockWitness = newWitness;
+
+    state.blockNumber = state.blockNumber.add(1);
+
+    // 4. Execute afterBlock hooks
+    // Witness root
+    const hasNoSTBatches = state.pendingSTBatches.commitment.equals(
+      startingPendingStBatches
+    );
+
+    // TODO Cover case when we witness root but pendingSTBatches is completely empty
+
+    state.witnessedRoots.witnessRoot(
+      {
+        appliedBatchListState: state.pendingSTBatches.commitment,
+        root: afterBlockRootWitness.witnessedRoot,
+      },
+      hasNoSTBatches.not()
+    );
+
+    // Switch state service to afterBlock one
+    this.stateServiceProvider.popCurrentStateService();
+
+    // Execute hooks
+    const afterBlockHookArgs = toAfterBlockHookArgument(
+      state,
+      afterBlockRootWitness.witnessedRoot,
+      transactionsHash
+    );
+    const afterBlockResult = await this.executeBlockHooks(
+      "after",
+      async (module, networkStateArg, hookArgs) =>
+        await module.afterBlock(networkStateArg, hookArgs),
+      {
+        ...afterBlockHookArgs,
+      },
+      beforeBlockResult.result,
+      isDummy
+    );
+
+    // Apply state and network state changes
+    state.pendingSTBatches.push(afterBlockResult.batch);
+    state.networkState = afterBlockResult.result;
+
+    return state;
+  }
+
+  @provableMethod()
+  public async merge(
+    publicInput: BlockProverPublicInput,
+    proof1: BlockProof,
+    proof2: BlockProof
+  ): Promise<BlockProverPublicOutput> {
+    proof1.verify();
+    proof2.verify();
+
+    function checkProperty<
+      Key extends keyof NonMethods<BlockProverPublicInput>,
+    >(key: Key) {
+      // Check state
+      publicInput[key].assertEquals(
+        proof1.publicInput[key],
+        errors.propertyNotMatchingStep(key, "publicInput.from -> proof1.from")
+      );
+      proof1.publicOutput[key].assertEquals(
+        proof2.publicInput[key],
+        errors.propertyNotMatchingStep(key, "proof1.to -> proof2.from")
+      );
+    }
+
+    checkProperty("stateRoot");
+    checkProperty("networkStateHash");
+    checkProperty("blockHashRoot");
+    checkProperty("eternalTransactionsHash");
+    checkProperty("incomingMessagesHash");
+    checkProperty("proverStateRemainder");
+
+    return proof2.publicOutput;
   }
 
   /**
@@ -238,38 +620,85 @@ export class BlockProver
   public zkProgramFactory(): PlainZkProgram<
     BlockProverPublicInput,
     BlockProverPublicOutput
-  > {
-    const StateTransitionProofClass =
-      this.stateTransitionProver.zkProgram.Proof;
-    const RuntimeProofClass = this.runtime.zkProgrammable.zkProgram.Proof;
-
-    const proveTransaction = this.proveTransaction.bind(this);
-    const merge = this.merge.bind(this);
-
-    const program = Experimental.ZkProgram({
+  >[] {
+    const { prover, stateTransitionProver, transactionProver } = this;
+    const StateTransitionProofClass = stateTransitionProver.zkProgram[0].Proof;
+    const TransactionProofClass = transactionProver.zkProgram[0].Proof;
+    const proveBlockBatchWithProofs =
+      prover.proveBlockBatchWithProofs.bind(prover);
+    const proveBlockBatchNoProofs = prover.proveBlockBatchNoProofs.bind(prover);
+    const merge = prover.merge.bind(prover);
+
+    const program = ZkProgram({
+      name: "BlockProver",
       publicInput: BlockProverPublicInput,
       publicOutput: BlockProverPublicOutput,
 
       methods: {
-        proveTransaction: {
+        proveBlockBatchWithProofs: {
           privateInputs: [
+            BlockProverStateInput,
+            NetworkState,
+            BlockHashMerkleTreeWitness,
+            BlockArgumentsBatch,
+            Bool,
+            Bool,
             StateTransitionProofClass,
-            RuntimeProofClass,
-            BlockProverExecutionData,
+            TransactionProofClass,
           ],
+          async method(
+            publicInput: BlockProverPublicInput,
+            stateWitness: BlockProverStateInput,
+            networkState: NetworkState,
+            blockWitness: BlockHashMerkleTreeWitness,
+            batch: BlockArgumentsBatch,
+            deferSTProof: Bool,
+            deferTransactionProof: Bool,
+            stateTransitionProof: StateTransitionProof,
+            transactionProof: TransactionProof
+          ) {
+            return {
+              publicOutput: await proveBlockBatchWithProofs(
+                publicInput,
+                stateWitness,
+                networkState,
+                blockWitness,
+                batch,
+                deferSTProof,
+                deferTransactionProof,
+                stateTransitionProof,
+                transactionProof
+              ),
+            };
+          },
+        },
 
-          method(
+        proveBlockBatchNoProofs: {
+          privateInputs: [
+            BlockProverStateInput,
+            NetworkState,
+            BlockHashMerkleTreeWitness,
+            BlockArgumentsBatch,
+            Bool,
+          ],
+          async method(
             publicInput: BlockProverPublicInput,
-            stateProof: StateTransitionProof,
-            appProof: Proof<void, MethodPublicOutput>,
-            executionData: BlockProverExecutionData
+            stateWitness: BlockProverStateInput,
+            networkState: NetworkState,
+            blockWitness: BlockHashMerkleTreeWitness,
+            batch: BlockArgumentsBatch,
+            finalize: Bool
           ) {
-            return proveTransaction(
-              publicInput,
-              stateProof,
-              appProof,
-              executionData
-            );
+            return {
+              publicOutput: await proveBlockBatchNoProofs(
+                publicInput,
+                stateWitness,
+                networkState,
+                blockWitness,
+                batch,
+                finalize
+              ),
+            };
           },
         },
 
@@ -279,29 +708,135 @@ export class BlockProver
             SelfProof<BlockProverPublicInput, BlockProverPublicOutput>,
           ],
 
-          method(
+          async method(
             publicInput: BlockProverPublicInput,
-            proof1: BlockProverProof,
-            proof2: BlockProverProof
+            proof1: BlockProof,
+            proof2: BlockProof
           ) {
-            return merge(publicInput, proof1, proof2);
+            return { publicOutput: await merge(publicInput, proof1, proof2) };
           },
         },
       },
     });
 
     const methods = {
-      proveTransaction: program.proveTransaction,
+      proveBlockBatchWithProofs: program.proveBlockBatchWithProofs,
+      proveBlockBatchNoProofs: program.proveBlockBatchNoProofs,
       merge: program.merge,
     };
 
-    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,
+      },
+    ];
+  }
+}
+
+/**
+ * BlockProver class, which aggregates a AppChainProof and
+ * a StateTransitionProof into a single BlockProof, that can
+ * then be merged to be committed to the base-layer contract
+ */
+@injectable()
+export class BlockProver
+  extends ProtocolModule
+  implements BlockProvable, CompilableModule
+{
+  public zkProgrammable: BlockProverProgrammable;
+
+  public constructor(
+    @inject("StateTransitionProver")
+    public readonly stateTransitionProver: WithZkProgrammable<
+      StateTransitionProverPublicInput,
+      StateTransitionProverPublicOutput
+    > &
+      StateTransitionProvable,
+    @inject("TransactionProver")
+    public readonly transactionProver: WithZkProgrammable<
+      TransactionProverPublicInput,
+      TransactionProverPublicOutput
+    > &
+      TransactionProvable,
+    @injectAll("ProvableBlockHook")
+    blockHooks: ProvableBlockHook<unknown>[],
+    @inject("StateServiceProvider")
+    stateServiceProvider: StateServiceProvider
+  ) {
+    super();
+    this.zkProgrammable = new BlockProverProgrammable(
+      this,
+      stateTransitionProver.zkProgrammable,
+      transactionProver.zkProgrammable,
+      blockHooks,
+      stateServiceProvider
+    );
+  }
+
+  public async compile(
+    registry: CompileRegistry
+  ): Promise<Record<string, CompileArtifact> | undefined> {
+    return await registry.forceProverExists(async () => {
+      await this.stateTransitionProver.compile(registry);
+      await this.transactionProver.compile(registry);
+      return await this.zkProgrammable.compile(registry);
+    });
+  }
+
+  public proveBlockBatchNoProofs(
+    publicInput: BlockProverPublicInput,
+    stateWitness: BlockProverStateInput,
+    networkState: NetworkState,
+    blockWitness: BlockHashMerkleTreeWitness,
+    batch: BlockArgumentsBatch,
+    finalize: Bool
+  ): Promise<BlockProverPublicOutput> {
+    return this.zkProgrammable.proveBlockBatchNoProofs(
+      publicInput,
+      stateWitness,
+      networkState,
+      blockWitness,
+      batch,
+      finalize
+    );
+  }
+
+  public proveBlockBatchWithProofs(
+    publicInput: BlockProverPublicInput,
+    stateWitness: BlockProverStateInput,
+    networkState: NetworkState,
+    blockWitness: BlockHashMerkleTreeWitness,
+    batch: BlockArgumentsBatch,
+    deferSTProof: Bool,
+    deferTransactionProof: Bool,
+    stateTransitionProof: StateTransitionProof,
+    transactionProof: TransactionProof
+  ): Promise<BlockProverPublicOutput> {
+    return this.zkProgrammable.proveBlockBatchWithProofs(
+      publicInput,
+      stateWitness,
+      networkState,
+      blockWitness,
+      batch,
+      deferSTProof,
+      deferTransactionProof,
+      stateTransitionProof,
+      transactionProof
+    );
+  }
+
+  public merge(
+    publicInput: BlockProverPublicInput,
+    proof1: BlockProof,
+    proof2: BlockProof
+  ): Promise<BlockProverPublicOutput> {
+    return this.zkProgrammable.merge(publicInput, proof1, proof2);
   }
 }