fix: return partial witnesses based on the content of read requests.

yarn-project/acir-simulator/src/client/client_execution_context.ts

@@ -1,5 +1,5 @@
 import { CircuitsWasm, HistoricBlockData, ReadRequestMembershipWitness, TxContext } from '@aztec/circuits.js';
-import { siloNullifier } from '@aztec/circuits.js/abis';
+import { computeUniqueCommitment, siloCommitment } from '@aztec/circuits.js/abis';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr, Point } from '@aztec/foundation/fields';
 import { createDebugLogger } from '@aztec/foundation/log';
@@ -12,40 +12,47 @@ import {
   toAcvmL1ToL2MessageLoadOracleInputs,
 } from '../acvm/index.js';
 import { PackedArgsCache } from '../common/packed_args_cache.js';
-import { DBOracle, PendingNoteData } from './db_oracle.js';
+import { DBOracle } from './db_oracle.js';
+import { ExecutionNoteCache } from './execution_note_cache.js';
+import { NewNoteData } from './execution_result.js';
 import { pickNotes } from './pick_notes.js';
 
 /**
  * The execution context for a client tx simulation.
  */
 export class ClientTxExecutionContext {
-  // Note: not forwarded to nested contexts via `extend()` because these witnesses
-  // are meant to be maintained on a per-call basis as they should mirror read requests
-  // output by an app circuit via public inputs.
-  private readRequestPartialWitnesses: ReadRequestMembershipWitness[] = [];
+  /**
+   * New notes created during this execution.
+   * It's possible that a note in this list has been nullified (in the same or other executions) and doen't exist in the ExecutionNoteCache and the final proof data.
+   * But we still include those notes in the execution result because their commitments are still in the public inputs of this execution.
+   * This information is only for references (currently used for tests), and is not used for any sort of constrains.
+   * Users can also use this to get a clearer idea of what's happened during a simulation.
+   */
+  private newNotes: NewNoteData[] = [];
+  /**
+   * Notes from previous transactions that are returned to the oracle call `getNotes` during this execution.
+   * The mapping maps from the unique siloed note hash to the index for notes created in private executions.
+   * It maps from siloed note hash to the index for notes created by public functions.
+   *
+   * They are not part of the ExecutionNoteCache and being forwarded to nested contexts via `extend()`
+   * because these notes are meant to be maintained on a per-call basis
+   * They should act as references for the read requests output by an app circuit via public inputs.
+   */
+  private gotNotes: Map<bigint, bigint> = new Map();
 
   /** Logger instance */
   private logger = createDebugLogger('aztec:simulator:execution_context');
 
   constructor(
     /**  The database oracle. */
     public db: DBOracle,
-    /** The tx nullifier, which is also the first nullifier. This will be used to compute the nonces for pending notes. */
-    private txNullifier: Fr,
     /** The tx context. */
     public txContext: TxContext,
     /** Data required to reconstruct the block hash, it contains historic roots. */
     public historicBlockData: HistoricBlockData,
     /** The cache of packed arguments */
     public packedArgsCache: PackedArgsCache,
-    /** Pending notes created (and not nullified) up to current point in execution.
-     *  If a nullifier for a note in this list is emitted, the note will be REMOVED. */
-    private pendingNotes: PendingNoteData[] = [],
-    /** The list of nullifiers created in this transaction. The commitment/note which is nullified
-     *  might be pending or not (i.e., was generated in a previous transaction)
-     *  Note that their value (bigint representation) is used because Frs cannot be looked up in Sets. */
-    private pendingNullifiers: Set<bigint> = new Set<bigint>(),
-
+    private noteCache: ExecutionNoteCache,
     private log = createDebugLogger('aztec:simulator:client_execution_context'),
   ) {}
 
@@ -56,21 +63,39 @@ export class ClientTxExecutionContext {
   public extend() {
     return new ClientTxExecutionContext(
       this.db,
-      this.txNullifier,
       this.txContext,
       this.historicBlockData,
       this.packedArgsCache,
-      this.pendingNotes,
-      this.pendingNullifiers,
+      this.noteCache,
     );
   }
 
   /**
-   * For getting accumulated data.
+   * This function will populate readRequestPartialWitnesses which
+   * here is just used to flag reads as "transient" for new notes created during this execution
+   * or to flag non-transient reads with their leafIndex.
+   * The KernelProver will use this to fully populate witnesses and provide hints to
+   * the kernel regarding which commitments each transient read request corresponds to.
+   * @param readRequests - Note hashed of the notes being read.
    * @returns An array of partially filled in read request membership witnesses.
    */
-  public getReadRequestPartialWitnesses() {
-    return this.readRequestPartialWitnesses;
+  public getReadRequestPartialWitnesses(readRequests: Fr[]) {
+    return readRequests
+      .filter(r => !r.isZero())
+      .map(r => {
+        const index = this.gotNotes.get(r.value);
+        return index !== undefined
+          ? ReadRequestMembershipWitness.empty(index)
+          : ReadRequestMembershipWitness.emptyTransient();
+      });
+  }
+
+  /**
+   * Get the data for the newly created notes.
+   * @param innerNoteHashes - Inner note hashes for the notes.
+   */
+  public getNewNotes(): NewNoteData[] {
+    return this.newNotes;
   }
 
   /**
@@ -98,12 +123,6 @@ export class ClientTxExecutionContext {
    * Real notes coming from DB will have a leafIndex which
    * represents their index in the private data tree.
    *
-   * This function will populate this.readRequestPartialWitnesses which
-   * here is just used to flag reads as "transient" (one in getPendingNotes)
-   * or to flag non-transient reads with their leafIndex.
-   * The KernelProver will use this to fully populate witnesses and provide hints to
-   * the kernel regarding which commitments each transient read request corresponds to.
-   *
    * @param contractAddress - The contract address.
    * @param storageSlot - The storage slot.
    * @param numSelects - The number of valid selects in selectBy and selectValues.
@@ -134,15 +153,13 @@ export class ClientTxExecutionContext {
   ): Promise<ACVMField[]> {
     const storageSlotField = fromACVMField(storageSlot);
 
-    const pendingNotes = this.pendingNotes.filter(
-      n => n.contractAddress.equals(contractAddress) && n.storageSlot.equals(storageSlotField),
-    );
+    // Nullified pending notes are already removed from the list.
+    const pendingNotes = this.noteCache.getNotes(contractAddress, storageSlotField);
 
+    const pendingNullifiers = this.noteCache.getNullifiers(contractAddress, storageSlotField);
     const dbNotes = await this.db.getNotes(contractAddress, storageSlotField);
+    const dbNotesFiltered = dbNotes.filter(n => !pendingNullifiers.has((n.siloedNullifier as Fr).value));
 
-    const dbNotesFiltered = dbNotes.filter(n => !this.pendingNullifiers.has((n.siloedNullifier as Fr).value));
-
-    // Nullified pending notes are already removed from the list.
     const notes = pickNotes([...dbNotesFiltered, ...pendingNotes], {
       selects: selectBy
         .slice(0, numSelects)
@@ -158,6 +175,15 @@ export class ClientTxExecutionContext {
         .join(', ')}`,
     );
 
+    const wasm = await CircuitsWasm.get();
+    notes.forEach(n => {
+      if (n.index !== undefined) {
+        const siloedNoteHash = siloCommitment(wasm, n.contractAddress, n.innerNoteHash);
+        const uniqueSiloedNoteHash = computeUniqueCommitment(wasm, n.nonce, siloedNoteHash);
+        this.gotNotes.set(uniqueSiloedNoteHash.value, n.index);
+      }
+    });
+
     // TODO: notice, that if we don't have a note in our DB, we don't know how big the preimage needs to be, and so we don't actually know how many dummy notes to return, or big to make those dummy notes, or where to position `is_some` booleans to inform the noir program that _all_ the notes should be dummies.
     // By a happy coincidence, a `0` field is interpreted as `is_none`, and since in this case (of an empty db) we'll return all zeros (paddedZeros), the noir program will treat the returned data as all dummies, but this is luck. Perhaps a preimage size should be conveyed by the get_notes noir oracle?
     const preimageLength = notes?.[0]?.preimage.length ?? 0;
@@ -173,14 +199,6 @@ export class ClientTxExecutionContext {
 
     const realNotePreimages = notes.flatMap(({ nonce, preimage }) => [nonce, isSome, ...preimage]);
 
-    // Add a partial witness for each note.
-    // It contains the note index for db notes. And flagged as transient for pending notes.
-    notes.forEach(({ index }) => {
-      this.readRequestPartialWitnesses.push(
-        index !== undefined ? ReadRequestMembershipWitness.empty(index) : ReadRequestMembershipWitness.emptyTransient(),
-      );
-    });
-
     const returnHeaderLength = 2; // is for the header values: `notes.length` and `contractAddress`.
     const extraPreimageLength = 2; // is for the nonce and isSome fields.
     const extendedPreimageLength = preimageLength + extraPreimageLength;
@@ -202,6 +220,58 @@ export class ClientTxExecutionContext {
     );
   }
 
+  /**
+   * Keep track of the new note created during execution.
+   * It can be used in subsequent calls (or transactions when chaining txs is possible).
+   * @param contractAddress - The contract address.
+   * @param storageSlot - The storage slot.
+   * @param preimage - The preimage of the new note.
+   * @param innerNoteHash - The inner note hash of the new note.
+   * @returns
+   */
+  public handleNewNote(
+    contractAddress: AztecAddress,
+    storageSlot: ACVMField,
+    preimage: ACVMField[],
+    innerNoteHash: ACVMField,
+  ) {
+    this.noteCache.addNewNote({
+      contractAddress,
+      storageSlot: fromACVMField(storageSlot),
+      nonce: Fr.ZERO, // Nonce cannot be known during private execution.
+      preimage: preimage.map(f => fromACVMField(f)),
+      siloedNullifier: undefined, // Siloed nullifier cannot be known for newly created note.
+      innerNoteHash: fromACVMField(innerNoteHash),
+    });
+    this.newNotes.push({
+      storageSlot: fromACVMField(storageSlot),
+      preimage: preimage.map(f => fromACVMField(f)),
+    });
+  }
+
+  /**
+   * Adding a siloed nullifier into the current set of all pending nullifiers created
+   * within the current transaction/execution.
+   * @param contractAddress - The contract address.
+   * @param storageSlot - The storage slot.
+   * @param innerNullifier - The pending nullifier to add in the list (not yet siloed by contract address).
+   * @param innerNoteHash - The inner note hash of the new note.
+   * @param contractAddress - The contract address
+   */
+  public async handleNullifiedNote(
+    contractAddress: AztecAddress,
+    storageSlot: ACVMField,
+    innerNullifier: ACVMField,
+    innerNoteHash: ACVMField,
+  ) {
+    await this.noteCache.nullifyNote(
+      contractAddress,
+      fromACVMField(storageSlot),
+      fromACVMField(innerNullifier),
+      fromACVMField(innerNoteHash),
+    );
+  }
+
   /**
    * Fetches the a message from the db, given its key.
    * @param msgKey - A buffer representing the message key.
@@ -219,71 +289,16 @@ export class ClientTxExecutionContext {
    * @returns The commitment data.
    */
   public async getCommitment(contractAddress: AztecAddress, commitment: ACVMField) {
+    const innerNoteHash = fromACVMField(commitment);
     // TODO(https://github.com/AztecProtocol/aztec-packages/issues/1386): only works
     // for noteHashes/commitments created by public functions! Once public kernel or
     // base rollup circuit injects nonces, this can be used with commitments created by
     // private functions as well.
-    const commitmentInputs = await this.db.getCommitmentOracle(contractAddress, fromACVMField(commitment));
+    const commitmentInputs = await this.db.getCommitmentOracle(contractAddress, innerNoteHash);
     // TODO(https://github.com/AztecProtocol/aztec-packages/issues/1029): support pending commitments here
-    this.readRequestPartialWitnesses.push(ReadRequestMembershipWitness.empty(commitmentInputs.index));
-    return toAcvmCommitmentLoadOracleInputs(commitmentInputs, this.historicBlockData.privateDataTreeRoot);
-  }
-
-  /**
-   * Process new note created during execution.
-   * @param contractAddress - The contract address.
-   * @param storageSlot - The storage slot.
-   * @param preimage - new note preimage.
-   * @param innerNoteHash - inner note hash
-   */
-  public pushNewNote(contractAddress: AztecAddress, storageSlot: Fr, preimage: Fr[], innerNoteHash: Fr) {
-    this.pendingNotes.push({
-      contractAddress,
-      storageSlot: storageSlot,
-      nonce: Fr.ZERO, // nonce is cannot be known during private execution
-      preimage,
-      innerNoteHash,
-    });
-  }
-
-  /**
-   * Adding a siloed nullifier into the current set of all pending nullifiers created
-   * within the current transaction/execution.
-   * @param innerNullifier - The pending nullifier to add in the list (not yet siloed by contract address).
-   * @param contractAddress - The contract address
-   */
-  public async pushNewNullifier(innerNullifier: Fr, contractAddress: AztecAddress) {
     const wasm = await CircuitsWasm.get();
-    const siloedNullifier = siloNullifier(wasm, contractAddress, innerNullifier);
-    this.pendingNullifiers.add(siloedNullifier.value);
-  }
-
-  /**
-   * Update the list of pending notes by chopping a note which is nullified.
-   * The identifier used to determine matching is the inner note hash value.
-   * However, we adopt a defensive approach and ensure that the contract address
-   * and storage slot do match.
-   * Note that this method might be called with an innerNoteHash referring to
-   * a note created in a previous transaction which will result in this array
-   * of pending notes left unchanged.
-   * @param innerNoteHash - The inner note hash value to which note will be chopped.
-   * @param contractAddress - The contract address
-   * @param storageSlot - The storage slot as a Field Fr element
-   */
-  public nullifyPendingNotes(innerNoteHash: Fr, contractAddress: AztecAddress, storageSlot: Fr) {
-    // IMPORTANT: We do need an in-place array mutation of this.pendingNotes as this array is shared
-    // by reference among the nested calls. That is the main reason for 'splice' usage below.
-    this.pendingNotes.splice(
-      0,
-      this.pendingNotes.length,
-      ...this.pendingNotes.filter(
-        n =>
-          !(
-            n.innerNoteHash.equals(innerNoteHash) &&
-            n.contractAddress.equals(contractAddress) &&
-            n.storageSlot.equals(storageSlot)
-          ),
-      ),
-    );
+    const siloedNoteHash = siloCommitment(wasm, contractAddress, innerNoteHash);
+    this.gotNotes.set(siloedNoteHash.value, commitmentInputs.index);
+    return toAcvmCommitmentLoadOracleInputs(commitmentInputs, this.historicBlockData.privateDataTreeRoot);
   }
 }

yarn-project/acir-simulator/src/client/db_oracle.ts

@@ -18,20 +18,14 @@ export interface NoteData {
   nonce: Fr;
   /** The preimage of the note */
   preimage: Fr[];
-  /** The corresponding nullifier of the note */
+  /** The inner note hash of the note. */
+  innerNoteHash: Fr;
+  /** The corresponding nullifier of the note. Undefined for pending notes. */
   siloedNullifier?: Fr;
   /** The note's leaf index in the private data tree. Undefined for pending notes. */
   index?: bigint;
 }
 
-/**
- * Information about a note needed during execution.
- */
-export interface PendingNoteData extends NoteData {
-  /** The inner note hash (used as a nullified commitment). */
-  innerNoteHash: Fr;
-}
-
 /**
  * The format that noir uses to get L1 to L2 Messages.
  */