feat(avm-simulator): consider previous pending nullifiers across enqueued calls

yarn-project/end-to-end/src/e2e_avm_simulator.test.ts

@@ -1,4 +1,4 @@
-import { type AccountWallet, AztecAddress, Fr, FunctionSelector, TxStatus } from '@aztec/aztec.js';
+import { type AccountWallet, AztecAddress, BatchCall, Fr, FunctionSelector, TxStatus } from '@aztec/aztec.js';
 import { GasSettings } from '@aztec/circuits.js';
 import {
   AvmAcvmInteropTestContract,
@@ -60,6 +60,19 @@ describe('e2e_avm_simulator', () => {
         await avmContract.methods.add_storage_map(address, 100).send().wait();
         expect(await avmContract.methods.view_storage_map(address).simulate()).toEqual(200n);
       });
+
+      it('Preserves storage across enqueued public calls', async () => {
+        const address = AztecAddress.fromBigInt(9090n);
+        // This will create 1 tx with 2 public calls in it.
+        await new BatchCall(wallet, [
+          avmContract.methods.set_storage_map(address, 100).request(),
+          avmContract.methods.add_storage_map(address, 100).request(),
+        ])
+          .send()
+          .wait();
+        // On a separate tx, we check the result.
+        expect(await avmContract.methods.view_storage_map(address).simulate()).toEqual(200n);
+      });
     });
 
     describe('Contract instance', () => {
@@ -85,6 +98,18 @@ describe('e2e_avm_simulator', () => {
         tx = await avmContract.methods.assert_nullifier_exists(nullifier).send().wait();
         expect(tx.status).toEqual(TxStatus.MINED);
       });
+
+      it('Emit and check in separate enqueued calls but same tx', async () => {
+        const nullifier = new Fr(123456);
+
+        // This will create 1 tx with 2 public calls in it.
+        await new BatchCall(wallet, [
+          avmContract.methods.new_nullifier(nullifier).request(),
+          avmContract.methods.assert_nullifier_exists(nullifier).request(),
+        ])
+          .send()
+          .wait();
+      });
     });
   });
 

yarn-project/prover-client/src/mocks/test_context.ts

@@ -1,5 +1,5 @@
 import { type BlockProver, type ProcessedTx, type Tx, type TxValidator } from '@aztec/circuit-types';
-import { type Gas, GlobalVariables, Header, type TxContext } from '@aztec/circuits.js';
+import { type Gas, GlobalVariables, Header, type Nullifier, type TxContext } from '@aztec/circuits.js';
 import { type Fr } from '@aztec/foundation/fields';
 import { type DebugLogger } from '@aztec/foundation/log';
 import { openTmpStore } from '@aztec/kv-store/utils';
@@ -129,6 +129,7 @@ export class TestContext {
       _globalVariables: GlobalVariables,
       availableGas: Gas,
       _txContext: TxContext,
+      _pendingNullifiers: Nullifier[],
       transactionFee?: Fr,
       _sideEffectCounter?: number,
     ) => {
@@ -166,6 +167,7 @@ export class TestContext {
       globalVariables: GlobalVariables,
       availableGas: Gas,
       txContext: TxContext,
+      pendingNullifiers: Nullifier[],
       transactionFee?: Fr,
       sideEffectCounter?: number,
     ) => Promise<PublicExecutionResult>,

yarn-project/simulator/src/avm/journal/journal.ts

@@ -50,6 +50,8 @@ export type JournalData = {
   newLogsHashes: TracedUnencryptedL2Log[];
   /** contract address -\> key -\> value */
   currentStorageValue: Map<bigint, Map<bigint, Fr>>;
+
+  sideEffectCounter: number;
 };
 
 // TRANSITIONAL: This should be removed once the kernel handles and entire enqueued call per circuit
@@ -143,6 +145,15 @@ export class AvmPersistableStateManager {
     this.publicStorage.write(storageAddress, slot, value);
 
     // TRANSITIONAL: This should be removed once the kernel handles and entire enqueued call per circuit
+    // The current info to the kernel clears any previous read or write request.
+    this.transitionalExecutionResult.contractStorageReads =
+      this.transitionalExecutionResult.contractStorageReads.filter(
+        read => !read.storageSlot.equals(slot) || !read.contractAddress!.equals(storageAddress),
+      );
+    this.transitionalExecutionResult.contractStorageUpdateRequests =
+      this.transitionalExecutionResult.contractStorageUpdateRequests.filter(
+        update => !update.storageSlot.equals(slot) || !update.contractAddress!.equals(storageAddress),
+      );
     this.transitionalExecutionResult.contractStorageUpdateRequests.push(
       new ContractStorageUpdateRequest(slot, value, this.trace.accessCounter, storageAddress),
     );
@@ -159,16 +170,24 @@ export class AvmPersistableStateManager {
    * @returns the latest value written to slot, or 0 if never written to before
    */
   public async readStorage(storageAddress: Fr, slot: Fr): Promise<Fr> {
-    const [exists, value] = await this.publicStorage.read(storageAddress, slot);
-    this.log.debug(`storage(${storageAddress})@${slot} ?? value: ${value}, exists: ${exists}.`);
+    const { value, exists, cached } = await this.publicStorage.read(storageAddress, slot);
+    this.log.debug(`storage(${storageAddress})@${slot} ?? value: ${value}, exists: ${exists}, cached: ${cached}.`);
 
     // TRANSITIONAL: This should be removed once the kernel handles and entire enqueued call per circuit
-    this.transitionalExecutionResult.contractStorageReads.push(
-      new ContractStorageRead(slot, value, this.trace.accessCounter, storageAddress),
-    );
+    // The current info to the kernel kernel does not consider cached reads.
+    if (!cached) {
+      // The current info to the kernel removes any previous reads to the same slot.
+      this.transitionalExecutionResult.contractStorageReads =
+        this.transitionalExecutionResult.contractStorageReads.filter(
+          read => !read.storageSlot.equals(slot) || !read.contractAddress!.equals(storageAddress),
+        );
+      this.transitionalExecutionResult.contractStorageReads.push(
+        new ContractStorageRead(slot, value, this.trace.accessCounter, storageAddress),
+      );
+    }
 
     // We want to keep track of all performed reads (even reverted ones)
-    this.trace.tracePublicStorageRead(storageAddress, slot, value, exists);
+    this.trace.tracePublicStorageRead(storageAddress, slot, value, exists, cached);
     return Promise.resolve(value);
   }
 
@@ -348,6 +367,7 @@ export class AvmPersistableStateManager {
       currentStorageValue: this.publicStorage.getCache().cachePerContract,
       storageReads: this.trace.publicStorageReads,
       storageWrites: this.trace.publicStorageWrites,
+      sideEffectCounter: this.trace.accessCounter,
     };
   }
 }

yarn-project/simulator/src/avm/journal/nullifiers.ts

@@ -1,3 +1,4 @@
+import { AztecAddress } from '@aztec/circuits.js';
 import { siloNullifier } from '@aztec/circuits.js/hash';
 import { Fr } from '@aztec/foundation/fields';
 
@@ -10,7 +11,7 @@ import type { CommitmentsDB } from '../../index.js';
  */
 export class Nullifiers {
   /** Cached nullifiers. */
-  private cache: NullifierCache;
+  public cache: NullifierCache;
   /** Parent's nullifier cache. Checked on cache-miss. */
   private readonly parentCache: NullifierCache | undefined;
   /** Reference to node storage. Checked on parent cache-miss. */
@@ -95,6 +96,7 @@ export class NullifierCache {
    * each entry being a nullifier.
    */
   private cachePerContract: Map<bigint, Set<bigint>> = new Map();
+  private siloedNullifiers: Set<bigint> = new Set();
 
   /**
    * Check whether a nullifier exists in the cache.
@@ -104,8 +106,10 @@ export class NullifierCache {
    * @returns whether the nullifier is found in the cache
    */
   public exists(storageAddress: Fr, nullifier: Fr): boolean {
-    const exists = this.cachePerContract.get(storageAddress.toBigInt())?.has(nullifier.toBigInt());
-    return exists ? true : false;
+    const exists =
+      this.cachePerContract.get(storageAddress.toBigInt())?.has(nullifier.toBigInt()) ||
+      this.siloedNullifiers.has(siloNullifier(AztecAddress.fromField(storageAddress), nullifier).toBigInt());
+    return !!exists;
   }
 
   /**
@@ -115,20 +119,25 @@ export class NullifierCache {
    * @param nullifier - the nullifier to stage
    */
   public append(storageAddress: Fr, nullifier: Fr) {
+    if (this.exists(storageAddress, nullifier)) {
+      throw new NullifierCollisionError(
+        `Nullifier ${nullifier} at contract ${storageAddress} already exists in cache.`,
+      );
+    }
+
     let nullifiersForContract = this.cachePerContract.get(storageAddress.toBigInt());
     // If this contract's nullifier set has no cached nullifiers, create a new Set to store them
     if (!nullifiersForContract) {
       nullifiersForContract = new Set();
       this.cachePerContract.set(storageAddress.toBigInt(), nullifiersForContract);
     }
-    if (nullifiersForContract.has(nullifier.toBigInt())) {
-      throw new NullifierCollisionError(
-        `Nullifier ${nullifier} at contract ${storageAddress} already exists in cache.`,
-      );
-    }
     nullifiersForContract.add(nullifier.toBigInt());
   }
 
+  public appendSiloed(siloedNullifier: Fr) {
+    this.siloedNullifiers.add(siloedNullifier.toBigInt());
+  }
+
   /**
    * Merge another cache's nullifiers into this instance's.
    *
@@ -139,6 +148,8 @@ export class NullifierCache {
    * @param incomingNullifiers - the incoming cached nullifiers to merge into this instance's
    */
   public acceptAndMerge(incomingNullifiers: NullifierCache) {
+    // Merge siloed nullifiers.
+    this.siloedNullifiers = new Set([...this.siloedNullifiers, ...incomingNullifiers.siloedNullifiers]);
     // Iterate over all contracts with staged writes in the child.
     for (const [incomingAddress, incomingCacheAtContract] of incomingNullifiers.cachePerContract) {
       const thisCacheAtContract = this.cachePerContract.get(incomingAddress);

yarn-project/simulator/src/avm/journal/public_storage.test.ts

@@ -19,46 +19,54 @@ describe('avm public storage', () => {
       const contractAddress = new Fr(1);
       const slot = new Fr(2);
       // never written!
-      const [exists, gotValue] = await publicStorage.read(contractAddress, slot);
+      const { exists, value: gotValue, cached } = await publicStorage.read(contractAddress, slot);
       // doesn't exist, value is zero
       expect(exists).toEqual(false);
       expect(gotValue).toEqual(Fr.ZERO);
+      expect(cached).toEqual(false);
     });
+
     it('Should cache storage write, reading works after write', async () => {
       const contractAddress = new Fr(1);
       const slot = new Fr(2);
       const value = new Fr(3);
       // Write to cache
       publicStorage.write(contractAddress, slot, value);
-      const [exists, gotValue] = await publicStorage.read(contractAddress, slot);
+      const { exists, value: gotValue, cached } = await publicStorage.read(contractAddress, slot);
       // exists because it was previously written
       expect(exists).toEqual(true);
       expect(gotValue).toEqual(value);
+      expect(cached).toEqual(true);
     });
+
     it('Reading works on fallback to host (gets value & exists)', async () => {
       const contractAddress = new Fr(1);
       const slot = new Fr(2);
       const storedValue = new Fr(420);
       // ensure that fallback to host gets a value
       publicDb.storageRead.mockResolvedValue(Promise.resolve(storedValue));
 
-      const [exists, gotValue] = await publicStorage.read(contractAddress, slot);
+      const { exists, value: gotValue, cached } = await publicStorage.read(contractAddress, slot);
       // it exists in the host, so it must've been written before
       expect(exists).toEqual(true);
       expect(gotValue).toEqual(storedValue);
+      expect(cached).toEqual(false);
     });
+
     it('Reading works on fallback to parent (gets value & exists)', async () => {
       const contractAddress = new Fr(1);
       const slot = new Fr(2);
       const value = new Fr(3);
       const childStorage = new PublicStorage(publicDb, publicStorage);
 
       publicStorage.write(contractAddress, slot, value);
-      const [exists, gotValue] = await childStorage.read(contractAddress, slot);
+      const { exists, value: gotValue, cached } = await childStorage.read(contractAddress, slot);
       // exists because it was previously written!
       expect(exists).toEqual(true);
       expect(gotValue).toEqual(value);
+      expect(cached).toEqual(true);
     });
+
     it('When reading from storage, should check cache, then parent, then host', async () => {
       // Store a different value in storage vs the cache, and make sure the cache is returned
       const contractAddress = new Fr(1);
@@ -71,21 +79,24 @@ describe('avm public storage', () => {
       const childStorage = new PublicStorage(publicDb, publicStorage);
 
       // Cache miss falls back to host
-      const [, cacheMissResult] = await childStorage.read(contractAddress, slot);
-      expect(cacheMissResult).toEqual(storedValue);
+      const { cached: cachedCacheMiss, value: valueCacheMiss } = await childStorage.read(contractAddress, slot);
+      expect(valueCacheMiss).toEqual(storedValue);
+      expect(cachedCacheMiss).toEqual(false);
 
       // Write to storage
       publicStorage.write(contractAddress, slot, parentValue);
       // Reading from child should give value written in parent
-      const [, valueFromParent] = await childStorage.read(contractAddress, slot);
+      const { cached: cachedValueFromParent, value: valueFromParent } = await childStorage.read(contractAddress, slot);
       expect(valueFromParent).toEqual(parentValue);
+      expect(cachedValueFromParent).toEqual(true);
 
       // Now write a value directly in child
       childStorage.write(contractAddress, slot, cachedValue);
 
       // Reading should now give the value written in child
-      const [, cachedResult] = await childStorage.read(contractAddress, slot);
+      const { cached: cachedChild, value: cachedResult } = await childStorage.read(contractAddress, slot);
       expect(cachedResult).toEqual(cachedValue);
+      expect(cachedChild).toEqual(true);
     });
   });
 
@@ -109,7 +120,7 @@ describe('avm public storage', () => {
     publicStorage.acceptAndMerge(childStorage);
 
     // Read from parent gives latest value written in child before merge (valueT1)
-    const [exists, result] = await publicStorage.read(contractAddress, slot);
+    const { exists, value: result } = await publicStorage.read(contractAddress, slot);
     expect(exists).toEqual(true);
     expect(result).toEqual(valueT1);
   });

yarn-project/simulator/src/avm/journal/public_storage.ts

@@ -1,7 +1,14 @@
+import { AztecAddress } from '@aztec/circuits.js';
 import { Fr } from '@aztec/foundation/fields';
 
 import type { PublicStateDB } from '../../index.js';
 
+type PublicStorageReadResult = {
+  value: Fr;
+  exists: boolean;
+  cached: boolean;
+};
+
 /**
  * A class to manage public storage reads and writes during a contract call's AVM simulation.
  * Maintains a storage write cache, and ensures that reads fall back to the correct source.
@@ -39,7 +46,8 @@ export class PublicStorage {
    * @param slot - the slot in the contract's storage being read from
    * @returns exists: whether the slot has EVER been written to before, value: the latest value written to slot, or 0 if never written to before
    */
-  public async read(storageAddress: Fr, slot: Fr): Promise<[/*exists=*/ boolean, /*value=*/ Fr]> {
+  public async read(storageAddress: Fr, slot: Fr): Promise<PublicStorageReadResult> {
+    let cached = false;
     // First try check this storage cache
     let value = this.cache.read(storageAddress, slot);
     // Then try parent's storage cache (if it exists / written to earlier in this TX)
@@ -49,11 +57,13 @@ export class PublicStorage {
     // Finally try the host's Aztec state (a trip to the database)
     if (!value) {
       value = await this.hostPublicStorage.storageRead(storageAddress, slot);
+    } else {
+      cached = true;
     }
     // if value is undefined, that means this slot has never been written to!
     const exists = value !== undefined;
     const valueOrZero = exists ? value : Fr.ZERO;
-    return Promise.resolve([exists, valueOrZero]);
+    return Promise.resolve({ value: valueOrZero, exists, cached });
   }
 
   /**
@@ -75,6 +85,17 @@ export class PublicStorage {
   public acceptAndMerge(incomingPublicStorage: PublicStorage) {
     this.cache.acceptAndMerge(incomingPublicStorage.cache);
   }
+
+  /**
+   * Commits ALL staged writes to the host's state.
+   */
+  public async commitToDB() {
+    for (const [storageAddress, cacheAtContract] of this.cache.cachePerContract) {
+      for (const [slot, value] of cacheAtContract) {
+        await this.hostPublicStorage.storageWrite(AztecAddress.fromBigInt(storageAddress), new Fr(slot), value);
+      }
+    }
+  }
 }
 
 /**