Fix bug in protocol and flow for merging blocks

Author: rpanic

packages/common/src/utils.ts

@@ -147,3 +147,5 @@ type NonMethodKeys<Type> = {
   [Key in keyof Type]: Type[Key] extends Function ? never : Key;
 }[keyof Type];
 export type NonMethods<Type> = Pick<Type, NonMethodKeys<Type>>;
+
+export const MAX_FIELD = Field(Field.ORDER - 1n);

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

@@ -24,6 +24,7 @@ export class BlockProverPublicInput extends Struct({
   blockHashRoot: Field,
   eternalTransactionsHash: Field,
   incomingMessagesHash: Field,
+  blockNumber: Field,
 }) {}
 
 export class BlockProverPublicOutput extends Struct({
@@ -36,15 +37,10 @@ export class BlockProverPublicOutput extends Struct({
   closed: Bool,
   blockNumber: Field,
 }) {
-  public equals(
-    input: BlockProverPublicInput,
-    closed: Bool,
-    blockNumber: Field
-  ): Bool {
+  public equals(input: BlockProverPublicInput, closed: Bool): Bool {
     const output2 = BlockProverPublicOutput.toFields({
       ...input,
       closed,
-      blockNumber,
     });
     const output1 = BlockProverPublicOutput.toFields(this);
     return output1

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

@@ -12,6 +12,7 @@ import {
 import { container, inject, injectable, injectAll } from "tsyringe";
 import {
   AreProofsEnabled,
+  MAX_FIELD,
   PlainZkProgram,
   provableMethod,
   WithZkProgrammable,
@@ -124,10 +125,6 @@ export interface BlockProverState {
   incomingMessagesHash: Field;
 }
 
-function maxField() {
-  return Field(Field.ORDER - 1n);
-}
-
 export type BlockProof = Proof<BlockProverPublicInput, BlockProverPublicOutput>;
 export type RuntimeProof = Proof<void, MethodPublicOutput>;
 
@@ -416,6 +413,11 @@ export class BlockProverProgrammable extends ZkProgrammable<
       "ExecutionData Networkstate doesn't equal public input hash"
     );
 
+    publicInput.blockNumber.assertEquals(
+      MAX_FIELD,
+      "blockNumber has to be MAX for transaction proofs"
+    );
+
     // Verify the [methodId, vk] tuple against the baked-in vk tree root
     const { verificationKey, witness: verificationKeyTreeWitness } =
       verificationKeyWitness;
@@ -445,7 +447,7 @@ export class BlockProverProgrammable extends ZkProgrammable<
 
     return new BlockProverPublicOutput({
       ...stateTo,
-      blockNumber: maxField(),
+      blockNumber: publicInput.blockNumber,
       closed: Bool(false),
     });
   }
@@ -525,16 +527,16 @@ export class BlockProverProgrammable extends ZkProgrammable<
       .not();
     stateTransitionProof.verifyIf(stsEmitted);
 
-    // Verify Transaction proof if it has at least 1 tx
+    // Verify Transaction proof if it has at least 1 tx - i.e. the
+    // input and output doesn't match fully
     // We have to compare the whole input and output because we can make no
     // assumptions about the values, since it can be an arbitrary dummy-proof
     const txProofOutput = transactionProof.publicOutput;
     const verifyTransactionProof = txProofOutput.equals(
       transactionProof.publicInput,
-      txProofOutput.closed,
-      txProofOutput.blockNumber
+      txProofOutput.closed
     );
-    transactionProof.verifyIf(verifyTransactionProof);
+    transactionProof.verifyIf(verifyTransactionProof.not());
 
     // 2. Execute beforeBlock hooks
     const beforeBlockResult = await this.executeBlockHooks(
@@ -616,6 +618,8 @@ export class BlockProverProgrammable extends ZkProgrammable<
     // Calculate the new block index
     const blockIndex = blockWitness.calculateIndex();
 
+    blockIndex.assertEquals(publicInput.blockNumber);
+
     blockWitness
       .calculateRoot(Field(0))
       .assertEquals(
@@ -633,7 +637,7 @@ export class BlockProverProgrammable extends ZkProgrammable<
 
     return new BlockProverPublicOutput({
       ...state,
-      blockNumber: blockIndex,
+      blockNumber: blockIndex.add(1),
       closed: Bool(true),
     });
   }
@@ -740,19 +744,25 @@ export class BlockProverProgrammable extends ZkProgrammable<
     //   assert proof1.height == proof2.height
     // }
 
-    const proof1Height = proof1.publicOutput.blockNumber;
     const proof1Closed = proof1.publicOutput.closed;
-    const proof2Height = proof2.publicOutput.blockNumber;
     const proof2Closed = proof2.publicOutput.closed;
 
-    const isValidTransactionMerge = proof1Height
-      .equals(maxField())
-      .and(proof2Height.equals(proof1Height))
+    const blockNumberProgressionValid = publicInput.blockNumber
+      .equals(proof1.publicInput.blockNumber)
+      .and(
+        proof1.publicOutput.blockNumber.equals(proof2.publicInput.blockNumber)
+      );
+
+    // For tx proofs, we check that the progression starts and end with MAX
+    // in addition to that both proofs are non-closed
+    const isValidTransactionMerge = publicInput.blockNumber
+      .equals(MAX_FIELD)
+      .and(blockNumberProgressionValid)
       .and(proof1Closed.or(proof2Closed).not());
 
     const isValidClosedMerge = proof1Closed
       .and(proof2Closed)
-      .and(proof1Height.add(1).equals(proof2Height));
+      .and(blockNumberProgressionValid);
 
     isValidTransactionMerge
       .or(isValidClosedMerge)
@@ -765,9 +775,8 @@ export class BlockProverProgrammable extends ZkProgrammable<
       blockHashRoot: proof2.publicOutput.blockHashRoot,
       eternalTransactionsHash: proof2.publicOutput.eternalTransactionsHash,
       incomingMessagesHash: proof2.publicOutput.incomingMessagesHash,
-      // Provable.if(isValidClosedMerge, Bool(true), Bool(false));
       closed: isValidClosedMerge,
-      blockNumber: proof2Height,
+      blockNumber: proof2.publicOutput.blockNumber,
     });
   }
 

packages/sequencer/src/protocol/production/BlockTaskFlowService.ts

@@ -9,7 +9,7 @@ import {
   Protocol,
   StateTransitionProof,
 } from "@proto-kit/protocol";
-import { log, MOCK_PROOF } from "@proto-kit/common";
+import { log, MAX_FIELD, MOCK_PROOF } from "@proto-kit/common";
 
 import { TaskQueue } from "../../worker/queue/TaskQueue";
 import { Flow, FlowCreator } from "../../worker/flow/Flow";
@@ -171,9 +171,9 @@ export class BlockTaskFlowService {
         mappingTask: this.blockProvingTask,
         reductionTask: this.blockReductionTask,
 
-        mergableFunction: (a, b) =>
+        mergableFunction: (a, b) => {
           // TODO Proper replication of merge logic
-          a.publicOutput.stateRoot
+          const part1 = a.publicOutput.stateRoot
             .equals(b.publicInput.stateRoot)
             .and(
               a.publicOutput.blockHashRoot.equals(b.publicInput.blockHashRoot)
@@ -189,7 +189,28 @@ export class BlockTaskFlowService {
               )
             )
             .and(a.publicOutput.closed.equals(b.publicOutput.closed))
-            .toBoolean(),
+            .toBoolean();
+
+          const proof1Closed = a.publicOutput.closed;
+          const proof2Closed = b.publicOutput.closed;
+
+          const blockNumberProgressionValid = a.publicOutput.blockNumber.equals(
+            b.publicInput.blockNumber
+          );
+
+          const isValidTransactionMerge = a.publicInput.blockNumber
+            .equals(MAX_FIELD)
+            .and(blockNumberProgressionValid)
+            .and(proof1Closed.or(proof2Closed).not());
+
+          const isValidClosedMerge = proof1Closed
+            .and(proof2Closed)
+            .and(blockNumberProgressionValid);
+
+          return (
+            part1 && isValidClosedMerge.or(isValidTransactionMerge).toBoolean()
+          );
+        },
       },
       this.flowCreator
     );
@@ -286,13 +307,13 @@ export class BlockTaskFlowService {
               blockTrace.block.publicInput.eternalTransactionsHash,
             incomingMessagesHash:
               blockTrace.block.publicInput.incomingMessagesHash,
+            blockNumber: MAX_FIELD,
           };
           const publicInput = new BlockProverPublicInput(piObject);
 
           // TODO Set publicInput.stateRoot to result after block hooks!
           const publicOutput = new BlockProverPublicOutput({
             ...piObject,
-            blockNumber: Field(Field.ORDER - 1n),
             closed: Bool(true),
           });
 

packages/sequencer/src/protocol/production/TransactionTraceService.ts

@@ -10,7 +10,7 @@ import {
   StateTransitionProverPublicInput,
   StateTransitionType,
 } from "@proto-kit/protocol";
-import { RollupMerkleTree } from "@proto-kit/common";
+import { MAX_FIELD, RollupMerkleTree } from "@proto-kit/common";
 import { Bool, Field } from "o1js";
 import chunk from "lodash/chunk";
 
@@ -133,6 +133,7 @@ export class TransactionTraceService {
       blockHashRoot: block.block.fromBlockHashRoot,
       eternalTransactionsHash: block.block.fromEternalTransactionsHash,
       incomingMessagesHash: block.block.fromMessagesHash,
+      blockNumber: block.block.height,
     });
 
     return {
@@ -241,6 +242,7 @@ export class TransactionTraceService {
           incomingMessagesHash,
           networkStateHash: networkState.hash(),
           blockHashRoot: Field(0),
+          blockNumber: MAX_FIELD,
         },
 
         executionData: {

packages/sequencer/test/integration/BlockProduction.test.ts

@@ -473,10 +473,17 @@ describe("block production", () => {
     [1, 2, 1],
     [1, 1, 2],
     [2, 2, 2],
+    [1, 14, 0],
   ])(
     "should produce multiple blocks with multiple batches with multiple transactions",
     async (batches, blocksPerBatch, txsPerBlock) => {
-      expect.assertions(2 * batches + 3 * batches * blocksPerBatch);
+      expect.assertions(
+        2 * batches +
+          1 * batches * blocksPerBatch +
+          2 * batches * blocksPerBatch * txsPerBlock
+      );
+
+      log.setLevel("DEBUG");
 
       const sender = PrivateKey.random();
 
@@ -511,8 +518,11 @@ describe("block production", () => {
           const block = await blockTrigger.produceBlock();
 
           expect(block).toBeDefined();
-          expect(block!.transactions).toHaveLength(txsPerBlock);
-          expect(block!.transactions[0].status.toBoolean()).toBe(true);
+
+          for (let k = 0; k < txsPerBlock; k++) {
+            expect(block!.transactions).toHaveLength(txsPerBlock);
+            expect(block!.transactions[0].status.toBoolean()).toBe(true);
+          }
         }
 
         const batch = await blockTrigger.produceBatch();