@aztec/archiver 0.56.0 → 0.58.0

package/src/archiver/kv_archiver_store/block_store.ts

@@ -26,6 +26,12 @@ export class BlockStore {
   /** Stores L1 block number in which the last processed L2 block was included */
   #lastSynchedL1Block: AztecSingleton<bigint>;
 
+  /** Stores l2 block number of the last proven block */
+  #lastProvenL2Block: AztecSingleton<number>;
+
+  /** Stores l2 epoch number of the last proven epoch */
+  #lastProvenL2Epoch: AztecSingleton<number>;
+
   /** Index mapping transaction hash (as a string) to its location in a block */
   #txIndex: AztecMap<string, BlockIndexValue>;
 
@@ -40,6 +46,8 @@ export class BlockStore {
     this.#txIndex = db.openMap('archiver_tx_index');
     this.#contractIndex = db.openMap('archiver_contract_index');
     this.#lastSynchedL1Block = db.openSingleton('archiver_last_synched_l1_block');
+    this.#lastProvenL2Block = db.openSingleton('archiver_last_proven_l2_block');
+    this.#lastProvenL2Epoch = db.openSingleton('archiver_last_proven_l2_epoch');
   }
 
   /**
@@ -73,6 +81,38 @@ export class BlockStore {
     });
   }
 
+  /**
+   * Unwinds blocks from the database
+   * @param from -  The tip of the chain, passed for verification purposes,
+   *                ensuring that we don't end up deleting something we did not intend
+   * @param blocksToUnwind - The number of blocks we are to unwind
+   * @returns True if the operation is successful
+   */
+  unwindBlocks(from: number, blocksToUnwind: number) {
+    return this.db.transaction(() => {
+      const last = this.getSynchedL2BlockNumber();
+      if (from != last) {
+        throw new Error(`Can only remove from the tip`);
+      }
+
+      for (let i = 0; i < blocksToUnwind; i++) {
+        const blockNumber = from - i;
+        const block = this.getBlock(blockNumber);
+
+        if (block === undefined) {
+          throw new Error(`Cannot remove block ${blockNumber} from the store, we don't have it`);
+        }
+        void this.#blocks.delete(block.data.number);
+        block.data.body.txEffects.forEach(tx => {
+          void this.#txIndex.delete(tx.txHash.toString());
+        });
+        void this.#blockBodies.delete(block.data.body.getTxsEffectsHash().toString('hex'));
+      }
+
+      return true;
+    });
+  }
+
   /**
    * Gets up to `limit` amount of L2 blocks starting from `from`.
    * @param start - Number of the first block to return (inclusive).
@@ -99,6 +139,18 @@ export class BlockStore {
     return this.getBlockFromBlockStorage(blockStorage);
   }
 
+  /**
+   * Gets the headers for a sequence of L2 blocks.
+   * @param start - Number of the first block to return (inclusive).
+   * @param limit - The number of blocks to return.
+   * @returns The requested L2 block headers
+   */
+  *getBlockHeaders(start: number, limit: number): IterableIterator<Header> {
+    for (const blockStorage of this.#blocks.values(this.#computeBlockRange(start, limit))) {
+      yield Header.fromBuffer(blockStorage.header);
+    }
+  }
+
   private getBlockFromBlockStorage(blockStorage: BlockStorage) {
     const header = Header.fromBuffer(blockStorage.header);
     const archive = AppendOnlyTreeSnapshot.fromBuffer(blockStorage.archive);
@@ -109,7 +161,7 @@ export class BlockStore {
     }
     const body = Body.fromBuffer(blockBodyBuffer);
 
-    const l2Block = L2Block.fromFields({ header, archive, body });
+    const l2Block = new L2Block(archive, header, body);
     return { data: l2Block, l1: blockStorage.l1 };
   }
 
@@ -191,13 +243,29 @@ export class BlockStore {
     void this.#lastSynchedL1Block.set(l1BlockNumber);
   }
 
+  getProvenL2BlockNumber(): number {
+    return this.#lastProvenL2Block.get() ?? 0;
+  }
+
+  setProvenL2BlockNumber(blockNumber: number) {
+    void this.#lastProvenL2Block.set(blockNumber);
+  }
+
+  getProvenL2EpochNumber(): number | undefined {
+    return this.#lastProvenL2Epoch.get();
+  }
+
+  setProvenL2EpochNumber(epochNumber: number) {
+    void this.#lastProvenL2Epoch.set(epochNumber);
+  }
+
   #computeBlockRange(start: number, limit: number): Required<Pick<Range<number>, 'start' | 'end'>> {
     if (limit < 1) {
       throw new Error(`Invalid limit: ${limit}`);
     }
 
     if (start < INITIAL_L2_BLOCK_NUM) {
-      start = INITIAL_L2_BLOCK_NUM;
+      throw new Error(`Invalid start: ${start}`);
     }
 
     const end = start + limit;

package/src/archiver/kv_archiver_store/contract_class_store.ts

@@ -1,11 +1,14 @@
-import { Fr, FunctionSelector, Vector } from '@aztec/circuits.js';
-import { BufferReader, numToUInt8, serializeToBuffer } from '@aztec/foundation/serialize';
-import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
 import {
   type ContractClassPublic,
+  type ContractClassPublicWithBlockNumber,
   type ExecutablePrivateFunctionWithMembershipProof,
+  Fr,
+  FunctionSelector,
   type UnconstrainedFunctionWithMembershipProof,
-} from '@aztec/types/contracts';
+  Vector,
+} from '@aztec/circuits.js';
+import { BufferReader, numToUInt8, serializeToBuffer } from '@aztec/foundation/serialize';
+import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
 
 /**
  * LMDB implementation of the ArchiverDataStore interface.
@@ -17,8 +20,18 @@ export class ContractClassStore {
     this.#contractClasses = db.openMap('archiver_contract_classes');
   }
 
-  addContractClass(contractClass: ContractClassPublic): Promise<void> {
-    return this.#contractClasses.set(contractClass.id.toString(), serializeContractClassPublic(contractClass));
+  async addContractClass(contractClass: ContractClassPublic, blockNumber: number): Promise<void> {
+    await this.#contractClasses.setIfNotExists(
+      contractClass.id.toString(),
+      serializeContractClassPublic({ ...contractClass, l2BlockNumber: blockNumber }),
+    );
+  }
+
+  async deleteContractClasses(contractClass: ContractClassPublic, blockNumber: number): Promise<void> {
+    const restoredContractClass = this.#contractClasses.get(contractClass.id.toString());
+    if (restoredContractClass && deserializeContractClassPublic(restoredContractClass).l2BlockNumber >= blockNumber) {
+      await this.#contractClasses.delete(contractClass.id.toString());
+    }
   }
 
   getContractClass(id: Fr): ContractClassPublic | undefined {
@@ -44,7 +57,7 @@ export class ContractClassStore {
       const existingClass = deserializeContractClassPublic(existingClassBuffer);
       const { privateFunctions: existingPrivateFns, unconstrainedFunctions: existingUnconstrainedFns } = existingClass;
 
-      const updatedClass: Omit<ContractClassPublic, 'id'> = {
+      const updatedClass: Omit<ContractClassPublicWithBlockNumber, 'id'> = {
         ...existingClass,
         privateFunctions: [
           ...existingPrivateFns,
@@ -63,8 +76,9 @@ export class ContractClassStore {
   }
 }
 
-function serializeContractClassPublic(contractClass: Omit<ContractClassPublic, 'id'>): Buffer {
+function serializeContractClassPublic(contractClass: Omit<ContractClassPublicWithBlockNumber, 'id'>): Buffer {
   return serializeToBuffer(
+    contractClass.l2BlockNumber,
     numToUInt8(contractClass.version),
     contractClass.artifactHash,
     contractClass.publicFunctions.length,
@@ -108,9 +122,10 @@ function serializeUnconstrainedFunction(fn: UnconstrainedFunctionWithMembershipP
   );
 }
 
-function deserializeContractClassPublic(buffer: Buffer): Omit<ContractClassPublic, 'id'> {
+function deserializeContractClassPublic(buffer: Buffer): Omit<ContractClassPublicWithBlockNumber, 'id'> {
   const reader = BufferReader.asReader(buffer);
   return {
+    l2BlockNumber: reader.readNumber(),
     version: reader.readUInt8() as 1,
     artifactHash: reader.readObject(Fr),
     publicFunctions: reader.readVector({

package/src/archiver/kv_archiver_store/contract_instance_store.ts

@@ -1,6 +1,5 @@
-import { type AztecAddress } from '@aztec/circuits.js';
+import { type AztecAddress, type ContractInstanceWithAddress, SerializableContractInstance } from '@aztec/circuits.js';
 import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
-import { type ContractInstanceWithAddress, SerializableContractInstance } from '@aztec/types/contracts';
 
 /**
  * LMDB implementation of the ArchiverDataStore interface.
@@ -19,6 +18,10 @@ export class ContractInstanceStore {
     );
   }
 
+  deleteContractInstance(contractInstance: ContractInstanceWithAddress): Promise<void> {
+    return this.#contractInstances.delete(contractInstance.address.toString());
+  }
+
   getContractInstance(address: AztecAddress): ContractInstanceWithAddress | undefined {
     const contractInstance = this.#contractInstances.get(address.toString());
     return contractInstance && SerializableContractInstance.fromBuffer(contractInstance).withAddress(address);

package/src/archiver/kv_archiver_store/kv_archiver_store.ts

@@ -1,6 +1,4 @@
 import {
-  type EncryptedL2BlockL2Logs,
-  type EncryptedNoteL2BlockL2Logs,
   type FromLogType,
   type GetUnencryptedLogsResponse,
   type InboxLeaf,
@@ -11,22 +9,22 @@ import {
   type TxEffect,
   type TxHash,
   type TxReceipt,
-  type UnencryptedL2BlockL2Logs,
 } from '@aztec/circuit-types';
-import { type Fr } from '@aztec/circuits.js';
-import { type ContractArtifact } from '@aztec/foundation/abi';
-import { type AztecAddress } from '@aztec/foundation/aztec-address';
-import { createDebugLogger } from '@aztec/foundation/log';
-import { type AztecKVStore } from '@aztec/kv-store';
 import {
   type ContractClassPublic,
   type ContractInstanceWithAddress,
   type ExecutablePrivateFunctionWithMembershipProof,
+  type Fr,
+  type Header,
   type UnconstrainedFunctionWithMembershipProof,
-} from '@aztec/types/contracts';
+} from '@aztec/circuits.js';
+import { type ContractArtifact } from '@aztec/foundation/abi';
+import { type AztecAddress } from '@aztec/foundation/aztec-address';
+import { createDebugLogger } from '@aztec/foundation/log';
+import { type AztecKVStore } from '@aztec/kv-store';
 
 import { type ArchiverDataStore, type ArchiverL1SynchPoint } from '../archiver_store.js';
-import { type DataRetrieval, type SingletonDataRetrieval } from '../structs/data_retrieval.js';
+import { type DataRetrieval } from '../structs/data_retrieval.js';
 import { type L1Published } from '../structs/published.js';
 import { BlockStore } from './block_store.js';
 import { ContractArtifactsStore } from './contract_artifacts_store.js';
@@ -34,14 +32,12 @@ import { ContractClassStore } from './contract_class_store.js';
 import { ContractInstanceStore } from './contract_instance_store.js';
 import { LogStore } from './log_store.js';
 import { MessageStore } from './message_store.js';
-import { ProvenStore } from './proven_store.js';
 
 /**
  * LMDB implementation of the ArchiverDataStore interface.
  */
 export class KVArchiverDataStore implements ArchiverDataStore {
   #blockStore: BlockStore;
-  #provenStore: ProvenStore;
   #logStore: LogStore;
   #messageStore: MessageStore;
   #contractClassStore: ContractClassStore;
@@ -52,7 +48,6 @@ export class KVArchiverDataStore implements ArchiverDataStore {
 
   constructor(db: AztecKVStore, logsMaxPageSize: number = 1000) {
     this.#blockStore = new BlockStore(db);
-    this.#provenStore = new ProvenStore(db);
     this.#logStore = new LogStore(db, this.#blockStore, logsMaxPageSize);
     this.#messageStore = new MessageStore(db);
     this.#contractClassStore = new ContractClassStore(db);
@@ -80,8 +75,14 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     return Promise.resolve(this.#contractInstanceStore.getContractInstance(address));
   }
 
-  async addContractClasses(data: ContractClassPublic[], _blockNumber: number): Promise<boolean> {
-    return (await Promise.all(data.map(c => this.#contractClassStore.addContractClass(c)))).every(Boolean);
+  async addContractClasses(data: ContractClassPublic[], blockNumber: number): Promise<boolean> {
+    return (await Promise.all(data.map(c => this.#contractClassStore.addContractClass(c, blockNumber)))).every(Boolean);
+  }
+
+  async deleteContractClasses(data: ContractClassPublic[], blockNumber: number): Promise<boolean> {
+    return (await Promise.all(data.map(c => this.#contractClassStore.deleteContractClasses(c, blockNumber)))).every(
+      Boolean,
+    );
   }
 
   addFunctions(
@@ -96,6 +97,10 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     return (await Promise.all(data.map(c => this.#contractInstanceStore.addContractInstance(c)))).every(Boolean);
   }
 
+  async deleteContractInstances(data: ContractInstanceWithAddress[], _blockNumber: number): Promise<boolean> {
+    return (await Promise.all(data.map(c => this.#contractInstanceStore.deleteContractInstance(c)))).every(Boolean);
+  }
+
   /**
    * Append new blocks to the store's list.
    * @param blocks - The L2 blocks to be added to the store and the last processed L1 block.
@@ -105,6 +110,17 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     return this.#blockStore.addBlocks(blocks);
   }
 
+  /**
+   * Unwinds blocks from the database
+   * @param from -  The tip of the chain, passed for verification purposes,
+   *                ensuring that we don't end up deleting something we did not intend
+   * @param blocksToUnwind - The number of blocks we are to unwind
+   * @returns True if the operation is successful
+   */
+  unwindBlocks(from: number, blocksToUnwind: number): Promise<boolean> {
+    return this.#blockStore.unwindBlocks(from, blocksToUnwind);
+  }
+
   /**
    * Gets up to `limit` amount of L2 blocks starting from `from`.
    *
@@ -121,6 +137,22 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     }
   }
 
+  /**
+   * Gets up to `limit` amount of L2 blocks headers starting from `from`.
+   *
+   * @param start - Number of the first block to return (inclusive).
+   * @param limit - The number of blocks to return.
+   * @returns The requested L2 blocks
+   */
+  getBlockHeaders(start: number, limit: number): Promise<Header[]> {
+    try {
+      return Promise.resolve(Array.from(this.#blockStore.getBlockHeaders(start, limit)));
+    } catch (err) {
+      // this function is sync so if any errors are thrown we need to make sure they're passed on as rejected Promises
+      return Promise.reject(err);
+    }
+  }
+
   /**
    * Gets a tx effect.
    * @param txHash - The txHash of the tx corresponding to the tx effect.
@@ -141,18 +173,19 @@ export class KVArchiverDataStore implements ArchiverDataStore {
 
   /**
    * Append new logs to the store's list.
-   * @param encryptedLogs - The logs to be added to the store.
-   * @param unencryptedLogs - The type of the logs to be added to the store.
-   * @param blockNumber - The block for which to add the logs.
+   * @param blocks - The blocks for which to add the logs.
    * @returns True if the operation is successful.
    */
-  addLogs(
-    noteEncryptedLogs: EncryptedNoteL2BlockL2Logs | undefined,
-    encryptedLogs: EncryptedL2BlockL2Logs | undefined,
-    unencryptedLogs: UnencryptedL2BlockL2Logs | undefined,
-    blockNumber: number,
-  ): Promise<boolean> {
-    return this.#logStore.addLogs(noteEncryptedLogs, encryptedLogs, unencryptedLogs, blockNumber);
+  addLogs(blocks: L2Block[]): Promise<boolean> {
+    return this.#logStore.addLogs(blocks);
+  }
+
+  deleteLogs(blocks: L2Block[]): Promise<boolean> {
+    return this.#logStore.deleteLogs(blocks);
+  }
+
+  getTotalL1ToL2MessageCount(): Promise<bigint> {
+    return Promise.resolve(this.#messageStore.getTotalL1ToL2MessageCount());
   }
 
   /**
@@ -228,11 +261,21 @@ export class KVArchiverDataStore implements ArchiverDataStore {
   }
 
   getProvenL2BlockNumber(): Promise<number> {
-    return Promise.resolve(this.#provenStore.getProvenL2BlockNumber());
+    return Promise.resolve(this.#blockStore.getProvenL2BlockNumber());
+  }
+
+  getProvenL2EpochNumber(): Promise<number | undefined> {
+    return Promise.resolve(this.#blockStore.getProvenL2EpochNumber());
   }
 
-  async setProvenL2BlockNumber(blockNumber: SingletonDataRetrieval<number>) {
-    await this.#provenStore.setProvenL2BlockNumber(blockNumber);
+  setProvenL2BlockNumber(blockNumber: number) {
+    this.#blockStore.setProvenL2BlockNumber(blockNumber);
+    return Promise.resolve();
+  }
+
+  setProvenL2EpochNumber(epochNumber: number) {
+    this.#blockStore.setProvenL2EpochNumber(epochNumber);
+    return Promise.resolve();
   }
 
   setBlockSynchedL1BlockNumber(l1BlockNumber: bigint) {
@@ -252,7 +295,6 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     return Promise.resolve({
       blocksSynchedTo: this.#blockStore.getSynchedL1BlockNumber(),
       messagesSynchedTo: this.#messageStore.getSynchedL1BlockNumber(),
-      provenLogsSynchedTo: this.#provenStore.getSynchedL1BlockNumber(),
     });
   }
 }

package/src/archiver/kv_archiver_store/log_store.ts

@@ -4,6 +4,7 @@ import {
   ExtendedUnencryptedL2Log,
   type FromLogType,
   type GetUnencryptedLogsResponse,
+  type L2Block,
   type L2BlockL2Logs,
   type LogFilter,
   LogId,
@@ -37,29 +38,28 @@ export class LogStore {
 
   /**
    * Append new logs to the store's list.
-   * @param encryptedLogs - The logs to be added to the store.
-   * @param unencryptedLogs - The type of the logs to be added to the store.
-   * @param blockNumber - The block for which to add the logs.
+   * @param blocks - The blocks for which to add the logs.
    * @returns True if the operation is successful.
    */
-  addLogs(
-    noteEncryptedLogs: EncryptedNoteL2BlockL2Logs | undefined,
-    encryptedLogs: EncryptedL2BlockL2Logs | undefined,
-    unencryptedLogs: UnencryptedL2BlockL2Logs | undefined,
-    blockNumber: number,
-  ): Promise<boolean> {
+  addLogs(blocks: L2Block[]): Promise<boolean> {
     return this.db.transaction(() => {
-      if (noteEncryptedLogs) {
-        void this.#noteEncryptedLogs.set(blockNumber, noteEncryptedLogs.toBuffer());
-      }
+      blocks.forEach(block => {
+        void this.#noteEncryptedLogs.set(block.number, block.body.noteEncryptedLogs.toBuffer());
+        void this.#encryptedLogs.set(block.number, block.body.encryptedLogs.toBuffer());
+        void this.#unencryptedLogs.set(block.number, block.body.unencryptedLogs.toBuffer());
+      });
 
-      if (encryptedLogs) {
-        void this.#encryptedLogs.set(blockNumber, encryptedLogs.toBuffer());
-      }
+      return true;
+    });
+  }
 
-      if (unencryptedLogs) {
-        void this.#unencryptedLogs.set(blockNumber, unencryptedLogs.toBuffer());
-      }
+  deleteLogs(blocks: L2Block[]): Promise<boolean> {
+    return this.db.transaction(() => {
+      blocks.forEach(block => {
+        void this.#noteEncryptedLogs.delete(block.number);
+        void this.#encryptedLogs.delete(block.number);
+        void this.#unencryptedLogs.delete(block.number);
+      });
 
       return true;
     });

package/src/archiver/kv_archiver_store/message_store.ts

@@ -1,10 +1,5 @@
-import { type InboxLeaf } from '@aztec/circuit-types';
-import {
-  Fr,
-  INITIAL_L2_BLOCK_NUM,
-  L1_TO_L2_MSG_SUBTREE_HEIGHT,
-  NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
-} from '@aztec/circuits.js';
+import { InboxLeaf } from '@aztec/circuit-types';
+import { Fr, L1_TO_L2_MSG_SUBTREE_HEIGHT } from '@aztec/circuits.js';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { type AztecKVStore, type AztecMap, type AztecSingleton } from '@aztec/kv-store';
 
@@ -17,6 +12,7 @@ export class MessageStore {
   #l1ToL2Messages: AztecMap<string, Buffer>;
   #l1ToL2MessageIndices: AztecMap<string, bigint[]>; // We store array of bigints here because there can be duplicate messages
   #lastSynchedL1Block: AztecSingleton<bigint>;
+  #totalMessageCount: AztecSingleton<bigint>;
 
   #log = createDebugLogger('aztec:archiver:message_store');
 
@@ -26,6 +22,11 @@ export class MessageStore {
     this.#l1ToL2Messages = db.openMap('archiver_l1_to_l2_messages');
     this.#l1ToL2MessageIndices = db.openMap('archiver_l1_to_l2_message_indices');
     this.#lastSynchedL1Block = db.openSingleton('archiver_last_l1_block_new_messages');
+    this.#totalMessageCount = db.openSingleton('archiver_l1_to_l2_message_count');
+  }
+
+  getTotalL1ToL2MessageCount(): bigint {
+    return this.#totalMessageCount.get() ?? 0n;
   }
 
   /**
@@ -55,21 +56,17 @@ export class MessageStore {
       void this.#lastSynchedL1Block.set(messages.lastProcessedL1BlockNumber);
 
       for (const message of messages.retrievedData) {
-        if (message.index >= this.#l1ToL2MessagesSubtreeSize) {
-          throw new Error(`Message index ${message.index} out of subtree range`);
-        }
-        const key = `${message.blockNumber}-${message.index}`;
+        const key = `${message.index}`;
         void this.#l1ToL2Messages.setIfNotExists(key, message.leaf.toBuffer());
 
-        const indexInTheWholeTree =
-          (message.blockNumber - BigInt(INITIAL_L2_BLOCK_NUM)) * BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP) +
-          message.index;
-
         const indices = this.#l1ToL2MessageIndices.get(message.leaf.toString()) ?? [];
-        indices.push(indexInTheWholeTree);
+        indices.push(message.index);
         void this.#l1ToL2MessageIndices.set(message.leaf.toString(), indices);
       }
 
+      const lastTotalMessageCount = this.getTotalL1ToL2MessageCount();
+      void this.#totalMessageCount.set(lastTotalMessageCount + BigInt(messages.retrievedData.length));
+
       return true;
     });
   }
@@ -89,9 +86,10 @@ export class MessageStore {
   getL1ToL2Messages(blockNumber: bigint): Fr[] {
     const messages: Fr[] = [];
     let undefinedMessageFound = false;
-    for (let messageIndex = 0; messageIndex < this.#l1ToL2MessagesSubtreeSize; messageIndex++) {
+    const startIndex = Number(InboxLeaf.smallestIndexFromL2Block(blockNumber));
+    for (let i = startIndex; i < startIndex + this.#l1ToL2MessagesSubtreeSize; i++) {
       // This is inefficient but probably fine for now.
-      const key = `${blockNumber}-${messageIndex}`;
+      const key = `${i}`;
       const message = this.#l1ToL2Messages.get(key);
       if (message) {
         if (undefinedMessageFound) {

package/src/archiver/memory_archiver_store/l1_to_l2_message_store.ts

@@ -1,9 +1,5 @@
-import { type InboxLeaf } from '@aztec/circuit-types';
-import {
-  INITIAL_L2_BLOCK_NUM,
-  L1_TO_L2_MSG_SUBTREE_HEIGHT,
-  NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
-} from '@aztec/circuits.js/constants';
+import { InboxLeaf } from '@aztec/circuit-types';
+import { L1_TO_L2_MSG_SUBTREE_HEIGHT } from '@aztec/circuits.js/constants';
 import { type Fr } from '@aztec/foundation/fields';
 
 /**
@@ -11,7 +7,7 @@ import { type Fr } from '@aztec/foundation/fields';
  */
 export class L1ToL2MessageStore {
   /**
-   * A map pointing from a key in a "blockNum-messageIndex" format to the corresponding L1 to L2 message hash.
+   * A map pointing from a key in a "messageIndex" format to the corresponding L1 to L2 message hash.
    */
   protected store: Map<string, Fr> = new Map();
 
@@ -19,21 +15,22 @@ export class L1ToL2MessageStore {
 
   constructor() {}
 
+  getTotalL1ToL2MessageCount(): bigint {
+    return BigInt(this.store.size);
+  }
+
   addMessage(message: InboxLeaf) {
-    if (message.index >= this.#l1ToL2MessagesSubtreeSize) {
-      throw new Error(`Message index ${message.index} out of subtree range`);
-    }
-    const key = `${message.blockNumber}-${message.index}`;
-    this.store.set(key, message.leaf);
+    this.store.set(`${message.index}`, message.leaf);
   }
 
   getMessages(blockNumber: bigint): Fr[] {
     const messages: Fr[] = [];
     let undefinedMessageFound = false;
-    for (let messageIndex = 0; messageIndex < this.#l1ToL2MessagesSubtreeSize; messageIndex++) {
+    const startIndex = Number(InboxLeaf.smallestIndexFromL2Block(blockNumber));
+
+    for (let i = startIndex; i < startIndex + this.#l1ToL2MessagesSubtreeSize; i++) {
       // This is inefficient but probably fine for now.
-      const key = `${blockNumber}-${messageIndex}`;
-      const message = this.store.get(key);
+      const message = this.store.get(`${i}`);
       if (message) {
         if (undefinedMessageFound) {
           throw new Error(`L1 to L2 message gap found in block ${blockNumber}`);
@@ -57,10 +54,7 @@ export class L1ToL2MessageStore {
   getMessageIndex(l1ToL2Message: Fr, startIndex: bigint): bigint | undefined {
     for (const [key, message] of this.store.entries()) {
       if (message.equals(l1ToL2Message)) {
-        const keyParts = key.split('-');
-        const [blockNumber, messageIndex] = [BigInt(keyParts[0]), BigInt(keyParts[1])];
-        const indexInTheWholeTree =
-          (blockNumber - BigInt(INITIAL_L2_BLOCK_NUM)) * BigInt(NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP) + messageIndex;
+        const indexInTheWholeTree = BigInt(key);
         if (indexInTheWholeTree < startIndex) {
           continue;
         }