@aztec/archiver 0.0.1-commit.c2595eba → 0.0.1-commit.c2eed6949

package/src/test/fake_l1_state.ts

@@ -1,5 +1,6 @@
 import type { BlobClientInterface } from '@aztec/blob-client/client';
 import { type Blob, getBlobsPerL1Block, getPrefixedEthBlobCommitments } from '@aztec/blob-lib';
+import { INITIAL_CHECKPOINT_NUMBER } from '@aztec/constants';
 import type { CheckpointProposedLog, InboxContract, MessageSentLog, RollupContract } from '@aztec/ethereum/contracts';
 import { MULTI_CALL_3_ADDRESS } from '@aztec/ethereum/contracts';
 import type { ViemPublicClient } from '@aztec/ethereum/types';
@@ -14,6 +15,7 @@ import { CommitteeAttestation, CommitteeAttestationsAndSigners, L2Block } from '
 import { Checkpoint } from '@aztec/stdlib/checkpoint';
 import { getSlotAtTimestamp } from '@aztec/stdlib/epoch-helpers';
 import { InboxLeaf } from '@aztec/stdlib/messaging';
+import { ConsensusPayload, SignatureDomainSeparator } from '@aztec/stdlib/p2p';
 import {
   makeAndSignCommitteeAttestationsAndSigners,
   makeCheckpointAttestationFromCheckpoint,
@@ -22,7 +24,16 @@ import {
 import { AppendOnlyTreeSnapshot } from '@aztec/stdlib/trees';
 
 import { type MockProxy, mock } from 'jest-mock-extended';
-import { type FormattedBlock, type Transaction, encodeFunctionData, multicall3Abi, toHex } from 'viem';
+import {
+  type AbiParameter,
+  type FormattedBlock,
+  type Transaction,
+  encodeAbiParameters,
+  encodeFunctionData,
+  keccak256,
+  multicall3Abi,
+  toHex,
+} from 'viem';
 
 import { updateRollingHash } from '../structs/inbox_message.js';
 
@@ -87,6 +98,10 @@ type CheckpointData = {
   blobHashes: `0x${string}`[];
   blobs: Blob[];
   signers: Secp256k1Signer[];
+  /** Hash of the packed attestations, matching what the L1 event emits. */
+  attestationsHash: Buffer32;
+  /** Payload digest, matching what the L1 event emits. */
+  payloadDigest: Buffer32;
   /** If true, archiveAt will ignore it */
   pruned?: boolean;
 };
@@ -131,12 +146,17 @@ export class FakeL1State {
   private provenCheckpointNumber: CheckpointNumber = CheckpointNumber(0);
   private targetCommitteeSize: number = 0;
   private version: bigint = 1n;
+  private canPruneResult: boolean = false;
 
   // Computed from checkpoints based on L1 block visibility
   private pendingCheckpointNumber: CheckpointNumber = CheckpointNumber(0);
 
+  // The L1 block number reported as "finalized" (defaults to the start block)
+  private finalizedL1BlockNumber: bigint;
+
   constructor(private readonly config: FakeL1StateConfig) {
     this.l1BlockNumber = config.l1StartBlock;
+    this.finalizedL1BlockNumber = config.l1StartBlock;
     this.lastArchive = new AppendOnlyTreeSnapshot(config.genesisArchiveRoot, 1);
   }
 
@@ -193,10 +213,10 @@ export class FakeL1State {
     // Store the messages internally so they match the checkpoint's inHash
     this.addMessages(checkpointNumber, messagesL1BlockNumber, messages);
 
-    // Create the transaction and blobs
-    const tx = this.makeRollupTx(checkpoint, signers);
-    const blobHashes = this.makeVersionedBlobHashes(checkpoint);
-    const blobs = this.makeBlobsFromCheckpoint(checkpoint);
+    // Create the transaction, blobs, and event hashes
+    const { tx, attestationsHash, payloadDigest } = await this.makeRollupTx(checkpoint, signers);
+    const blobHashes = await this.makeVersionedBlobHashes(checkpoint);
+    const blobs = await this.makeBlobsFromCheckpoint(checkpoint);
 
     // Store the checkpoint data
     this.checkpoints.push({
@@ -207,6 +227,8 @@ export class FakeL1State {
       blobHashes,
       blobs,
       signers,
+      attestationsHash,
+      payloadDigest,
     });
 
     // Update last archive for auto-chaining
@@ -266,16 +288,40 @@ export class FakeL1State {
     this.updatePendingCheckpointNumber();
   }
 
+  /** Sets the L1 block number that will be reported as "finalized". */
+  setFinalizedL1BlockNumber(blockNumber: bigint): void {
+    this.finalizedL1BlockNumber = blockNumber;
+  }
+
   /** Marks a checkpoint as proven. Updates provenCheckpointNumber. */
   markCheckpointAsProven(checkpointNumber: CheckpointNumber): void {
     this.provenCheckpointNumber = checkpointNumber;
   }
 
+  /**
+   * Simulates what `rollup.getProvenCheckpointNumber({ blockNumber: atL1Block })` would return.
+   */
+  getProvenCheckpointNumberAtL1Block(atL1Block: bigint): CheckpointNumber {
+    if (this.provenCheckpointNumber === 0) {
+      return CheckpointNumber(0);
+    }
+    const checkpoint = this.checkpoints.find(cp => cp.checkpointNumber === this.provenCheckpointNumber);
+    if (checkpoint && checkpoint.l1BlockNumber <= atL1Block) {
+      return this.provenCheckpointNumber;
+    }
+    return CheckpointNumber(0);
+  }
+
   /** Sets the target committee size for attestation validation. */
   setTargetCommitteeSize(size: number): void {
     this.targetCommitteeSize = size;
   }
 
+  /** Sets whether the rollup contract would allow pruning at the next block. */
+  setCanPrune(value: boolean): void {
+    this.canPruneResult = value;
+  }
+
   /**
    * Removes all entries for a checkpoint number (simulates L1 reorg or prune).
    * Note: Does NOT remove messages for this checkpoint (use numL1ToL2Messages: 0 when re-adding).
@@ -384,6 +430,13 @@ export class FakeL1State {
       });
     });
 
+    mockRollup.getProvenCheckpointNumber.mockImplementation((options?: { blockNumber?: bigint }) => {
+      const atBlock = options?.blockNumber ?? this.l1BlockNumber;
+      return Promise.resolve(this.getProvenCheckpointNumberAtL1Block(atBlock));
+    });
+
+    mockRollup.canPruneAtTime.mockImplementation(() => Promise.resolve(this.canPruneResult));
+
     // Mock the wrapper method for fetching checkpoint events
     mockRollup.getCheckpointProposedEvents.mockImplementation((fromBlock: bigint, toBlock: bigint) =>
       Promise.resolve(this.getCheckpointProposedLogs(fromBlock, toBlock)),
@@ -398,13 +451,22 @@ export class FakeL1State {
   createMockInboxContract(_publicClient: MockProxy<ViemPublicClient>): MockProxy<InboxContract> {
     const mockInbox = mock<InboxContract>();
 
-    mockInbox.getState.mockImplementation(() =>
-      Promise.resolve({
+    mockInbox.getState.mockImplementation(() => {
+      // treeInProgress must be > any sealed checkpoint. On L1, a checkpoint can only be proposed
+      // after its messages are sealed, so treeInProgress > checkpointNumber for all published checkpoints.
+      const maxFromMessages =
+        this.messages.length > 0 ? Math.max(...this.messages.map(m => Number(m.checkpointNumber))) + 1 : 0;
+      const maxFromCheckpoints =
+        this.checkpoints.length > 0
+          ? Math.max(...this.checkpoints.filter(cp => !cp.pruned).map(cp => Number(cp.checkpointNumber))) + 1
+          : 0;
+      const treeInProgress = Math.max(maxFromMessages, maxFromCheckpoints, INITIAL_CHECKPOINT_NUMBER);
+      return Promise.resolve({
         messagesRollingHash: this.messagesRollingHash,
         totalMessagesInserted: BigInt(this.messages.length),
-        treeInProgress: 0n,
-      }),
-    );
+        treeInProgress: BigInt(treeInProgress),
+      });
+    });
 
     // Mock the wrapper methods for fetching message events
     mockInbox.getMessageSentEvents.mockImplementation((fromBlock: bigint, toBlock: bigint) =>
@@ -425,10 +487,13 @@ export class FakeL1State {
     publicClient.getChainId.mockResolvedValue(1);
     publicClient.getBlockNumber.mockImplementation(() => Promise.resolve(this.l1BlockNumber));
 
-    // Use async function pattern that existing test uses for getBlock
-
-    publicClient.getBlock.mockImplementation((async (args: { blockNumber?: bigint } = {}) => {
-      const blockNum = args.blockNumber ?? (await publicClient.getBlockNumber());
+    publicClient.getBlock.mockImplementation((async (args: { blockNumber?: bigint; blockTag?: string } = {}) => {
+      let blockNum: bigint;
+      if (args.blockTag === 'finalized') {
+        blockNum = this.finalizedL1BlockNumber;
+      } else {
+        blockNum = args.blockNumber ?? (await publicClient.getBlockNumber());
+      }
       return {
         number: blockNum,
         timestamp: BigInt(blockNum) * BigInt(this.config.ethereumSlotDuration) + this.config.l1GenesisTime,
@@ -502,10 +567,8 @@ export class FakeL1State {
           checkpointNumber: cpData.checkpointNumber,
           archive: cpData.checkpoint.archive.root,
           versionedBlobHashes: cpData.blobHashes.map(h => Buffer.from(h.slice(2), 'hex')),
-          // These are intentionally undefined to skip hash validation in the archiver
-          // (validation is skipped when these fields are falsy)
-          payloadDigest: undefined,
-          attestationsHash: undefined,
+          attestationsHash: cpData.attestationsHash,
+          payloadDigest: cpData.payloadDigest,
         },
       }));
   }
@@ -531,14 +594,17 @@ export class FakeL1State {
       }));
   }
 
-  private makeRollupTx(checkpoint: Checkpoint, signers: Secp256k1Signer[]): Transaction {
+  private async makeRollupTx(
+    checkpoint: Checkpoint,
+    signers: Secp256k1Signer[],
+  ): Promise<{ tx: Transaction; attestationsHash: Buffer32; payloadDigest: Buffer32 }> {
     const attestations = signers
       .map(signer => makeCheckpointAttestationFromCheckpoint(checkpoint, signer))
       .map(attestation => CommitteeAttestation.fromSignature(attestation.signature))
       .map(committeeAttestation => committeeAttestation.toViem());
 
     const header = checkpoint.header.toViem();
-    const blobInput = getPrefixedEthBlobCommitments(getBlobsPerL1Block(checkpoint.toBlobFields()));
+    const blobInput = getPrefixedEthBlobCommitments(await getBlobsPerL1Block(checkpoint.toBlobFields()));
     const archive = toHex(checkpoint.archive.root.toBuffer());
     const attestationsAndSigners = new CommitteeAttestationsAndSigners(
       attestations.map(attestation => CommitteeAttestation.fromViem(attestation)),
@@ -549,6 +615,8 @@ export class FakeL1State {
       signers[0],
     );
 
+    const packedAttestations = attestationsAndSigners.getPackedAttestations();
+
     const rollupInput = encodeFunctionData({
       abi: RollupAbi,
       functionName: 'propose',
@@ -558,7 +626,7 @@ export class FakeL1State {
           archive,
           oracleInput: { feeAssetPriceModifier: 0n },
         },
-        attestationsAndSigners.getPackedAttestations(),
+        packedAttestations,
         attestationsAndSigners.getSigners().map(signer => signer.toString()),
         attestationsAndSignersSignature.toViemSignature(),
         blobInput,
@@ -579,21 +647,52 @@ export class FakeL1State {
       ],
     });
 
-    return {
+    // Compute attestationsHash (same logic as CalldataRetriever)
+    const attestationsHash = Buffer32.fromString(
+      keccak256(encodeAbiParameters([this.getCommitteeAttestationsStructDef()], [packedAttestations])),
+    );
+
+    // Compute payloadDigest (same logic as CalldataRetriever)
+    const consensusPayload = ConsensusPayload.fromCheckpoint(checkpoint);
+    const payloadToSign = consensusPayload.getPayloadToSign(SignatureDomainSeparator.checkpointAttestation);
+    const payloadDigest = Buffer32.fromString(keccak256(payloadToSign));
+
+    const tx = {
       input: multiCallInput,
       hash: archive,
       blockHash: archive,
       to: MULTI_CALL_3_ADDRESS as `0x${string}`,
     } as Transaction<bigint, number>;
+
+    return { tx, attestationsHash, payloadDigest };
+  }
+
+  /** Extracts the CommitteeAttestations struct definition from RollupAbi for hash computation. */
+  private getCommitteeAttestationsStructDef(): AbiParameter {
+    const proposeFunction = RollupAbi.find(item => item.type === 'function' && item.name === 'propose') as
+      | { type: 'function'; name: string; inputs: readonly AbiParameter[] }
+      | undefined;
+
+    if (!proposeFunction) {
+      throw new Error('propose function not found in RollupAbi');
+    }
+
+    const attestationsParam = proposeFunction.inputs.find(param => param.name === '_attestations');
+    if (!attestationsParam) {
+      throw new Error('_attestations parameter not found in propose function');
+    }
+
+    const tupleParam = attestationsParam as unknown as { type: 'tuple'; components?: readonly AbiParameter[] };
+    return { type: 'tuple', components: tupleParam.components || [] } as AbiParameter;
   }
 
-  private makeVersionedBlobHashes(checkpoint: Checkpoint): `0x${string}`[] {
-    return getBlobsPerL1Block(checkpoint.toBlobFields()).map(
+  private async makeVersionedBlobHashes(checkpoint: Checkpoint): Promise<`0x${string}`[]> {
+    return (await getBlobsPerL1Block(checkpoint.toBlobFields())).map(
       b => `0x${b.getEthVersionedBlobHash().toString('hex')}` as `0x${string}`,
     );
   }
 
-  private makeBlobsFromCheckpoint(checkpoint: Checkpoint): Blob[] {
-    return getBlobsPerL1Block(checkpoint.toBlobFields());
+  private async makeBlobsFromCheckpoint(checkpoint: Checkpoint): Promise<Blob[]> {
+    return await getBlobsPerL1Block(checkpoint.toBlobFields());
   }
 }

package/src/test/mock_archiver.ts

@@ -56,8 +56,9 @@ export class MockPrefilledArchiver extends MockArchiver {
     }
 
     const fromBlock = this.l2Blocks.length;
-    // TODO: Add L2 blocks and checkpoints separately once archiver has the apis for that.
-    this.addProposedBlocks(this.prefilled.slice(fromBlock, fromBlock + numBlocks).flatMap(c => c.blocks));
+    const checkpointsToAdd = this.prefilled.slice(fromBlock, fromBlock + numBlocks);
+    this.addProposedBlocks(checkpointsToAdd.flatMap(c => c.blocks));
+    this.checkpointList.push(...checkpointsToAdd);
     return Promise.resolve();
   }
 }

package/src/test/mock_l2_block_source.ts

@@ -8,6 +8,7 @@ import { createLogger } from '@aztec/foundation/log';
 import type { FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import {
+  type BlockData,
   BlockHash,
   CheckpointedL2Block,
   L2Block,
@@ -15,9 +16,16 @@ import {
   type L2Tips,
   type ValidateCheckpointResult,
 } from '@aztec/stdlib/block';
-import { Checkpoint, L1PublishedData, PublishedCheckpoint } from '@aztec/stdlib/checkpoint';
+import { Checkpoint, type CheckpointData, L1PublishedData, PublishedCheckpoint } from '@aztec/stdlib/checkpoint';
 import type { ContractClassPublic, ContractDataSource, ContractInstanceWithAddress } from '@aztec/stdlib/contract';
-import { EmptyL1RollupConstants, type L1RollupConstants, getSlotRangeForEpoch } from '@aztec/stdlib/epoch-helpers';
+import {
+  EmptyL1RollupConstants,
+  type L1RollupConstants,
+  getEpochAtSlot,
+  getSlotRangeForEpoch,
+} from '@aztec/stdlib/epoch-helpers';
+import { computeCheckpointOutHash } from '@aztec/stdlib/messaging';
+import { CheckpointHeader } from '@aztec/stdlib/rollup';
 import { type BlockHeader, TxExecutionResult, TxHash, TxReceipt, TxStatus } from '@aztec/stdlib/tx';
 import type { UInt64 } from '@aztec/stdlib/types';
 
@@ -26,6 +34,7 @@ import type { UInt64 } from '@aztec/stdlib/types';
  */
 export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
   protected l2Blocks: L2Block[] = [];
+  protected checkpointList: Checkpoint[] = [];
 
   private provenBlockNumber: number = 0;
   private finalizedBlockNumber: number = 0;
@@ -33,14 +42,36 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
 
   private log = createLogger('archiver:mock_l2_block_source');
 
+  /** Creates blocks grouped into single-block checkpoints. */
   public async createBlocks(numBlocks: number) {
-    for (let i = 0; i < numBlocks; i++) {
-      const blockNum = this.l2Blocks.length + 1;
-      const block = await L2Block.random(BlockNumber(blockNum), { slotNumber: SlotNumber(blockNum) });
-      this.l2Blocks.push(block);
+    await this.createCheckpoints(numBlocks, 1);
+  }
+
+  public getCheckpointNumber(): Promise<CheckpointNumber> {
+    return Promise.resolve(
+      this.checkpointList.length === 0 ? CheckpointNumber.ZERO : CheckpointNumber(this.checkpointList.length),
+    );
+  }
+
+  /** Creates checkpoints, each containing `blocksPerCheckpoint` blocks. */
+  public async createCheckpoints(numCheckpoints: number, blocksPerCheckpoint: number = 1) {
+    for (let c = 0; c < numCheckpoints; c++) {
+      const checkpointNum = CheckpointNumber(this.checkpointList.length + 1);
+      const startBlockNum = this.l2Blocks.length + 1;
+      const slotNumber = SlotNumber(Number(checkpointNum));
+      const checkpoint = await Checkpoint.random(checkpointNum, {
+        numBlocks: blocksPerCheckpoint,
+        startBlockNumber: startBlockNum,
+        slotNumber,
+        checkpointNumber: checkpointNum,
+      });
+      this.checkpointList.push(checkpoint);
+      this.l2Blocks.push(...checkpoint.blocks);
     }
 
-    this.log.verbose(`Created ${numBlocks} blocks in the mock L2 block source`);
+    this.log.verbose(
+      `Created ${numCheckpoints} checkpoints with ${blocksPerCheckpoint} blocks each in the mock L2 block source`,
+    );
   }
 
   public addProposedBlocks(blocks: L2Block[]) {
@@ -50,6 +81,16 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
 
   public removeBlocks(numBlocks: number) {
     this.l2Blocks = this.l2Blocks.slice(0, -numBlocks);
+    const maxBlockNum = this.l2Blocks.length;
+    // Remove any checkpoint whose last block is beyond the remaining blocks.
+    this.checkpointList = this.checkpointList.filter(c => {
+      const lastBlockNum = c.blocks[0].number + c.blocks.length - 1;
+      return lastBlockNum <= maxBlockNum;
+    });
+    // Keep tip numbers consistent with remaining blocks.
+    this.checkpointedBlockNumber = Math.min(this.checkpointedBlockNumber, maxBlockNum);
+    this.provenBlockNumber = Math.min(this.provenBlockNumber, maxBlockNum);
+    this.finalizedBlockNumber = Math.min(this.finalizedBlockNumber, maxBlockNum);
     this.log.verbose(`Removed ${numBlocks} blocks from the mock L2 block source`);
   }
 
@@ -65,7 +106,33 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
   }
 
   public setCheckpointedBlockNumber(checkpointedBlockNumber: number) {
+    const prevCheckpointed = this.checkpointedBlockNumber;
     this.checkpointedBlockNumber = checkpointedBlockNumber;
+    // Auto-create single-block checkpoints for newly checkpointed blocks that don't have one yet.
+    // This handles blocks added via addProposedBlocks that are now being marked as checkpointed.
+    const newCheckpoints: Checkpoint[] = [];
+    for (let blockNum = prevCheckpointed + 1; blockNum <= checkpointedBlockNumber; blockNum++) {
+      const block = this.l2Blocks[blockNum - 1];
+      if (!block) {
+        continue;
+      }
+      if (this.checkpointList.some(c => c.blocks.some(b => b.number === block.number))) {
+        continue;
+      }
+      const checkpointNum = CheckpointNumber(this.checkpointList.length + newCheckpoints.length + 1);
+      const checkpoint = new Checkpoint(
+        block.archive,
+        CheckpointHeader.random({ slotNumber: block.header.globalVariables.slotNumber }),
+        [block],
+        checkpointNum,
+      );
+      newCheckpoints.push(checkpoint);
+    }
+    // Insert new checkpoints in order by number.
+    if (newCheckpoints.length > 0) {
+      this.checkpointList.push(...newCheckpoints);
+      this.checkpointList.sort((a, b) => a.number - b.number);
+    }
   }
 
   /**
@@ -112,13 +179,7 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     if (!block) {
       return Promise.resolve(undefined);
     }
-    const checkpointedBlock = new CheckpointedL2Block(
-      CheckpointNumber.fromBlockNumber(number),
-      block,
-      new L1PublishedData(BigInt(number), BigInt(number), `0x${number.toString(16).padStart(64, '0')}`),
-      [],
-    );
-    return Promise.resolve(checkpointedBlock);
+    return Promise.resolve(this.toCheckpointedBlock(block));
   }
 
   public async getCheckpointedBlocks(from: BlockNumber, limit: number): Promise<CheckpointedL2Block[]> {
@@ -167,44 +228,22 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
   }
 
   public getCheckpoints(from: CheckpointNumber, limit: number) {
-    // TODO(mbps): Implement this properly. This only works when we have one block per checkpoint.
-    const blocks = this.l2Blocks.slice(from - 1, from - 1 + limit);
-    return Promise.all(
-      blocks.map(async block => {
-        // Create a checkpoint from the block - manually construct since L2Block doesn't have toCheckpoint()
-        const checkpoint = await Checkpoint.random(block.checkpointNumber, { numBlocks: 1 });
-        checkpoint.blocks = [block];
-        return new PublishedCheckpoint(
-          checkpoint,
-          new L1PublishedData(BigInt(block.number), BigInt(block.number), Buffer32.random().toString()),
-          [],
-        );
-      }),
+    const checkpoints = this.checkpointList.slice(from - 1, from - 1 + limit);
+    return Promise.resolve(
+      checkpoints.map(checkpoint => new PublishedCheckpoint(checkpoint, this.mockL1DataForCheckpoint(checkpoint), [])),
     );
   }
 
-  public async getCheckpointByArchive(archive: Fr): Promise<Checkpoint | undefined> {
-    // TODO(mbps): Implement this properly. This only works when we have one block per checkpoint.
-    const block = this.l2Blocks.find(b => b.archive.root.equals(archive));
-    if (!block) {
-      return undefined;
-    }
-    // Create a checkpoint from the block - manually construct since L2Block doesn't have toCheckpoint()
-    const checkpoint = await Checkpoint.random(block.checkpointNumber, { numBlocks: 1 });
-    checkpoint.blocks = [block];
-    return checkpoint;
+  public getCheckpointByArchive(archive: Fr): Promise<Checkpoint | undefined> {
+    const checkpoint = this.checkpointList.find(c => c.archive.root.equals(archive));
+    return Promise.resolve(checkpoint);
   }
 
   public async getCheckpointedBlockByHash(blockHash: BlockHash): Promise<CheckpointedL2Block | undefined> {
     for (const block of this.l2Blocks) {
       const hash = await block.hash();
       if (hash.equals(blockHash)) {
-        return CheckpointedL2Block.fromFields({
-          checkpointNumber: CheckpointNumber.fromBlockNumber(block.number),
-          block,
-          l1: new L1PublishedData(BigInt(block.number), BigInt(block.number), Buffer32.random().toString()),
-          attestations: [],
-        });
+        return this.toCheckpointedBlock(block);
       }
     }
     return undefined;
@@ -215,14 +254,7 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     if (!block) {
       return Promise.resolve(undefined);
     }
-    return Promise.resolve(
-      CheckpointedL2Block.fromFields({
-        checkpointNumber: CheckpointNumber.fromBlockNumber(block.number),
-        block,
-        l1: new L1PublishedData(BigInt(block.number), BigInt(block.number), Buffer32.random().toString()),
-        attestations: [],
-      }),
-    );
+    return Promise.resolve(this.toCheckpointedBlock(block));
   }
 
   public async getL2BlockByHash(blockHash: BlockHash): Promise<L2Block | undefined> {
@@ -255,47 +287,69 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     return Promise.resolve(block?.header);
   }
 
+  public async getBlockData(number: BlockNumber): Promise<BlockData | undefined> {
+    const block = this.l2Blocks[number - 1];
+    if (!block) {
+      return undefined;
+    }
+    return {
+      header: block.header,
+      archive: block.archive,
+      blockHash: await block.hash(),
+      checkpointNumber: block.checkpointNumber,
+      indexWithinCheckpoint: block.indexWithinCheckpoint,
+    };
+  }
+
+  public async getBlockDataByArchive(archive: Fr): Promise<BlockData | undefined> {
+    const block = this.l2Blocks.find(b => b.archive.root.equals(archive));
+    if (!block) {
+      return undefined;
+    }
+    return {
+      header: block.header,
+      archive: block.archive,
+      blockHash: await block.hash(),
+      checkpointNumber: block.checkpointNumber,
+      indexWithinCheckpoint: block.indexWithinCheckpoint,
+    };
+  }
+
   getBlockHeader(number: number | 'latest'): Promise<BlockHeader | undefined> {
     return Promise.resolve(this.l2Blocks.at(typeof number === 'number' ? number - 1 : -1)?.header);
   }
 
   getCheckpointsForEpoch(epochNumber: EpochNumber): Promise<Checkpoint[]> {
-    // TODO(mbps): Implement this properly. This only works when we have one block per checkpoint.
-    const epochDuration = DefaultL1ContractsConfig.aztecEpochDuration;
-    const [start, end] = getSlotRangeForEpoch(epochNumber, { epochDuration });
-    const blocks = this.l2Blocks.filter(b => {
-      const slot = b.header.globalVariables.slotNumber;
-      return slot >= start && slot <= end;
-    });
-    // Create checkpoints from blocks - manually construct since L2Block doesn't have toCheckpoint()
-    return Promise.all(
-      blocks.map(async block => {
-        const checkpoint = await Checkpoint.random(block.checkpointNumber, { numBlocks: 1 });
-        checkpoint.blocks = [block];
-        return checkpoint;
-      }),
-    );
+    return Promise.resolve(this.getCheckpointsInEpoch(epochNumber));
   }
 
-  getCheckpointedBlocksForEpoch(epochNumber: EpochNumber): Promise<CheckpointedL2Block[]> {
-    const epochDuration = DefaultL1ContractsConfig.aztecEpochDuration;
-    const [start, end] = getSlotRangeForEpoch(epochNumber, { epochDuration });
-    const blocks = this.l2Blocks.filter(b => {
-      const slot = b.header.globalVariables.slotNumber;
-      return slot >= start && slot <= end;
-    });
+  getCheckpointsDataForEpoch(epochNumber: EpochNumber): Promise<CheckpointData[]> {
+    const checkpoints = this.getCheckpointsInEpoch(epochNumber);
     return Promise.resolve(
-      blocks.map(block =>
-        CheckpointedL2Block.fromFields({
-          checkpointNumber: CheckpointNumber.fromBlockNumber(block.number),
-          block,
-          l1: new L1PublishedData(BigInt(block.number), BigInt(block.number), Buffer32.random().toString()),
+      checkpoints.map(
+        (checkpoint): CheckpointData => ({
+          checkpointNumber: checkpoint.number,
+          header: checkpoint.header,
+          archive: checkpoint.archive,
+          checkpointOutHash: computeCheckpointOutHash(
+            checkpoint.blocks.map(b => b.body.txEffects.map(tx => tx.l2ToL1Msgs)),
+          ),
+          startBlock: checkpoint.blocks[0].number,
+          blockCount: checkpoint.blocks.length,
           attestations: [],
+          l1: this.mockL1DataForCheckpoint(checkpoint),
         }),
       ),
     );
   }
 
+  getCheckpointedBlocksForEpoch(epochNumber: EpochNumber): Promise<CheckpointedL2Block[]> {
+    const checkpoints = this.getCheckpointsInEpoch(epochNumber);
+    return Promise.resolve(
+      checkpoints.flatMap(checkpoint => checkpoint.blocks.map(block => this.toCheckpointedBlock(block))),
+    );
+  }
+
   getBlocksForSlot(slotNumber: SlotNumber): Promise<L2Block[]> {
     const blocks = this.l2Blocks.filter(b => b.header.globalVariables.slotNumber === slotNumber);
     return Promise.resolve(blocks);
@@ -345,6 +399,7 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
             txEffect.transactionFee.toBigInt(),
             await block.hash(),
             block.number,
+            getEpochAtSlot(block.slot, EmptyL1RollupConstants),
           );
         }
       }
@@ -384,7 +439,10 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
 
     const makeTipId = (blockId: typeof latestBlockId) => ({
       block: blockId,
-      checkpoint: { number: CheckpointNumber.fromBlockNumber(blockId.number), hash: blockId.hash },
+      checkpoint: {
+        number: this.findCheckpointNumberForBlock(blockId.number) ?? CheckpointNumber(0),
+        hash: blockId.hash,
+      },
     });
 
     return {
@@ -395,11 +453,11 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     };
   }
 
-  getL2EpochNumber(): Promise<EpochNumber> {
+  getSyncedL2EpochNumber(): Promise<EpochNumber> {
     throw new Error('Method not implemented.');
   }
 
-  getL2SlotNumber(): Promise<SlotNumber> {
+  getSyncedL2SlotNumber(): Promise<SlotNumber> {
     throw new Error('Method not implemented.');
   }
 
@@ -472,4 +530,38 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
   getPendingChainValidationStatus(): Promise<ValidateCheckpointResult> {
     return Promise.resolve({ valid: true });
   }
+
+  /** Returns checkpoints whose slot falls within the given epoch. */
+  private getCheckpointsInEpoch(epochNumber: EpochNumber): Checkpoint[] {
+    const epochDuration = DefaultL1ContractsConfig.aztecEpochDuration;
+    const [start, end] = getSlotRangeForEpoch(epochNumber, { epochDuration });
+    return this.checkpointList.filter(c => c.header.slotNumber >= start && c.header.slotNumber <= end);
+  }
+
+  /** Creates a mock L1PublishedData for a checkpoint. */
+  private mockL1DataForCheckpoint(checkpoint: Checkpoint): L1PublishedData {
+    return new L1PublishedData(BigInt(checkpoint.number), BigInt(checkpoint.number), Buffer32.random().toString());
+  }
+
+  /** Creates a CheckpointedL2Block from a block using stored checkpoint info. */
+  private toCheckpointedBlock(block: L2Block): CheckpointedL2Block {
+    const checkpoint = this.checkpointList.find(c => c.blocks.some(b => b.number === block.number));
+    const checkpointNumber = checkpoint?.number ?? block.checkpointNumber;
+    return new CheckpointedL2Block(
+      checkpointNumber,
+      block,
+      new L1PublishedData(
+        BigInt(block.number),
+        BigInt(block.number),
+        `0x${block.number.toString(16).padStart(64, '0')}`,
+      ),
+      [],
+    );
+  }
+
+  /** Finds the checkpoint number for a block, or undefined if the block is not in any checkpoint. */
+  private findCheckpointNumberForBlock(blockNumber: BlockNumber): CheckpointNumber | undefined {
+    const checkpoint = this.checkpointList.find(c => c.blocks.some(b => b.number === blockNumber));
+    return checkpoint?.number;
+  }
 }