feat(avm-simulator): create cache for pending nullifiers and existence checks

yarn-project/sequencer-client/src/simulator/public_executor.ts

@@ -166,4 +166,8 @@ export class WorldStateDB implements CommitmentsDB {
   public async getCommitmentIndex(commitment: Fr): Promise<bigint | undefined> {
     return await this.db.findLeafIndex(MerkleTreeId.NOTE_HASH_TREE, commitment.toBuffer());
   }
+
+  public async getNullifierIndex(nullifier: Fr): Promise<bigint | undefined> {
+    return await this.db.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBuffer());
+  }
 }

yarn-project/simulator/src/avm/fixtures/index.ts

@@ -18,21 +18,33 @@ import { AvmPersistableStateManager } from '../journal/journal.js';
  * Create a new AVM context with default values.
  */
 export function initContext(overrides?: {
-  worldState?: AvmPersistableStateManager;
+  persistableState?: AvmPersistableStateManager;
   env?: AvmExecutionEnvironment;
   machineState?: AvmMachineState;
 }): AvmContext {
   return new AvmContext(
-    overrides?.worldState || initMockWorldStateJournal(),
+    overrides?.persistableState || initMockPersistableStateManager(),
     overrides?.env || initExecutionEnvironment(),
     overrides?.machineState || initMachineState(),
   );
 }
 
-/** Creates an empty world state with mocked storage. */
-export function initMockWorldStateJournal(): AvmPersistableStateManager {
-  const hostStorage = new HostStorage(mock<PublicStateDB>(), mock<PublicContractsDB>(), mock<CommitmentsDB>());
-  return new AvmPersistableStateManager(hostStorage);
+/** Creates an empty host storage with mocked dbs. */
+export function initHostStorage(overrides?: {
+  publicDb?: PublicStateDB;
+  contractsDb?: PublicContractsDB;
+  commitmentsDb?: CommitmentsDB;
+}): HostStorage {
+  return new HostStorage(
+    overrides?.publicDb || mock<PublicStateDB>(),
+    overrides?.contractsDb || mock<PublicContractsDB>(),
+    overrides?.commitmentsDb || mock<CommitmentsDB>(),
+  );
+}
+
+/** Creates an empty state manager with mocked storage. */
+export function initMockPersistableStateManager(): AvmPersistableStateManager {
+  return new AvmPersistableStateManager(initHostStorage());
 }
 
 /**

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

@@ -6,15 +6,9 @@ import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../public/db
  * A wrapper around the node dbs
  */
 export class HostStorage {
-  public readonly publicStateDb: PublicStateDB;
-
-  public readonly contractsDb: PublicContractsDB;
-
-  public readonly commitmentsDb: CommitmentsDB;
-
-  constructor(publicStateDb: PublicStateDB, contractsDb: PublicContractsDB, commitmentsDb: CommitmentsDB) {
-    this.publicStateDb = publicStateDb;
-    this.contractsDb = contractsDb;
-    this.commitmentsDb = commitmentsDb;
-  }
+  constructor(
+    public readonly publicStateDb: PublicStateDB,
+    public readonly contractsDb: PublicContractsDB,
+    public readonly commitmentsDb: CommitmentsDB,
+  ) {}
 }

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

@@ -12,8 +12,8 @@ describe('journal', () => {
 
   beforeEach(() => {
     publicDb = mock<PublicStateDB>();
-    const commitmentsDb = mock<CommitmentsDB>();
     const contractsDb = mock<PublicContractsDB>();
+    const commitmentsDb = mock<CommitmentsDB>();
 
     const hostStorage = new HostStorage(publicDb, contractsDb, commitmentsDb);
     journal = new AvmPersistableStateManager(hostStorage);
@@ -60,22 +60,21 @@ describe('journal', () => {
       const journalUpdates = journal.flush();
       expect(journalUpdates.newNoteHashes).toEqual([utxo]);
     });
+    it('Should maintain nullifiers', async () => {
+      const contractAddress = new Fr(1);
+      const utxo = new Fr(2);
+      await journal.writeNullifier(contractAddress, utxo);
 
+      const journalUpdates = journal.flush();
+      expect(journalUpdates.newNullifiers).toEqual([utxo]);
+    });
     it('Should maintain l1 messages', () => {
       const utxo = [new Fr(1)];
       journal.writeL1Message(utxo);
 
       const journalUpdates = journal.flush();
       expect(journalUpdates.newL1Messages).toEqual([utxo]);
     });
-
-    it('Should maintain nullifiers', () => {
-      const utxo = new Fr(1);
-      journal.writeNullifier(utxo);
-
-      const journalUpdates = journal.flush();
-      expect(journalUpdates.newNullifiers).toEqual([utxo]);
-    });
   });
 
   it('Should merge two successful journals together', async () => {
@@ -100,15 +99,15 @@ describe('journal', () => {
     journal.writeNoteHash(commitment);
     journal.writeLog(logs);
     journal.writeL1Message(logs);
-    journal.writeNullifier(commitment);
+    await journal.writeNullifier(contractAddress, commitment);
 
     const childJournal = new AvmPersistableStateManager(journal.hostStorage, journal);
     childJournal.writeStorage(contractAddress, key, valueT1);
     await childJournal.readStorage(contractAddress, key);
     childJournal.writeNoteHash(commitmentT1);
     childJournal.writeLog(logsT1);
     childJournal.writeL1Message(logsT1);
-    childJournal.writeNullifier(commitmentT1);
+    await childJournal.writeNullifier(contractAddress, commitmentT1);
 
     journal.acceptNestedCallState(childJournal);
 
@@ -158,15 +157,15 @@ describe('journal', () => {
     journal.writeStorage(contractAddress, key, value);
     await journal.readStorage(contractAddress, key);
     journal.writeNoteHash(commitment);
-    journal.writeNullifier(commitment);
+    await journal.writeNullifier(contractAddress, commitment);
     journal.writeLog(logs);
     journal.writeL1Message(logs);
 
     const childJournal = new AvmPersistableStateManager(journal.hostStorage, journal);
     childJournal.writeStorage(contractAddress, key, valueT1);
     await childJournal.readStorage(contractAddress, key);
     childJournal.writeNoteHash(commitmentT1);
-    childJournal.writeNullifier(commitmentT1);
+    await childJournal.writeNullifier(contractAddress, commitmentT1);
     childJournal.writeLog(logsT1);
     childJournal.writeL1Message(logsT1);
 

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

@@ -1,6 +1,7 @@
 import { Fr } from '@aztec/foundation/fields';
 
 import { HostStorage } from './host_storage.js';
+import { Nullifiers } from './nullifiers.js';
 import { PublicStorage } from './public_storage.js';
 import { WorldStateAccessTrace } from './trace.js';
 
@@ -10,6 +11,7 @@ import { WorldStateAccessTrace } from './trace.js';
 export type JournalData = {
   newNoteHashes: Fr[];
   newNullifiers: Fr[];
+
   newL1Messages: Fr[][];
   newLogs: Fr[][];
 
@@ -38,8 +40,8 @@ export class AvmPersistableStateManager {
   /** World State */
   /** Public storage, including cached writes */
   private publicStorage: PublicStorage;
-  ///** Nullifier set, including cached/recently-emitted nullifiers */
-  //private nullifiers: NullifiersDB;
+  /** Nullifier set, including cached/recently-emitted nullifiers */
+  private nullifiers: Nullifiers;
 
   /** World State Access Trace */
   private trace: WorldStateAccessTrace;
@@ -51,6 +53,7 @@ export class AvmPersistableStateManager {
   constructor(hostStorage: HostStorage, parent?: AvmPersistableStateManager) {
     this.hostStorage = hostStorage;
     this.publicStorage = new PublicStorage(hostStorage.publicStateDb, parent?.publicStorage);
+    this.nullifiers = new Nullifiers(hostStorage.commitmentsDb, parent?.nullifiers);
     this.trace = new WorldStateAccessTrace(parent?.trace);
   }
 
@@ -69,8 +72,9 @@ export class AvmPersistableStateManager {
    * @param value - the value being written to the slot
    */
   public writeStorage(storageAddress: Fr, slot: Fr, value: Fr) {
+    // Cache storage writes for later reference/reads
     this.publicStorage.write(storageAddress, slot, value);
-    // We want to keep track of all performed writes in the journal
+    // Trace all storage writes (even reverted ones)
     this.trace.tracePublicStorageWrite(storageAddress, slot, value);
   }
 
@@ -83,7 +87,7 @@ export class AvmPersistableStateManager {
    */
   public async readStorage(storageAddress: Fr, slot: Fr): Promise<Fr> {
     const [_exists, value] = await this.publicStorage.read(storageAddress, slot);
-    // We want to keep track of all performed reads
+    // We want to keep track of all performed reads (even reverted ones)
     this.trace.tracePublicStorageRead(storageAddress, slot, value);
     return Promise.resolve(value);
   }
@@ -92,9 +96,17 @@ export class AvmPersistableStateManager {
     this.trace.traceNewNoteHash(/*storageAddress*/ Fr.ZERO, noteHash);
   }
 
-  public writeNullifier(nullifier: Fr) {
-    // TODO track pending nullifiers in set per-contract
-    this.trace.traceNewNullifier(/*storageAddress*/ Fr.ZERO, nullifier);
+  public async checkNullifierExists(storageAddress: Fr, nullifier: Fr) {
+    const [exists, _isPending, _leafIndex] = await this.nullifiers.checkExists(storageAddress, nullifier);
+    //this.trace.traceNullifierCheck(storageAddress, nullifier, exists, isPending, leafIndex);
+    return Promise.resolve(exists);
+  }
+
+  public async writeNullifier(storageAddress: Fr, nullifier: Fr) {
+    // Cache pending nullifiers for later access
+    await this.nullifiers.append(storageAddress, nullifier);
+    // Trace all nullifier creations (even reverted ones)
+    this.trace.traceNewNullifier(storageAddress, nullifier);
   }
 
   public writeL1Message(message: Fr[]) {

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

@@ -0,0 +1,147 @@
+import { Fr } from '@aztec/foundation/fields';
+
+import { MockProxy, mock } from 'jest-mock-extended';
+
+import { CommitmentsDB } from '../../index.js';
+import { Nullifiers } from './nullifiers.js';
+
+describe('avm nullifier caching', () => {
+  let commitmentsDb: MockProxy<CommitmentsDB>;
+  let nullifiers: Nullifiers;
+
+  beforeEach(() => {
+    commitmentsDb = mock<CommitmentsDB>();
+    nullifiers = new Nullifiers(commitmentsDb);
+  });
+
+  describe('Nullifier caching and existence checks', () => {
+    it('Reading a non-existent nullifier works (gets zero & DNE)', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2);
+      // never written!
+      const [exists, isPending, gotIndex] = await nullifiers.checkExists(contractAddress, nullifier);
+      // doesn't exist, not pending, index is zero (non-existent)
+      expect(exists).toEqual(false);
+      expect(isPending).toEqual(false);
+      expect(gotIndex).toEqual(Fr.ZERO);
+    });
+    it('Should cache nullifier, existence check works after creation', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2);
+
+      // Write to cache
+      await nullifiers.append(contractAddress, nullifier);
+      const [exists, isPending, gotIndex] = await nullifiers.checkExists(contractAddress, nullifier);
+      // exists (in cache), isPending, index is zero (not in tree)
+      expect(exists).toEqual(true);
+      expect(isPending).toEqual(true);
+      expect(gotIndex).toEqual(Fr.ZERO);
+    });
+    it('Existence check works on fallback to host (gets index, exists, not-pending)', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2);
+      const storedLeafIndex = BigInt(420);
+
+      commitmentsDb.getNullifierIndex.mockResolvedValue(Promise.resolve(storedLeafIndex));
+
+      const [exists, isPending, gotIndex] = await nullifiers.checkExists(contractAddress, nullifier);
+      // exists (in host), not pending, tree index retrieved from host
+      expect(exists).toEqual(true);
+      expect(isPending).toEqual(false);
+      expect(gotIndex).toEqual(gotIndex);
+    });
+    it('Existence check works on fallback to parent (gets value, exists, is pending)', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2);
+      const childNullifiers = new Nullifiers(commitmentsDb, nullifiers);
+
+      // Write to parent cache
+      await nullifiers.append(contractAddress, nullifier);
+      // Get from child cache
+      const [exists, isPending, gotIndex] = await childNullifiers.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', () => {
+    it('Cant append nullifier that already exists in cache', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2); // same nullifier for both!
+
+      // Append a nullifier to cache
+      await nullifiers.append(contractAddress, nullifier);
+      // Can't append again
+      await expect(nullifiers.append(contractAddress, nullifier)).rejects.toThrowError(
+        `Nullifier ${nullifier} at contract ${contractAddress} already exists in parent cache or host.`,
+      );
+    });
+    it('Cant append nullifier that already exists in parent cache', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2); // same nullifier for both!
+
+      // Append a nullifier to parent
+      await nullifiers.append(contractAddress, nullifier);
+      const childNullifiers = new Nullifiers(commitmentsDb, nullifiers);
+      // Can't append again in child
+      await expect(childNullifiers.append(contractAddress, nullifier)).rejects.toThrowError(
+        `Nullifier ${nullifier} at contract ${contractAddress} already exists in parent cache or host.`,
+      );
+    });
+    it('Cant append nullifier that already exist in host', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2); // same nullifier for both!
+      const storedLeafIndex = BigInt(420);
+
+      // Nullifier exists in host
+      commitmentsDb.getNullifierIndex.mockResolvedValue(Promise.resolve(storedLeafIndex));
+      // Can't append to cache
+      await expect(nullifiers.append(contractAddress, nullifier)).rejects.toThrowError(
+        `Nullifier ${nullifier} at contract ${contractAddress} already exists in parent cache or host.`,
+      );
+    });
+  });
+
+  describe('Nullifier cache merging', () => {
+    it('Should be able to merge two nullifier caches together', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier0 = new Fr(2);
+      const nullifier1 = new Fr(3);
+
+      // Append a nullifier to parent
+      await nullifiers.append(contractAddress, nullifier0);
+
+      const childNullifiers = new Nullifiers(commitmentsDb, nullifiers);
+      // Append a nullifier to child
+      await childNullifiers.append(contractAddress, nullifier1);
+
+      // Parent accepts child's nullifiers
+      nullifiers.acceptAndMerge(childNullifiers);
+
+      // After merge, parent has both nullifiers
+      const results0 = await nullifiers.checkExists(contractAddress, nullifier0);
+      expect(results0).toEqual([/*exists=*/ true, /*isPending=*/ true, /*leafIndex=*/ Fr.ZERO]);
+      const results1 = await nullifiers.checkExists(contractAddress, nullifier1);
+      expect(results1).toEqual([/*exists=*/ true, /*isPending=*/ true, /*leafIndex=*/ Fr.ZERO]);
+    });
+    it('Cant merge two nullifier caches with colliding entries', async () => {
+      const contractAddress = new Fr(1);
+      const nullifier = new Fr(2);
+
+      // Append a nullifier to parent
+      await nullifiers.append(contractAddress, nullifier);
+
+      // Create child cache, don't derive from parent so we can concoct a collision on merge
+      const childNullifiers = new Nullifiers(commitmentsDb);
+      // Append a nullifier to child
+      await childNullifiers.append(contractAddress, nullifier);
+
+      // Parent accepts child's nullifiers
+      expect(() => nullifiers.acceptAndMerge(childNullifiers)).toThrowError(
+        `Failed to accept child call's nullifiers. Nullifier ${nullifier.toBigInt()} already exists at contract ${contractAddress.toBigInt()}.`,
+      );
+    });
+  });
+});