feat(avm-simulator): create cache for pending nullifiers and existenc…

…e checks

yarn-project/sequencer-client/src/sequencer/public_processor.ts

@@ -9,7 +9,12 @@ import { MerkleTreeOperations } from '@aztec/world-state';
 import { EmptyPublicProver } from '../prover/empty.js';
 import { PublicProver } from '../prover/index.js';
 import { PublicKernelCircuitSimulator } from '../simulator/index.js';
-import { ContractsDataSourcePublicDB, WorldStateDB, WorldStatePublicDB } from '../simulator/public_executor.js';
+import {
+  ContractsDataSourcePublicDB,
+  NullifiersDB,
+  WorldStateDB,
+  WorldStatePublicDB,
+} from '../simulator/public_executor.js';
 import { RealPublicKernelCircuitSimulator } from '../simulator/public_kernel.js';
 import { AbstractPhaseManager } from './abstract_phase_manager.js';
 import { PhaseManagerFactory } from './phase_manager_factory.js';
@@ -41,7 +46,14 @@ export class PublicProcessorFactory {
     const publicContractsDB = new ContractsDataSourcePublicDB(this.contractDataSource);
     const worldStatePublicDB = new WorldStatePublicDB(this.merkleTree);
     const worldStateDB = new WorldStateDB(this.merkleTree, this.l1Tol2MessagesDataSource);
-    const publicExecutor = new PublicExecutor(worldStatePublicDB, publicContractsDB, worldStateDB, historicalHeader);
+    const nullifiersDB = new NullifiersDB(this.merkleTree);
+    const publicExecutor = new PublicExecutor(
+      worldStatePublicDB,
+      publicContractsDB,
+      worldStateDB,
+      nullifiersDB,
+      historicalHeader,
+    );
     return new PublicProcessor(
       this.merkleTree,
       publicExecutor,

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

@@ -8,7 +8,7 @@ import {
   PublicDataTreeLeafPreimage,
 } from '@aztec/circuits.js';
 import { computePublicDataTreeLeafSlot } from '@aztec/circuits.js/hash';
-import { CommitmentsDB, MessageLoadOracleInputs, PublicContractsDB, PublicStateDB } from '@aztec/simulator';
+import { CommitmentsDB, MessageLoadOracleInputs, NullifiersDB as NDB, PublicContractsDB, PublicStateDB } from '@aztec/simulator';
 import { MerkleTreeOperations } from '@aztec/world-state';
 
 /**
@@ -167,3 +167,14 @@ export class WorldStateDB implements CommitmentsDB {
     return await this.db.findLeafIndex(MerkleTreeId.NOTE_HASH_TREE, commitment.toBuffer());
   }
 }
+
+/**
+ * Implements Nullifier db.
+ */
+export class NullifiersDB implements NDB {
+  constructor(private db: MerkleTreeOperations) {}
+
+  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

@@ -7,7 +7,7 @@ import { Fr } from '@aztec/foundation/fields';
 import { mock } from 'jest-mock-extended';
 import merge from 'lodash.merge';
 
-import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../index.js';
+import { CommitmentsDB, NullifiersDB, PublicContractsDB, PublicStateDB } from '../../index.js';
 import { AvmContext } from '../avm_context.js';
 import { AvmExecutionEnvironment } from '../avm_execution_environment.js';
 import { AvmMachineState } from '../avm_machine_state.js';
@@ -31,7 +31,12 @@ export function initContext(overrides?: {
 
 /** Creates an empty world state with mocked storage. */
 export function initMockWorldStateJournal(): AvmPersistableStateManager {
-  const hostStorage = new HostStorage(mock<PublicStateDB>(), mock<PublicContractsDB>(), mock<CommitmentsDB>());
+  const hostStorage = new HostStorage(
+    mock<PublicStateDB>(),
+    mock<PublicContractsDB>(),
+    mock<CommitmentsDB>(),
+    mock<NullifiersDB>(),
+  );
   return new AvmPersistableStateManager(hostStorage);
 }
 

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

@@ -1,20 +1,15 @@
-import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../public/db.js';
+import { CommitmentsDB, NullifiersDB, PublicContractsDB, PublicStateDB } from '../../public/db.js';
 
 /**
  * Host storage
  *
  * 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,
+    public readonly nullifiersDb: NullifiersDB,
+  ) {}
 }

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

@@ -2,7 +2,7 @@ import { Fr } from '@aztec/foundation/fields';
 
 import { MockProxy, mock } from 'jest-mock-extended';
 
-import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../index.js';
+import { CommitmentsDB, NullifiersDB, PublicContractsDB, PublicStateDB } from '../../index.js';
 import { HostStorage } from './host_storage.js';
 import { AvmPersistableStateManager, JournalData } from './journal.js';
 
@@ -12,10 +12,11 @@ describe('journal', () => {
 
   beforeEach(() => {
     publicDb = mock<PublicStateDB>();
-    const commitmentsDb = mock<CommitmentsDB>();
     const contractsDb = mock<PublicContractsDB>();
+    const commitmentsDb = mock<CommitmentsDB>();
+    const nullifiersDb = mock<NullifiersDB>();
 
-    const hostStorage = new HostStorage(publicDb, contractsDb, commitmentsDb);
+    const hostStorage = new HostStorage(publicDb, contractsDb, commitmentsDb, nullifiersDb);
     journal = new AvmPersistableStateManager(hostStorage);
   });
 
@@ -60,22 +61,20 @@ describe('journal', () => {
       const journalUpdates = journal.flush();
       expect(journalUpdates.newNoteHashes).toEqual([utxo]);
     });
+    it('Should maintain nullifiers', async () => {
+      const utxo = new Fr(1);
+      await journal.writeNullifier(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(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(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(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(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.nullifiersDb, 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,8 +96,10 @@ export class AvmPersistableStateManager {
     this.trace.traceNewNoteHash(/*storageAddress*/ Fr.ZERO, noteHash);
   }
 
-  public writeNullifier(nullifier: Fr) {
-    // TODO track pending nullifiers in set per-contract
+  public async writeNullifier(nullifier: Fr) {
+    // Cache pending nullifiers for later access
+    await this.nullifiers.append(/*storageAddress*/ Fr.ZERO, nullifier);
+    // Trace all nullifier creations (even reverted ones)
     this.trace.traceNewNullifier(/*storageAddress*/ Fr.ZERO, nullifier);
   }