@aztec/archiver 0.0.1-commit.3469e52 → 0.0.1-commit.381b1a9

package/src/test/mock_l2_block_source.ts

@@ -8,24 +8,33 @@ 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,
-  L2BlockHash,
-  L2BlockNew,
+  L2Block,
   type L2BlockSource,
   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 { type BlockHeader, TxHash, TxReceipt, TxStatus } from '@aztec/stdlib/tx';
+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';
 
 /**
  * A mocked implementation of L2BlockSource to be used in tests.
  */
 export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
-  protected l2Blocks: L2BlockNew[] = [];
+  protected l2Blocks: L2Block[] = [];
+  protected checkpointList: Checkpoint[] = [];
 
   private provenBlockNumber: number = 0;
   private finalizedBlockNumber: number = 0;
@@ -33,23 +42,55 @@ 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 L2BlockNew.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 addBlocks(blocks: L2BlockNew[]) {
+  public addProposedBlocks(blocks: L2Block[]) {
     this.l2Blocks.push(...blocks);
-    this.log.verbose(`Added ${blocks.length} blocks to the mock L2 block source`);
+    this.log.verbose(`Added ${blocks.length} proposed blocks to the mock L2 block source`);
   }
 
   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);
+    }
   }
 
   /**
@@ -96,6 +163,14 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     return Promise.resolve(BlockNumber(this.provenBlockNumber));
   }
 
+  public getCheckpointedL2BlockNumber() {
+    return Promise.resolve(BlockNumber(this.checkpointedBlockNumber));
+  }
+
+  public getFinalizedL2BlockNumber() {
+    return Promise.resolve(BlockNumber(this.finalizedBlockNumber));
+  }
+
   public getCheckpointedBlock(number: BlockNumber): Promise<CheckpointedL2Block | undefined> {
     if (number > this.checkpointedBlockNumber) {
       return Promise.resolve(undefined);
@@ -104,20 +179,10 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     if (!block) {
       return Promise.resolve(undefined);
     }
-    const checkpointedBlock = new CheckpointedL2Block(
-      CheckpointNumber(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,
-    _proven?: boolean,
-  ): Promise<CheckpointedL2Block[]> {
+  public async getCheckpointedBlocks(from: BlockNumber, limit: number): Promise<CheckpointedL2Block[]> {
     const result: CheckpointedL2Block[] = [];
     for (let i = 0; i < limit; i++) {
       const blockNum = from + i;
@@ -137,7 +202,7 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
    * @param number - The block number to return (inclusive).
    * @returns The requested L2 block.
    */
-  public getBlock(number: number): Promise<L2BlockNew | undefined> {
+  public getBlock(number: number): Promise<L2Block | undefined> {
     const block = this.l2Blocks[number - 1];
     return Promise.resolve(block);
   }
@@ -147,7 +212,7 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
    * @param number - The block number to return.
    * @returns The requested L2 block.
    */
-  public getL2BlockNew(number: BlockNumber): Promise<L2BlockNew | undefined> {
+  public getL2Block(number: BlockNumber): Promise<L2Block | undefined> {
     const block = this.l2Blocks[number - 1];
     return Promise.resolve(block);
   }
@@ -158,95 +223,41 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
    * @param limit - The maximum number of blocks to return.
    * @returns The requested mocked L2 blocks.
    */
-  public getBlocks(from: number, limit: number, proven?: boolean): Promise<L2BlockNew[]> {
-    return Promise.resolve(
-      this.l2Blocks
-        .slice(from - 1, from - 1 + limit)
-        .filter(b => !proven || this.provenBlockNumber === undefined || b.number <= this.provenBlockNumber),
-    );
-  }
-
-  public getPublishedCheckpoints(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 L2BlockNew 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()),
-          [],
-        );
-      }),
-    );
-  }
-
-  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 L2BlockNew doesn't have toCheckpoint()
-    const checkpoint = await Checkpoint.random(block.checkpointNumber, { numBlocks: 1 });
-    checkpoint.blocks = [block];
-    return checkpoint;
+  public getBlocks(from: number, limit: number): Promise<L2Block[]> {
+    return Promise.resolve(this.l2Blocks.slice(from - 1, from - 1 + limit));
   }
 
-  public getPublishedBlocks(from: number, limit: number, proven?: boolean): Promise<CheckpointedL2Block[]> {
-    const blocks = this.l2Blocks
-      .slice(from - 1, from - 1 + limit)
-      .filter(b => !proven || this.provenBlockNumber === undefined || b.number <= this.provenBlockNumber);
+  public getCheckpoints(from: CheckpointNumber, limit: number) {
+    const checkpoints = this.checkpointList.slice(from - 1, from - 1 + limit);
     return Promise.resolve(
-      blocks.map(block =>
-        CheckpointedL2Block.fromFields({
-          checkpointNumber: CheckpointNumber(block.number),
-          block,
-          l1: new L1PublishedData(BigInt(block.number), BigInt(block.number), Buffer32.random().toString()),
-          attestations: [],
-        }),
-      ),
+      checkpoints.map(checkpoint => new PublishedCheckpoint(checkpoint, this.mockL1DataForCheckpoint(checkpoint), [])),
     );
   }
 
-  getL2BlocksNew(from: BlockNumber, limit: number, proven?: boolean): Promise<L2BlockNew[]> {
-    // getBlocks already returns L2BlockNew[], so just return directly
-    return this.getBlocks(from, limit, proven);
+  public getCheckpointByArchive(archive: Fr): Promise<Checkpoint | undefined> {
+    const checkpoint = this.checkpointList.find(c => c.archive.root.equals(archive));
+    return Promise.resolve(checkpoint);
   }
 
-  public async getPublishedBlockByHash(blockHash: Fr): Promise<CheckpointedL2Block | undefined> {
+  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(block.number),
-          block,
-          l1: new L1PublishedData(BigInt(block.number), BigInt(block.number), Buffer32.random().toString()),
-          attestations: [],
-        });
+        return this.toCheckpointedBlock(block);
       }
     }
     return undefined;
   }
 
-  public getPublishedBlockByArchive(archive: Fr): Promise<CheckpointedL2Block | undefined> {
+  public getCheckpointedBlockByArchive(archive: Fr): Promise<CheckpointedL2Block | undefined> {
     const block = this.l2Blocks.find(b => b.archive.root.equals(archive));
     if (!block) {
       return Promise.resolve(undefined);
     }
-    return Promise.resolve(
-      CheckpointedL2Block.fromFields({
-        checkpointNumber: CheckpointNumber(block.number),
-        block,
-        l1: new L1PublishedData(BigInt(block.number), BigInt(block.number), Buffer32.random().toString()),
-        attestations: [],
-      }),
-    );
+    return Promise.resolve(this.toCheckpointedBlock(block));
   }
 
-  public async getL2BlockNewByHash(blockHash: Fr): Promise<L2BlockNew | undefined> {
+  public async getL2BlockByHash(blockHash: BlockHash): Promise<L2Block | undefined> {
     for (const block of this.l2Blocks) {
       const hash = await block.hash();
       if (hash.equals(blockHash)) {
@@ -256,12 +267,12 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     return undefined;
   }
 
-  public getL2BlockNewByArchive(archive: Fr): Promise<L2BlockNew | undefined> {
+  public getL2BlockByArchive(archive: Fr): Promise<L2Block | undefined> {
     const block = this.l2Blocks.find(b => b.archive.root.equals(archive));
     return Promise.resolve(block);
   }
 
-  public async getBlockHeaderByHash(blockHash: Fr): Promise<BlockHeader | undefined> {
+  public async getBlockHeaderByHash(blockHash: BlockHash): Promise<BlockHeader | undefined> {
     for (const block of this.l2Blocks) {
       const hash = await block.hash();
       if (hash.equals(blockHash)) {
@@ -276,46 +287,77 @@ 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 L2BlockNew 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));
+  }
+
+  getCheckpointsDataForEpoch(epochNumber: EpochNumber): Promise<CheckpointData[]> {
+    const checkpoints = this.getCheckpointsInEpoch(epochNumber);
+    return Promise.resolve(
+      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),
+        }),
+      ),
     );
   }
 
-  getBlocksForEpoch(epochNumber: EpochNumber): Promise<L2BlockNew[]> {
-    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;
-    });
-    return Promise.resolve(blocks);
+  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<L2BlockNew[]> {
+  getBlocksForSlot(slotNumber: SlotNumber): Promise<L2Block[]> {
     const blocks = this.l2Blocks.filter(b => b.header.globalVariables.slotNumber === slotNumber);
     return Promise.resolve(blocks);
   }
 
-  async getBlockHeadersForEpoch(epochNumber: EpochNumber): Promise<BlockHeader[]> {
-    const blocks = await this.getBlocksForEpoch(epochNumber);
-    return blocks.map(b => b.header);
+  async getCheckpointedBlockHeadersForEpoch(epochNumber: EpochNumber): Promise<BlockHeader[]> {
+    const checkpointedBlocks = await this.getCheckpointedBlocksForEpoch(epochNumber);
+    return checkpointedBlocks.map(b => b.block.header);
   }
 
   /**
@@ -334,7 +376,7 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     return {
       data: txEffect,
       l2BlockNumber: block.number,
-      l2BlockHash: L2BlockHash.fromField(await block.hash()),
+      l2BlockHash: await block.hash(),
       txIndexInBlock: block.body.txEffects.indexOf(txEffect),
     };
   }
@@ -348,13 +390,16 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
     for (const block of this.l2Blocks) {
       for (const txEffect of block.body.txEffects) {
         if (txEffect.txHash.equals(txHash)) {
+          // In mock, assume all txs are checkpointed with successful execution
           return new TxReceipt(
             txHash,
-            TxStatus.SUCCESS,
-            '',
+            TxStatus.CHECKPOINTED,
+            TxExecutionResult.SUCCESS,
+            undefined,
             txEffect.transactionFee.toBigInt(),
-            L2BlockHash.fromField(await block.hash()),
+            await block.hash(),
             block.number,
+            getEpochAtSlot(block.slot, EmptyL1RollupConstants),
           );
         }
       }
@@ -394,7 +439,10 @@ export class MockL2BlockSource implements L2BlockSource, ContractDataSource {
 
     const makeTipId = (blockId: typeof latestBlockId) => ({
       block: blockId,
-      checkpoint: { number: CheckpointNumber(blockId.number), hash: blockId.hash },
+      checkpoint: {
+        number: this.findCheckpointNumberForBlock(blockId.number) ?? CheckpointNumber(0),
+        hash: blockId.hash,
+      },
     });
 
     return {
@@ -405,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.');
   }
 
@@ -482,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;
+  }
 }

package/src/test/mock_structs.ts

@@ -12,7 +12,7 @@ import type { Secp256k1Signer } from '@aztec/foundation/crypto/secp256k1-signer'
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { EthAddress } from '@aztec/foundation/eth-address';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { CommitteeAttestation, L2BlockNew } from '@aztec/stdlib/block';
+import { CommitteeAttestation, L2Block } from '@aztec/stdlib/block';
 import { Checkpoint, L1PublishedData, PublishedCheckpoint } from '@aztec/stdlib/checkpoint';
 import { PrivateLog, PublicLog, SiloedTag, Tag } from '@aztec/stdlib/logs';
 import { InboxLeaf } from '@aztec/stdlib/messaging';
@@ -46,24 +46,40 @@ export function makeInboxMessage(
 }
 
 export function makeInboxMessages(
-  count: number,
+  totalCount: number,
   opts: {
     initialHash?: Buffer16;
     initialCheckpointNumber?: CheckpointNumber;
+    messagesPerCheckpoint?: number;
     overrideFn?: (msg: InboxMessage, index: number) => InboxMessage;
   } = {},
 ): InboxMessage[] {
-  const { initialHash = Buffer16.ZERO, overrideFn = msg => msg, initialCheckpointNumber = 1 } = opts;
+  const {
+    initialHash = Buffer16.ZERO,
+    overrideFn = msg => msg,
+    initialCheckpointNumber = CheckpointNumber(1),
+    messagesPerCheckpoint = 1,
+  } = opts;
+
   const messages: InboxMessage[] = [];
   let rollingHash = initialHash;
-  for (let i = 0; i < count; i++) {
+  for (let i = 0; i < totalCount; i++) {
+    const msgIndex = i % messagesPerCheckpoint;
+    const checkpointNumber = CheckpointNumber.fromBigInt(
+      BigInt(initialCheckpointNumber) + BigInt(i) / BigInt(messagesPerCheckpoint),
+    );
     const leaf = Fr.random();
-    const checkpointNumber = CheckpointNumber(i + initialCheckpointNumber);
-    const message = overrideFn(makeInboxMessage(rollingHash, { leaf, checkpointNumber }), i);
+    const message = overrideFn(
+      makeInboxMessage(rollingHash, {
+        leaf,
+        checkpointNumber,
+        index: InboxLeaf.smallestIndexForCheckpoint(checkpointNumber) + BigInt(msgIndex),
+      }),
+      i,
+    );
     rollingHash = message.rollingHash;
     messages.push(message);
   }
-
   return messages;
 }
 
@@ -111,6 +127,25 @@ export function makeL1PublishedData(l1BlockNumber: number): L1PublishedData {
   return new L1PublishedData(BigInt(l1BlockNumber), BigInt(l1BlockNumber * 1000), makeBlockHash(l1BlockNumber));
 }
 
+/** Creates a Checkpoint from a list of blocks with a header that matches the blocks' structure. */
+export function makeCheckpoint(blocks: L2Block[], checkpointNumber = CheckpointNumber(1)): Checkpoint {
+  const firstBlock = blocks[0];
+  const { slotNumber, timestamp, coinbase, feeRecipient, gasFees } = firstBlock.header.globalVariables;
+  return new Checkpoint(
+    blocks.at(-1)!.archive,
+    CheckpointHeader.random({
+      lastArchiveRoot: firstBlock.header.lastArchive.root,
+      slotNumber,
+      timestamp,
+      coinbase,
+      feeRecipient,
+      gasFees,
+    }),
+    blocks,
+    checkpointNumber,
+  );
+}
+
 /** Wraps a Checkpoint with L1 published data and random attestations. */
 export function makePublishedCheckpoint(
   checkpoint: Checkpoint,
@@ -268,8 +303,8 @@ export async function makeCheckpointWithLogs(
 ): Promise<PublishedCheckpoint> {
   const { previousArchive, numTxsPerBlock = 4, privateLogs, publicLogs } = options;
 
-  const block = await L2BlockNew.random(BlockNumber(blockNumber), {
-    checkpointNumber: CheckpointNumber(blockNumber),
+  const block = await L2Block.random(BlockNumber(blockNumber), {
+    checkpointNumber: CheckpointNumber.fromBlockNumber(BlockNumber(blockNumber)),
     indexWithinCheckpoint: IndexWithinCheckpoint(0),
     state: makeStateForBlock(blockNumber, numTxsPerBlock),
     ...(previousArchive ? { lastArchive: previousArchive } : {}),
@@ -285,11 +320,6 @@ export async function makeCheckpointWithLogs(
     return txEffect;
   });
 
-  const checkpoint = new Checkpoint(
-    AppendOnlyTreeSnapshot.random(),
-    CheckpointHeader.random(),
-    [block],
-    CheckpointNumber(blockNumber),
-  );
+  const checkpoint = makeCheckpoint([block], CheckpointNumber.fromBlockNumber(BlockNumber(blockNumber)));
   return makePublishedCheckpoint(checkpoint, blockNumber);
 }