feat(fast-inbox): world-state per-block message insertion (A-1380)

yarn-project/txe/src/state_machine/synchronizer.ts

@@ -35,10 +35,14 @@ export class TXESynchronizer implements WorldStateSynchronizer {
   }
 
   public async handleL2Block(block: L2Block, l1ToL2Messages: Fr[] = []) {
-    await this.nativeWorldStateService.handleL2BlockAndMessages(
-      block,
-      padArrayEnd<Fr, number>(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP),
-    );
+    // Pad the bundle only for a first-in-checkpoint block, matching how the circuits (and native world state) build the
+    // message tree. TXE mines one block per checkpoint, so this is always the first block, but keep the condition
+    // explicit so the caller matches the per-block message-insertion semantics of handleL2BlockAndMessages.
+    const messages =
+      block.indexWithinCheckpoint === 0
+        ? padArrayEnd<Fr, number>(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP)
+        : l1ToL2Messages;
+    await this.nativeWorldStateService.handleL2BlockAndMessages(block, messages);
 
     this.blockNumber = block.header.globalVariables.blockNumber;
   }

yarn-project/world-state/src/native/native_world_state.test.ts

@@ -32,7 +32,14 @@ import { tmpdir } from 'os';
 import { join } from 'path';
 
 import type { WorldStateTreeMapSizes } from '../synchronizer/factory.js';
-import { assertSameState, compareChains, mockBlock, mockEmptyBlock, updateBlockState } from '../test/utils.js';
+import {
+  assertSameState,
+  compareChains,
+  mockBlock,
+  mockBlockWithIndex,
+  mockEmptyBlock,
+  updateBlockState,
+} from '../test/utils.js';
 import { INITIAL_NULLIFIER_TREE_SIZE, INITIAL_PUBLIC_DATA_TREE_SIZE } from '../world-state-db/merkle_tree_db.js';
 import type { WorldStateStatusSummary } from './message.js';
 import { NativeWorldStateService, WORLD_STATE_DB_VERSION, WORLD_STATE_DIR } from './native_world_state.js';
@@ -149,14 +156,140 @@ describe('NativeWorldState', () => {
       expect(status.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP));
     });
 
-    it('throws error if messages are provided for non-first block', async () => {
-      const isFirstBlock = false;
-      const numMessages = 1;
-      const { block, messages } = await mockBlock(BlockNumber(1), 1, fork, 1, numMessages, isFirstBlock);
+    it('appends a non-first block bundle without padding', async () => {
+      const numMessages = 3;
+      const { block, messages } = await mockBlockWithIndex(
+        BlockNumber(1),
+        /*indexWithinCheckpoint=*/ 1,
+        1,
+        fork,
+        numMessages,
+        1,
+      );
+
+      const status = await ws.handleL2BlockAndMessages(block, messages);
+
+      // Non-first blocks append their bundle exactly as given (no padding to NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP).
+      expect(status.meta.messageTreeMeta.size).toBe(BigInt(numMessages));
+    });
+  });
+
+  describe('Per-block message insertion', () => {
+    let ws: NativeWorldStateService;
+
+    beforeEach(async () => {
+      ws = await NativeWorldStateService.tmp();
+    });
+
+    afterEach(async () => {
+      await ws.close();
+    });
+
+    it('advances the L1-to-L2 message tree per block, including on non-first blocks', async () => {
+      const fork = await ws.fork();
+
+      // Block 1 is first-in-checkpoint, so its bundle is padded to NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP (this is how
+      // the circuits build the tree).
+      const { block: b1, messages: m1 } = await mockBlockWithIndex(BlockNumber(1), /*index=*/ 0, 1, fork, 3, 1);
+      const s1 = await ws.handleL2BlockAndMessages(b1, m1);
+      expect(s1.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP));
+      expect(s1.meta.messageTreeMeta.unfinalizedBlockHeight).toBe(1);
+
+      // Block 2 is non-first and carries no messages: the message tree size and root are unchanged, but the tree is
+      // still committed as a new block (so its per-block history stays in lockstep with the other trees).
+      const { block: b2, messages: m2 } = await mockBlockWithIndex(BlockNumber(2), /*index=*/ 1, 1, fork, 0, 1);
+      const s2 = await ws.handleL2BlockAndMessages(b2, m2);
+      expect(s2.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP));
+      expect(s2.meta.messageTreeMeta.root).toEqual(s1.meta.messageTreeMeta.root);
+      expect(s2.meta.messageTreeMeta.unfinalizedBlockHeight).toBe(2);
+
+      // Block 3 is non-first and carries messages: the bundle is appended unpadded, so the tree grows by exactly the
+      // bundle size and the root changes on a non-first block.
+      const { block: b3, messages: m3 } = await mockBlockWithIndex(BlockNumber(3), /*index=*/ 2, 1, fork, 5, 1);
+      const s3 = await ws.handleL2BlockAndMessages(b3, m3);
+      expect(s3.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP + 5));
+      expect(s3.meta.messageTreeMeta.root).not.toEqual(s2.meta.messageTreeMeta.root);
+      expect(s3.meta.messageTreeMeta.unfinalizedBlockHeight).toBe(3);
 
-      await expect(ws.handleL2BlockAndMessages(block, messages)).rejects.toThrow(
-        'L1 to L2 messages must be empty for non-first blocks',
+      await fork.close();
+
+      // A fork opened at block 2 sees exactly the first two bundles (3 padded + 0); at block 3 it also sees the third.
+      const forkAt2 = await ws.fork(BlockNumber(2));
+      expect((await forkAt2.getTreeInfo(MerkleTreeId.L1_TO_L2_MESSAGE_TREE)).size).toBe(
+        BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP),
+      );
+      await forkAt2.close();
+
+      const forkAt3 = await ws.fork(BlockNumber(3));
+      expect((await forkAt3.getTreeInfo(MerkleTreeId.L1_TO_L2_MESSAGE_TREE)).size).toBe(
+        BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP + 5),
       );
+      await forkAt3.close();
+    });
+
+    it('unwinds per block, reverting exactly the messages appended by a non-first block', async () => {
+      const fork = await ws.fork();
+
+      const { block: b1, messages: m1 } = await mockBlockWithIndex(BlockNumber(1), /*index=*/ 0, 1, fork, 2, 1);
+      const s1 = await ws.handleL2BlockAndMessages(b1, m1);
+
+      // Non-first block carrying 4 messages.
+      const { block: b2, messages: m2 } = await mockBlockWithIndex(BlockNumber(2), /*index=*/ 1, 1, fork, 4, 1);
+      const s2 = await ws.handleL2BlockAndMessages(b2, m2);
+      expect(s2.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP + 4));
+
+      // Non-first block carrying 5 messages.
+      const { block: b3, messages: m3 } = await mockBlockWithIndex(BlockNumber(3), /*index=*/ 2, 1, fork, 5, 1);
+      const s3 = await ws.handleL2BlockAndMessages(b3, m3);
+      expect(s3.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP + 9));
+      await fork.close();
+
+      // Unwind block 3: the message tree returns to the post-block-2 state.
+      const afterUnwind3 = await ws.unwindBlocks(BlockNumber(2));
+      expect(afterUnwind3.meta.messageTreeMeta.size).toBe(s2.meta.messageTreeMeta.size);
+      expect(afterUnwind3.meta.messageTreeMeta.root).toEqual(s2.meta.messageTreeMeta.root);
+
+      // Unwind block 2 (a message-carrying non-first block): its 4 messages are reverted, back to the post-block-1
+      // state — the pending-chain rollback does not assume the message tree only changes at checkpoint starts.
+      const afterUnwind2 = await ws.unwindBlocks(BlockNumber(1));
+      expect(afterUnwind2.meta.messageTreeMeta.size).toBe(s1.meta.messageTreeMeta.size);
+      expect(afterUnwind2.meta.messageTreeMeta.root).toEqual(s1.meta.messageTreeMeta.root);
+
+      // Re-syncing after the unwind reconverges: fresh non-first blocks append cleanly on top of block 1.
+      const resyncFork = await ws.fork();
+      const { block: b2b, messages: m2b } = await mockBlockWithIndex(BlockNumber(2), /*index=*/ 1, 1, resyncFork, 4, 1);
+      const s2b = await ws.handleL2BlockAndMessages(b2b, m2b);
+      expect(s2b.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP + 4));
+
+      const { block: b3b, messages: m3b } = await mockBlockWithIndex(BlockNumber(3), /*index=*/ 2, 1, resyncFork, 5, 1);
+      const s3b = await ws.handleL2BlockAndMessages(b3b, m3b);
+      expect(s3b.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP + 9));
+      await resyncFork.close();
+    });
+
+    it('leaves the message tree byte-identical on every non-first block of a legacy-shaped checkpoint', async () => {
+      const fork = await ws.fork();
+
+      // Legacy call shape: the whole (padded) checkpoint bundle is attached to the first block; non-first blocks carry
+      // an empty bundle. With this shape the code change is a no-op, so the message tree must match the pre-change
+      // behaviour (identical trees) at every block, not just at the checkpoint end.
+      const { block: b1, messages: m1 } = await mockBlockWithIndex(BlockNumber(1), /*index=*/ 0, 2, fork, 6, 2);
+      const s1 = await ws.handleL2BlockAndMessages(b1, m1);
+      expect(s1.meta.messageTreeMeta.size).toBe(BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP));
+
+      for (let index = 1; index <= 2; index++) {
+        const blockNumber = index + 1;
+        const { block, messages } = await mockBlockWithIndex(BlockNumber(blockNumber), index, 2, fork, 0, 2);
+        const status = await ws.handleL2BlockAndMessages(block, messages);
+
+        // The message tree is untouched by non-first blocks in the legacy shape.
+        expect(status.meta.messageTreeMeta.size).toEqual(s1.meta.messageTreeMeta.size);
+        expect(status.meta.messageTreeMeta.root).toEqual(s1.meta.messageTreeMeta.root);
+        // But the chain as a whole still advances: the archive tree grows with each block.
+        expect(status.meta.archiveTreeMeta.unfinalizedBlockHeight).toBe(blockNumber);
+      }
+
+      await fork.close();
     });
   });
 

yarn-project/world-state/src/native/native_world_state.ts

@@ -264,18 +264,20 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
   }
 
   public async handleL2BlockAndMessages(l2Block: L2Block, l1ToL2Messages: Fr[]): Promise<WorldStateStatusFull> {
+    // Any block may carry an L1-to-L2 message bundle and transition the L1-to-L2 message tree. Pre-flip the legacy
+    // synchronizer only ever attaches messages to the first block of a checkpoint (the whole padded checkpoint bundle),
+    // so first-in-checkpoint bundles are padded to NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP to match how the circuits build
+    // the tree, producing bit-identical trees. Non-first bundles are appended exactly as given (the post-flip compact
+    // append from the circuits). The "non-first blocks carry no messages" invariant of the legacy flow is now enforced
+    // by the caller (the synchronizer) rather than here, so the API accepts the new per-block shape ahead of the flip.
+    // Padding is the caller's transitional concern and moves entirely to the caller at the flip.
     const isFirstBlock = l2Block.indexWithinCheckpoint === 0;
-    if (!isFirstBlock && l1ToL2Messages.length > 0) {
-      throw new Error(
-        `L1 to L2 messages must be empty for non-first blocks, but got ${l1ToL2Messages.length} messages for block ${l2Block.number}.`,
-      );
-    }
 
-    // We have to pad the given l1 to l2 messages, and the note hashes and nullifiers within tx effects, because that's
-    // how the trees are built by circuits.
+    // We have to pad the note hashes and nullifiers within tx effects (and, for a first-in-checkpoint block, the l1 to
+    // l2 messages) because that's how the trees are built by circuits.
     const paddedL1ToL2Messages = isFirstBlock
       ? padArrayEnd<Fr, number>(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP)
-      : [];
+      : l1ToL2Messages;
 
     const paddedNoteHashes = l2Block.body.txEffects.flatMap(txEffect =>
       padArrayEnd(txEffect.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),

yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.test.ts

@@ -1,4 +1,4 @@
-import { BlockNumber, CheckpointNumber } from '@aztec/foundation/branded-types';
+import { BlockNumber, CheckpointNumber, IndexWithinCheckpoint } from '@aztec/foundation/branded-types';
 import { timesParallel } from '@aztec/foundation/collection';
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, createLogger } from '@aztec/foundation/log';
@@ -279,6 +279,17 @@ describe('ServerWorldStateSynchronizer', () => {
     expect(merkleTreeDb.handleL2BlockAndMessages.mock.calls[5][1]).toEqual([]);
   });
 
+  it('rejects a non-first block that carries L1->L2 messages (transitional invariant)', async () => {
+    // World state accepts a bundle on any block, but pre-flip the synchronizer must only attach messages to the first
+    // block of a checkpoint. The call-site guard enforces that until the flip switches to per-block derivation.
+    const nonFirstBlock = await L2Block.random(BlockNumber(2), { indexWithinCheckpoint: IndexWithinCheckpoint(1) });
+
+    await expect(server.callHandleL2Block(nonFirstBlock, [Fr.random()])).rejects.toThrow(
+      'L1 to L2 messages must be empty for non-first blocks',
+    );
+    expect(merkleTreeDb.handleL2BlockAndMessages).not.toHaveBeenCalled();
+  });
+
   describe('getVerifiedSnapshot', () => {
     let snapshot: MockProxy<MerkleTreeReadOperations>;
 
@@ -355,6 +366,10 @@ class TestWorldStateSynchronizer extends ServerWorldStateSynchronizer {
     return this.mockBlockStream;
   }
 
+  public callHandleL2Block(block: L2Block, messages: Fr[]) {
+    return this.handleL2Block(block, messages);
+  }
+
   public override getL2Tips() {
     return Promise.resolve({
       proposed: this.latest,

yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts

@@ -3,6 +3,7 @@ import type { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import { promiseWithResolvers } from '@aztec/foundation/promise';
 import { elapsed } from '@aztec/foundation/timer';
+import { assert } from '@aztec/foundation/validation';
 import {
   type BlockHash,
   EventDrivenL2BlockStream,
@@ -406,7 +407,14 @@ export class ServerWorldStateSynchronizer
    * @param l1ToL2Messages - The L1 to L2 messages for the block.
    * @returns Whether the block handled was produced by this same node.
    */
-  private async handleL2Block(l2Block: L2Block, l1ToL2Messages: Fr[]): Promise<WorldStateStatusFull> {
+  protected async handleL2Block(l2Block: L2Block, l1ToL2Messages: Fr[]): Promise<WorldStateStatusFull> {
+    // Transitional invariant (pre-flip): the legacy per-checkpoint fetch only ever attaches messages to the first
+    // block of a checkpoint, so no non-first block should carry a bundle here. World state itself now accepts a bundle
+    // on any block; this rule is owned by the synchronizer until the flip switches it to per-block message derivation.
+    assert(
+      l2Block.indexWithinCheckpoint === 0 || l1ToL2Messages.length === 0,
+      `L1 to L2 messages must be empty for non-first blocks, but got ${l1ToL2Messages.length} messages for block ${l2Block.number} (index ${l2Block.indexWithinCheckpoint} within checkpoint).`,
+    );
     this.log.debug(`Pushing L2 block ${l2Block.number} to merkle tree db `, {
       blockNumber: l2Block.number,
       blockHash: await l2Block.hash().then(h => h.toString()),

yarn-project/world-state/src/test/utils.ts

@@ -77,16 +77,32 @@ export async function updateBlockState(block: L2Block, l1ToL2Messages: Fr[], for
   block.archive = new AppendOnlyTreeSnapshot(Fr.fromBuffer(archiveState.root), Number(archiveState.size));
 }
 
-export async function mockBlock(
+export function mockBlock(
   blockNum: BlockNumber,
   size: number,
   fork: MerkleTreeWriteOperations,
   maxEffects: number | undefined = 1000, // Defaults to the maximum tx effects.
   numL1ToL2Messages: number = NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
   isFirstBlockInCheckpoint: boolean = true,
+) {
+  return mockBlockWithIndex(blockNum, isFirstBlockInCheckpoint ? 0 : 1, size, fork, numL1ToL2Messages, maxEffects);
+}
+
+/**
+ * Builds a mock L2 block at an explicit position within its checkpoint, applying its state (including its L1-to-L2
+ * message bundle) to the given fork. Unlike {@link mockBlock}, the caller chooses the `indexWithinCheckpoint`, so
+ * non-first blocks can carry message bundles — exercising the per-block message insertion path.
+ */
+export async function mockBlockWithIndex(
+  blockNum: BlockNumber,
+  indexWithinCheckpoint: number,
+  size: number,
+  fork: MerkleTreeWriteOperations,
+  numL1ToL2Messages: number = NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
+  maxEffects: number | undefined = 1000, // Defaults to the maximum tx effects.
 ) {
   const block = await L2Block.random(blockNum, {
-    indexWithinCheckpoint: isFirstBlockInCheckpoint ? IndexWithinCheckpoint(0) : IndexWithinCheckpoint(1),
+    indexWithinCheckpoint: IndexWithinCheckpoint(indexWithinCheckpoint),
     txsPerBlock: size,
     txOptions: { maxEffects },
   });

yarn-project/world-state/src/world-state-db/merkle_tree_db.ts

@@ -31,7 +31,11 @@ export const INITIAL_PUBLIC_DATA_TREE_SIZE = 2 * MAX_TOTAL_PUBLIC_DATA_UPDATE_RE
 
 export interface MerkleTreeAdminDatabase extends ForkMerkleTreeOperations, ReadonlyWorldStateAccess {
   /**
-   * Handles a single L2 block (i.e. Inserts the new note hashes into the merkle tree).
+   * Handles a single L2 block: inserts its note hashes, nullifiers, public data writes, and the block's L1-to-L2
+   * message bundle into the merkle trees. Any block may carry a message bundle and transition the L1-to-L2 message
+   * tree, not just the first block of a checkpoint. A first-in-checkpoint bundle is padded to
+   * NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP to match how the circuits build the tree; a non-first bundle is appended
+   * exactly as given. Padding is a transitional concern of this method that moves entirely to the caller at the flip.
    * @param block - The L2 block to handle.
    * @param l1ToL2Messages - The L1 to L2 messages for the block.
    */