fix(avm-simulator): pending storage and nullifiers should be accessible in grandchild nested calls

yarn-project/end-to-end/src/shared/uniswap_l1_l2.ts

@@ -940,7 +940,7 @@ export const uniswapL1L2TestSuite = (
         uniswapPortal.simulate.swapPublic(swapArgs, {
           account: ownerEthAddress.toString(),
         } as any),
-      ).rejects.toThrow('The contract function "swapPublic" reverted.');
+      ).rejects.toThrow('The contract function "swapPublic" reverted');
     });
 
     it("can't call swap_private on L1 if called swap_public on L2", async () => {

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

@@ -141,7 +141,7 @@ export class AvmPersistableStateManager {
    * @param value - the value being written to the slot
    */
   public writeStorage(storageAddress: Fr, slot: Fr, value: Fr) {
-    this.log.debug(`storage(${storageAddress})@${slot} <- ${value}`);
+    this.log.debug(`Storage write (address=${storageAddress}, slot=${slot}): value=${value}`);
     // Cache storage writes for later reference/reads
     this.publicStorage.write(storageAddress, slot, value);
 
@@ -172,7 +172,9 @@ export class AvmPersistableStateManager {
    */
   public async readStorage(storageAddress: Fr, slot: Fr): Promise<Fr> {
     const { value, exists, cached } = await this.publicStorage.read(storageAddress, slot);
-    this.log.debug(`storage(${storageAddress})@${slot} ?? value: ${value}, exists: ${exists}, cached: ${cached}.`);
+    this.log.debug(
+      `Storage read  (address=${storageAddress}, slot=${slot}): value=${value}, exists=${exists}, cached=${cached}`,
+    );
 
     // TRANSITIONAL: This should be removed once the kernel handles and entire enqueued call per circuit
     // The current info to the kernel kernel does not consider cached reads.

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

@@ -64,6 +64,21 @@ describe('avm nullifier caching', () => {
       expect(isPending).toEqual(true);
       expect(gotIndex).toEqual(Fr.ZERO);
     });
+    it('Existence check works on fallback to grandparent (gets value, exists, is pending)', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2);
+      const childNullifiers = new Nullifiers(commitmentsDb, nullifiers);
+      const grandChildNullifiers = new Nullifiers(commitmentsDb, childNullifiers);
+
+      // Write to parent cache
+      await nullifiers.append(contractAddress, nullifier);
+      // Get from child cache
+      const [exists, isPending, gotIndex] = await grandChildNullifiers.checkExists(contractAddress, nullifier);
+      // exists (in parent), isPending, index is zero (not in tree)
+      expect(exists).toEqual(true);
+      expect(isPending).toEqual(true);
+      expect(gotIndex).toEqual(Fr.ZERO);
+    });
   });
 
   describe('Nullifier collision failures', () => {

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

@@ -12,21 +12,38 @@ import type { CommitmentsDB } from '../../index.js';
 export class Nullifiers {
   /** Cached nullifiers. */
   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. */
-  private readonly hostNullifiers: CommitmentsDB;
 
-  constructor(hostNullifiers: CommitmentsDB, parent?: Nullifiers) {
-    this.hostNullifiers = hostNullifiers;
-    this.parentCache = parent?.cache;
+  constructor(
+    /** Reference to node storage. Checked on parent cache-miss. */
+    private readonly hostNullifiers: CommitmentsDB,
+    /** Parent's nullifiers. Checked on this' cache-miss. */
+    private readonly parent?: Nullifiers | undefined,
+  ) {
     this.cache = new NullifierCache();
   }
 
+  /**
+   * Get a nullifier's existence in this' cache or parent's (recursively).
+   * DOES NOT CHECK HOST STORAGE!
+   * @param storageAddress - the address of the contract whose storage is being read from
+   * @param nullifier - the nullifier to check for
+   * @returns exists: whether the nullifier exists in cache here or in parent's
+   */
+  private checkExistsHereOrParent(storageAddress: Fr, nullifier: Fr): boolean {
+    // First check this cache
+    let existsAsPending = this.cache.exists(storageAddress, nullifier);
+    // Then try parent's nullifier cache
+    if (!existsAsPending && this.parent) {
+      // Note: this will recurse to grandparent/etc until a cache-hit is encountered.
+      existsAsPending = this.parent.checkExistsHereOrParent(storageAddress, nullifier);
+    }
+    return existsAsPending;
+  }
+
   /**
    * Get a nullifier's existence status.
    * 1. Check cache.
-   * 2. Check parent's cache.
+   * 2. Check parent cache.
    * 3. Fall back to the host state.
    * 4. Not found! Nullifier does not exist.
    *
@@ -40,12 +57,8 @@ export class Nullifiers {
     storageAddress: Fr,
     nullifier: Fr,
   ): Promise<[/*exists=*/ boolean, /*isPending=*/ boolean, /*leafIndex=*/ Fr]> {
-    // First check this cache
-    let existsAsPending = this.cache.exists(storageAddress, nullifier);
-    // Then check parent's cache
-    if (!existsAsPending && this.parentCache) {
-      existsAsPending = this.parentCache?.exists(storageAddress, nullifier);
-    }
+    // Check this cache and parent's (recursively)
+    const existsAsPending = this.checkExistsHereOrParent(storageAddress, nullifier);
     // Finally try the host's Aztec state (a trip to the database)
     // If the value is found in the database, it will be associated with a leaf index!
     let leafIndex: bigint | undefined = undefined;

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

@@ -67,6 +67,21 @@ describe('avm public storage', () => {
       expect(cached).toEqual(true);
     });
 
+    it('Reading works on fallback to grandparent (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);
+      const grandChildStorage = new PublicStorage(publicDb, childStorage);
+
+      publicStorage.write(contractAddress, slot, value);
+      const { exists, value: gotValue, cached } = await grandChildStorage.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);

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

@@ -17,14 +17,13 @@ type PublicStorageReadResult = {
 export class PublicStorage {
   /** Cached storage writes. */
   private cache: PublicStorageCache;
-  /** Parent's storage cache. Checked on cache-miss. */
-  private readonly parentCache: PublicStorageCache | undefined;
-  /** Reference to node storage. Checked on parent cache-miss. */
-  private readonly hostPublicStorage: PublicStateDB;
 
-  constructor(hostPublicStorage: PublicStateDB, parent?: PublicStorage) {
-    this.hostPublicStorage = hostPublicStorage;
-    this.parentCache = parent?.cache;
+  constructor(
+    /** Reference to node storage. Checked on parent cache-miss. */
+    private readonly hostPublicStorage: PublicStateDB,
+    /** Parent's storage. Checked on this' cache-miss. */
+    private readonly parent?: PublicStorage,
+  ) {
     this.cache = new PublicStorageCache();
   }
 
@@ -35,10 +34,29 @@ export class PublicStorage {
     return this.cache;
   }
 
+  /**
+   * Read a storage value from this' cache or parent's (recursively).
+   * DOES NOT CHECK HOST STORAGE!
+   *
+   * @param storageAddress - the address of the contract whose storage is being read from
+   * @param slot - the slot in the contract's storage being read from
+   * @returns value: the latest value written according to this cache or the parent's. undefined on cache miss.
+   */
+  public readHereOrParent(storageAddress: Fr, slot: Fr): Fr | undefined {
+    // First try check this storage cache
+    let value = this.cache.read(storageAddress, slot);
+    // Then try parent's storage cache
+    if (!value && this.parent) {
+      // Note: this will recurse to grandparent/etc until a cache-hit is encountered.
+      value = this.parent.readHereOrParent(storageAddress, slot);
+    }
+    return value;
+  }
+
   /**
    * Read a value from storage.
    * 1. Check cache.
-   * 2. Check parent's cache.
+   * 2. Check parent cache.
    * 3. Fall back to the host state.
    * 4. Not found! Value has never been written to before. Flag it as non-existent and return value zero.
    *
@@ -48,12 +66,8 @@ export class PublicStorage {
    */
   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)
-    if (!value && this.parentCache) {
-      value = this.parentCache?.read(storageAddress, slot);
-    }
+    // Check this cache and parent's (recursively)
+    let value = this.readHereOrParent(storageAddress, slot);
     // Finally try the host's Aztec state (a trip to the database)
     if (!value) {
       value = await this.hostPublicStorage.storageRead(storageAddress, slot);
@@ -73,8 +87,8 @@ export class PublicStorage {
    * @param slot - the slot in the contract's storage being written to
    * @param value - the value being written to the slot
    */
-  public write(storageAddress: Fr, key: Fr, value: Fr) {
-    this.cache.write(storageAddress, key, value);
+  public write(storageAddress: Fr, slot: Fr, value: Fr) {
+    this.cache.write(storageAddress, slot, value);
   }
 
   /**