@aztec/pxe 0.0.1-commit.1142ef1 → 0.0.1-commit.18ccd8f0

package/src/storage/note_store/stored_note.ts

@@ -0,0 +1,48 @@
+import { BlockNumber } from '@aztec/foundation/branded-types';
+import { BufferReader, serializeToBuffer } from '@aztec/foundation/serialize';
+import { NoteDao } from '@aztec/stdlib/note';
+
+export class StoredNote {
+  constructor(
+    readonly noteDao: NoteDao,
+    readonly scopes: Set<string>,
+    private _nullifiedAt: BlockNumber | undefined = undefined,
+  ) {}
+
+  static fromBuffer(buffer: Buffer) {
+    const reader = BufferReader.asReader(buffer);
+
+    const noteDao = NoteDao.fromBuffer(reader);
+    const scopes = reader.readVector({ fromBuffer: (r: BufferReader) => r.readString() });
+
+    const nullifiedAtRaw = reader.readNumber();
+    const nullifiedAt = nullifiedAtRaw === 0 ? undefined : (nullifiedAtRaw as BlockNumber);
+
+    return new StoredNote(noteDao, new Set(scopes), nullifiedAt);
+  }
+
+  toBuffer(): Buffer {
+    const scopesArray = [...this.scopes];
+    return serializeToBuffer(this.noteDao, scopesArray.length, ...scopesArray, this._nullifiedAt ?? 0);
+  }
+
+  addScope(scope: string) {
+    this.scopes.add(scope);
+  }
+
+  markAsNullified(blockNumber: BlockNumber) {
+    this._nullifiedAt = blockNumber;
+  }
+
+  markAsActive() {
+    this._nullifiedAt = undefined;
+  }
+
+  isNullified() {
+    return this._nullifiedAt !== undefined;
+  }
+
+  get nullifiedAt() {
+    return this._nullifiedAt;
+  }
+}

package/src/storage/private_event_store/private_event_store.ts

@@ -1,14 +1,16 @@
 import { BlockNumber } from '@aztec/foundation/branded-types';
 import { Fr } from '@aztec/foundation/curves/bn254';
+import { toArray } from '@aztec/foundation/iterable';
 import { createLogger } from '@aztec/foundation/log';
-import { BufferReader, serializeToBuffer } from '@aztec/foundation/serialize';
-import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
+import { Semaphore } from '@aztec/foundation/queue';
+import type { AztecAsyncKVStore, AztecAsyncMap, AztecAsyncMultiMap } from '@aztec/kv-store';
 import type { EventSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { L2BlockHash } from '@aztec/stdlib/block';
-import { type InTx, TxHash } from '@aztec/stdlib/tx';
+import type { InTx, TxHash } from '@aztec/stdlib/tx';
 
+import type { StagedStore } from '../../job_coordinator/job_coordinator.js';
 import type { PackedPrivateEvent } from '../../pxe.js';
+import { StoredPrivateEvent } from './stored_private_event.js';
 
 export type PrivateEventStoreFilter = {
   contractAddress: AztecAddress;
@@ -18,55 +20,53 @@ export type PrivateEventStoreFilter = {
   txHash?: TxHash;
 };
 
-type PrivateEventEntry = {
-  randomness: Fr; // Note that this value is currently not being returned on queries and is therefore temporarily unused
-  msgContent: Buffer;
-  eventCommitmentIndex: number;
-  l2BlockNumber: number;
-  l2BlockHash: Buffer;
-  txHash: Buffer;
-  /** The lookup key for #eventsByContractScopeSelector, used for cleanup during rollback */
-  lookupKey: string;
-};
-
 type PrivateEventMetadata = InTx & {
   contractAddress: AztecAddress;
   scope: AztecAddress;
+  /** The index of the tx within the block */
+  txIndexInBlock: number;
+  /** The index of the event within the tx (based on nullifier position) */
+  eventIndexInTx: number;
 };
 
 /**
  * Stores decrypted private event logs.
  */
-export class PrivateEventStore {
+export class PrivateEventStore implements StagedStore {
+  readonly storeName: string = 'private_event';
+
   #store: AztecAsyncKVStore;
-  /** Map storing the actual private event log entries, keyed by eventCommitmentIndex */
-  #eventLogs: AztecAsyncMap<number, PrivateEventEntry>;
-  /** Map from contractAddress_scope_eventSelector to eventCommitmentIndex[] for efficient lookup */
-  #eventsByContractScopeSelector: AztecAsyncMap<string, number[]>;
-  /** Map from block number to eventCommitmentIndex[] for rollback support */
-  #eventsByBlockNumber: AztecAsyncMap<number, number[]>;
-  /** Map from eventCommitmentIndex to boolean indicating if log has been seen. */
-  #seenLogs: AztecAsyncMap<number, boolean>;
+  /** Actual private event log entries, keyed by siloedEventCommitment */
+  #events: AztecAsyncMap<string, Buffer>;
+  /** Multi-map from contractAddress_eventSelector to siloedEventCommitment for efficient lookup */
+  #eventsByContractAndEventSelector: AztecAsyncMultiMap<string, string>;
+  /** Multi-map from block number to siloedEventCommitment for rollback support */
+  #eventsByBlockNumber: AztecAsyncMultiMap<number, string>;
+
+  /** jobId => eventId (event siloed nullifier) => StoredPrivateEvent */
+  #eventsForJob: Map<string, Map<string, StoredPrivateEvent>>;
+
+  /** Per-job locks to prevent concurrent writes from affecting each other. */
+  #jobLocks: Map<string, Semaphore>;
 
   logger = createLogger('private_event_store');
 
   constructor(store: AztecAsyncKVStore) {
     this.#store = store;
-    this.#eventLogs = this.#store.openMap('private_event_logs');
-    this.#eventsByContractScopeSelector = this.#store.openMap('events_by_contract_scope_selector');
-    this.#seenLogs = this.#store.openMap('seen_logs');
-    this.#eventsByBlockNumber = this.#store.openMap('events_by_block_number');
-  }
+    this.#events = this.#store.openMap('private_event_logs');
+    this.#eventsByContractAndEventSelector = this.#store.openMultiMap('events_by_contract_selector');
+    this.#eventsByBlockNumber = this.#store.openMultiMap('events_by_block_number');
 
-  #keyFor(contractAddress: AztecAddress, scope: AztecAddress, eventSelector: EventSelector): string {
-    return `${contractAddress.toString()}_${scope.toString()}_${eventSelector.toString()}`;
+    this.#eventsForJob = new Map();
+    this.#jobLocks = new Map();
   }
 
   /**
    * Store a private event log.
    * @param eventSelector - The event selector of the event.
+   * @param randomness - The randomness used for the event commitment.
    * @param msgContent - The content of the event.
-   * @param eventCommitmentIndex - The index of the event commitment in the nullifier tree.
+   * @param siloedEventCommitment - The siloed event commitment (used as unique identifier).
    * @param metadata
    *  contractAddress - The address of the contract that emitted the event.
    *  scope - The address to which the event is scoped.
@@ -77,41 +77,46 @@ export class PrivateEventStore {
     eventSelector: EventSelector,
     randomness: Fr,
     msgContent: Fr[],
-    eventCommitmentIndex: number,
+    siloedEventCommitment: Fr,
     metadata: PrivateEventMetadata,
-  ): Promise<void> {
-    const { contractAddress, scope, txHash, l2BlockNumber, l2BlockHash } = metadata;
-
-    return this.#store.transactionAsync(async () => {
-      const key = this.#keyFor(contractAddress, scope, eventSelector);
-
-      // Check if this exact log has already been stored using eventCommitmentIndex as unique identifier
-      const hasBeenSeen = await this.#seenLogs.getAsync(eventCommitmentIndex);
-      if (hasBeenSeen) {
-        this.logger.verbose('Ignoring duplicate event log', { txHash: txHash.toString(), eventCommitmentIndex });
-        return;
-      }
-
-      this.logger.verbose('storing private event log', { contractAddress, scope, msgContent, l2BlockNumber });
+    jobId: string,
+  ) {
+    return this.#withJobLock(jobId, async () => {
+      const { contractAddress, scope, txHash, l2BlockNumber, l2BlockHash, txIndexInBlock, eventIndexInTx } = metadata;
+      const eventId = siloedEventCommitment.toString();
 
-      await this.#eventLogs.set(eventCommitmentIndex, {
-        randomness,
-        msgContent: serializeToBuffer(msgContent),
+      this.logger.verbose('storing private event log (job stage)', {
+        eventId,
+        contractAddress,
+        scope,
+        msgContent,
         l2BlockNumber,
-        l2BlockHash: l2BlockHash.toBuffer(),
-        eventCommitmentIndex,
-        txHash: txHash.toBuffer(),
-        lookupKey: key,
       });
 
-      const existingIndices = (await this.#eventsByContractScopeSelector.getAsync(key)) || [];
-      await this.#eventsByContractScopeSelector.set(key, [...existingIndices, eventCommitmentIndex]);
+      const existing = await this.#readEvent(eventId, jobId);
 
-      const existingBlockIndices = (await this.#eventsByBlockNumber.getAsync(l2BlockNumber)) || [];
-      await this.#eventsByBlockNumber.set(l2BlockNumber, [...existingBlockIndices, eventCommitmentIndex]);
-
-      // Mark this log as seen using eventCommitmentIndex
-      await this.#seenLogs.set(eventCommitmentIndex, true);
+      if (existing) {
+        // If we already stored this event, we still want to make sure to track it for the given scope
+        existing.addScope(scope.toString());
+        this.#writeEvent(eventId, existing, jobId);
+      } else {
+        this.#writeEvent(
+          eventId,
+          new StoredPrivateEvent(
+            randomness,
+            msgContent,
+            l2BlockNumber,
+            l2BlockHash,
+            txHash,
+            txIndexInBlock,
+            eventIndexInTx,
+            contractAddress,
+            eventSelector,
+            new Set([scope.toString()]),
+          ),
+          jobId,
+        );
+      }
     });
   }
 
@@ -123,94 +128,234 @@ export class PrivateEventStore {
    *  fromBlock: The block number to search from (inclusive).
    *  toBlock: The block number to search upto (exclusive).
    *  scope: - The addresses that decrypted the logs.
-   * @returns - The event log contents, augmented with metadata about
-   *  the transaction and block it the event was included in .
+   * @returns - The event log contents, augmented with metadata about the transaction and block in which the event was
+   * included.
    */
   public async getPrivateEvents(
     eventSelector: EventSelector,
     filter: PrivateEventStoreFilter,
   ): Promise<PackedPrivateEvent[]> {
-    const events: Array<{ eventCommitmentIndex: number; event: PackedPrivateEvent }> = [];
+    const events: Array<{
+      l2BlockNumber: number;
+      txIndexInBlock: number;
+      eventIndexInTx: number;
+      event: PackedPrivateEvent;
+    }> = [];
 
-    for (const scope of filter.scopes) {
-      const key = this.#keyFor(filter.contractAddress, scope, eventSelector);
-      const eventCommitmentIndices = (await this.#eventsByContractScopeSelector.getAsync(key)) || [];
+    const key = this.#keyFor(filter.contractAddress, eventSelector);
+    const targetScopes = new Set(filter.scopes.map(s => s.toString()));
 
-      for (const eventCommitmentIndex of eventCommitmentIndices) {
-        const entry = await this.#eventLogs.getAsync(eventCommitmentIndex);
-        if (!entry || entry.l2BlockNumber < filter.fromBlock || entry.l2BlockNumber >= filter.toBlock) {
-          continue;
-        }
+    const eventIds: string[] = await toArray(this.#eventsByContractAndEventSelector.getValuesAsync(key));
 
-        // Convert buffer back to Fr array
-        const reader = BufferReader.asReader(entry.msgContent);
-        const numFields = entry.msgContent.length / Fr.SIZE_IN_BYTES;
-        const msgContent = reader.readArray(numFields, Fr);
-        const txHash = TxHash.fromBuffer(entry.txHash);
-        const l2BlockHash = L2BlockHash.fromBuffer(entry.l2BlockHash);
+    for (const eventId of eventIds) {
+      const eventBuffer = await this.#events.getAsync(eventId);
 
-        if (filter.txHash && !txHash.equals(filter.txHash)) {
-          continue;
-        }
+      // Defensive, if it happens, there's a problem with how we're handling #eventsByContractAndEventSelector
+      if (!eventBuffer) {
+        this.logger.verbose(
+          `EventId ${eventId} does not exist in main index but it is referenced from contract event selector index`,
+        );
+        continue;
+      }
 
-        events.push({
-          eventCommitmentIndex: entry.eventCommitmentIndex,
-          event: {
-            packedEvent: msgContent,
-            l2BlockNumber: BlockNumber(entry.l2BlockNumber),
-            txHash,
-            l2BlockHash,
-            eventSelector,
-          },
-        });
+      const storedPrivateEvent = StoredPrivateEvent.fromBuffer(eventBuffer);
+
+      // Filter by block range
+      if (storedPrivateEvent.l2BlockNumber < filter.fromBlock || storedPrivateEvent.l2BlockNumber >= filter.toBlock) {
+        continue;
+      }
+
+      // Filter by scopes
+      if (storedPrivateEvent.scopes.intersection(targetScopes).size === 0) {
+        continue;
+      }
+
+      // Filter by txHash
+      if (filter.txHash && !storedPrivateEvent.txHash.equals(filter.txHash)) {
+        continue;
       }
+
+      events.push({
+        l2BlockNumber: storedPrivateEvent.l2BlockNumber,
+        txIndexInBlock: storedPrivateEvent.txIndexInBlock,
+        eventIndexInTx: storedPrivateEvent.eventIndexInTx,
+        event: {
+          packedEvent: storedPrivateEvent.msgContent,
+          l2BlockNumber: BlockNumber(storedPrivateEvent.l2BlockNumber),
+          txHash: storedPrivateEvent.txHash,
+          l2BlockHash: storedPrivateEvent.l2BlockHash,
+          eventSelector,
+        },
+      });
     }
 
-    // Sort by eventCommitmentIndex only
-    events.sort((a, b) => a.eventCommitmentIndex - b.eventCommitmentIndex);
+    // Sort by block number, then by tx index within block, then by event index within tx
+    events.sort((a, b) => {
+      if (a.l2BlockNumber !== b.l2BlockNumber) {
+        return a.l2BlockNumber - b.l2BlockNumber;
+      }
+      if (a.txIndexInBlock !== b.txIndexInBlock) {
+        return a.txIndexInBlock - b.txIndexInBlock;
+      }
+      return a.eventIndexInTx - b.eventIndexInTx;
+    });
+
     return events.map(ev => ev.event);
   }
 
   /**
    * Rolls back private events that were stored after a given `blockNumber` and up to `synchedBlockNumber` (the block
    * number up to which PXE managed to sync before the reorg happened).
+   *
+   * We don't need staged writes for a rollback since it's handled in the context of a blockchain rewind.
+   *
+   * Rollbacks are handled by the BlockSynchronizer, which runs a DB transaction across stores when it detects a
+   * re-org, including setting the new anchor block after rolling back.
+   *
+   * So if anything fails in the process of rolling back any store, all DB changes occurring during rollbacks will be
+   * lost and the anchor block will not be updated; which means this code will eventually need to run again
+   * (i.e.: PXE will detect it's basing it work on an invalid block hash, then which re-triggers rewind).
+   *
+   * For further details, refer to `BlockSynchronizer#handleBlockStreamEvent`.
+   *
+   * IMPORTANT: This method must be called within a transaction to ensure atomicity.
    */
-  public async rollbackEventsAfterBlock(blockNumber: number, synchedBlockNumber: number): Promise<void> {
-    await this.#store.transactionAsync(async () => {
-      let removedCount = 0;
-
-      for (let block = blockNumber + 1; block <= synchedBlockNumber; block++) {
-        const indices = await this.#eventsByBlockNumber.getAsync(block);
-        if (indices) {
-          await this.#eventsByBlockNumber.delete(block);
-
-          for (const eventCommitmentIndex of indices) {
-            const entry = await this.#eventLogs.getAsync(eventCommitmentIndex);
-            if (!entry) {
-              throw new Error(`Event log not found for eventCommitmentIndex ${eventCommitmentIndex}`);
-            }
-
-            await this.#eventLogs.delete(eventCommitmentIndex);
-            await this.#seenLogs.delete(eventCommitmentIndex);
-
-            // Update #eventsByContractScopeSelector using the stored lookupKey
-            const existingIndices = await this.#eventsByContractScopeSelector.getAsync(entry.lookupKey);
-            if (!existingIndices || existingIndices.length === 0) {
-              throw new Error(`No indices found in #eventsByContractScopeSelector for key ${entry.lookupKey}`);
-            }
-            const filteredIndices = existingIndices.filter(idx => idx !== eventCommitmentIndex);
-            if (filteredIndices.length === 0) {
-              await this.#eventsByContractScopeSelector.delete(entry.lookupKey);
-            } else {
-              await this.#eventsByContractScopeSelector.set(entry.lookupKey, filteredIndices);
-            }
-
-            removedCount++;
+  public async rollback(blockNumber: number, synchedBlockNumber: number): Promise<void> {
+    let removedCount = 0;
+
+    for (let block = blockNumber + 1; block <= synchedBlockNumber; block++) {
+      const eventIds: string[] = await toArray(this.#eventsByBlockNumber.getValuesAsync(block));
+
+      if (eventIds.length > 0) {
+        await this.#eventsByBlockNumber.delete(block);
+
+        for (const eventId of eventIds) {
+          const buffer = await this.#events.getAsync(eventId);
+          if (!buffer) {
+            throw new Error(`Event not found for eventId ${eventId}`);
           }
+
+          const entry = StoredPrivateEvent.fromBuffer(buffer);
+          await this.#events.delete(eventId);
+          await this.#eventsByContractAndEventSelector.deleteValue(
+            this.#keyFor(entry.contractAddress, entry.eventSelector),
+            eventId,
+          );
+
+          removedCount++;
         }
       }
+    }
+
+    this.logger.verbose(`Rolled back ${removedCount} private events after block ${blockNumber}`);
+  }
+
+  /**
+   * Commits in memory job data to persistent storage.
+   *
+   * Called by JobCoordinator when a job completes successfully.
+   *
+   * Note: JobCoordinator wraps all commits in a single transaction, so we don't need our own transactionAsync here
+   * (and using one would throw on IndexedDB as it does not support nested txs).
+   *
+   * @param jobId - The jobId identifying which staged data to commit
+   */
+  commit(jobId: string): Promise<void> {
+    return this.#withJobLock(jobId, async () => {
+      for (const [eventId, entry] of this.#getEventsForJob(jobId).entries()) {
+        const lookupKey = this.#keyFor(entry.contractAddress, entry.eventSelector);
+        this.logger.verbose('storing private event log', { eventId, lookupKey });
+
+        await Promise.all([
+          this.#events.set(eventId, entry.toBuffer()),
+          this.#eventsByContractAndEventSelector.set(lookupKey, eventId),
+          this.#eventsByBlockNumber.set(entry.l2BlockNumber, eventId),
+        ]);
+      }
 
-      this.logger.verbose(`Rolled back ${removedCount} private events after block ${blockNumber}`);
+      this.#clearJobData(jobId);
     });
   }
+
+  /**
+   * Discards in memory job data without persisting it.
+   */
+  discardStaged(jobId: string): Promise<void> {
+    return this.#withJobLock(jobId, () => Promise.resolve(this.#clearJobData(jobId)));
+  }
+
+  /**
+   * Reads an event from in-memory job data first, falling back to persistent storage if not found.
+   *
+   * Returns undefined if the event does not exist in the store overall.
+   */
+  async #readEvent(eventId: string, jobId: string): Promise<StoredPrivateEvent | undefined> {
+    const eventForJob = this.#getEventsForJob(jobId).get(eventId);
+    if (eventForJob) {
+      return eventForJob;
+    }
+
+    const buffer = await this.#events.getAsync(eventId);
+    return buffer ? StoredPrivateEvent.fromBuffer(buffer) : undefined;
+  }
+
+  /**
+   * Writes an event to in-memory job data.
+   *
+   * Writes are only allowed in a job context. Events modified during a job will only be persisted when `commit` is
+   * called.
+   */
+  #writeEvent(eventId: string, entry: StoredPrivateEvent, jobId: string) {
+    this.#getEventsForJob(jobId).set(eventId, entry);
+  }
+
+  /**
+   * Get in-memory data only visible to @param jobId
+   */
+  #getEventsForJob(jobId: string): Map<string, StoredPrivateEvent> {
+    let eventsForJob = this.#eventsForJob.get(jobId);
+    if (eventsForJob === undefined) {
+      eventsForJob = new Map();
+      this.#eventsForJob.set(jobId, eventsForJob);
+    }
+    return eventsForJob;
+  }
+
+  /**
+   * Clear data structures supporting a specific job.
+   */
+  #clearJobData(jobId: string) {
+    this.#eventsForJob.delete(jobId);
+    this.#jobLocks.delete(jobId);
+  }
+
+  /**
+   * Ensures a function can only run once it acquires a unique per-job lock, and handles proper lock release after it
+   * runs.
+   *
+   * This primitive allows concurrent writes on this store without risking data corruption due to unsound write
+   * interleaving.
+   */
+  async #withJobLock<T>(jobId: string, fn: () => Promise<T>): Promise<T> {
+    let lock = this.#jobLocks.get(jobId);
+    if (!lock) {
+      lock = new Semaphore(1);
+      this.#jobLocks.set(jobId, lock);
+    }
+    await lock.acquire();
+    try {
+      return await fn();
+    } finally {
+      lock.release();
+    }
+  }
+
+  /**
+   * Returns a string key based on @param contractAddress and @param eventSelector.
+   *
+   * The returned key is meant to be used when interacting with index #eventsByContractAndEventSelector.
+   */
+  #keyFor(contractAddress: AztecAddress, eventSelector: EventSelector): string {
+    return `${contractAddress.toString()}_${eventSelector.toString()}`;
+  }
 }

package/src/storage/private_event_store/stored_private_event.ts

@@ -0,0 +1,73 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { BufferReader, serializeToBuffer } from '@aztec/foundation/serialize';
+import { EventSelector } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { BlockHash } from '@aztec/stdlib/block';
+import { TxHash } from '@aztec/stdlib/tx';
+
+/** Serializable private event entry with scope tracking. */
+export class StoredPrivateEvent {
+  constructor(
+    readonly randomness: Fr,
+    readonly msgContent: Fr[],
+    readonly l2BlockNumber: number,
+    readonly l2BlockHash: BlockHash,
+    readonly txHash: TxHash,
+    readonly txIndexInBlock: number,
+    readonly eventIndexInTx: number,
+    readonly contractAddress: AztecAddress,
+    readonly eventSelector: EventSelector,
+    readonly scopes: Set<string>,
+  ) {}
+
+  addScope(scope: string) {
+    this.scopes.add(scope);
+  }
+
+  toBuffer(): Buffer {
+    const scopesArray = [...this.scopes];
+    return serializeToBuffer(
+      this.randomness,
+      this.msgContent.length,
+      ...this.msgContent,
+      this.l2BlockNumber,
+      this.l2BlockHash,
+      this.txHash,
+      this.txIndexInBlock,
+      this.eventIndexInTx,
+      this.contractAddress,
+      this.eventSelector.toBuffer(),
+      scopesArray.length,
+      ...scopesArray,
+    );
+  }
+
+  static fromBuffer(buffer: Buffer): StoredPrivateEvent {
+    const reader = BufferReader.asReader(buffer);
+
+    const randomness = Fr.fromBuffer(reader);
+    const msgContentLength = reader.readNumber();
+    const msgContent = reader.readArray(msgContentLength, Fr);
+    const l2BlockNumber = reader.readNumber();
+    const l2BlockHash = BlockHash.fromBuffer(reader);
+    const txHash = TxHash.fromBuffer(reader);
+    const txIndexInBlock = reader.readNumber();
+    const eventIndexInTx = reader.readNumber();
+    const contractAddress = AztecAddress.fromBuffer(reader);
+    const eventSelector = EventSelector.fromBuffer(reader);
+    const scopes = reader.readVector({ fromBuffer: (r: BufferReader) => r.readString() });
+
+    return new StoredPrivateEvent(
+      randomness,
+      msgContent,
+      l2BlockNumber,
+      l2BlockHash,
+      txHash,
+      txIndexInBlock,
+      eventIndexInTx,
+      contractAddress,
+      eventSelector,
+      new Set(scopes),
+    );
+  }
+}