From af80beb831bc0ad3dc8a579b019098cc5bf75f99 Mon Sep 17 00:00:00 2001 From: dbanks12 Date: Fri, 23 Feb 2024 15:38:49 +0000 Subject: [PATCH] feat(avm-simulator): create cache for pending nullifiers and existence checks --- .../src/simulator/public_executor.ts | 4 + .../simulator/src/avm/fixtures/index.ts | 24 ++- .../simulator/src/avm/journal/host_storage.ts | 16 +- .../simulator/src/avm/journal/journal.test.ts | 25 ++- .../simulator/src/avm/journal/journal.ts | 26 ++- .../src/avm/journal/nullifiers.test.ts | 147 +++++++++++++++ .../simulator/src/avm/journal/nullifiers.ts | 170 ++++++++++++++++++ .../src/avm/opcodes/accrued_substate.test.ts | 41 ++++- .../src/avm/opcodes/accrued_substate.ts | 21 ++- .../src/avm/opcodes/external_calls.test.ts | 2 +- .../simulator/src/avm/opcodes/storage.test.ts | 7 +- yarn-project/simulator/src/public/db.ts | 7 + 12 files changed, 448 insertions(+), 42 deletions(-) create mode 100644 yarn-project/simulator/src/avm/journal/nullifiers.test.ts create mode 100644 yarn-project/simulator/src/avm/journal/nullifiers.ts diff --git a/yarn-project/sequencer-client/src/simulator/public_executor.ts b/yarn-project/sequencer-client/src/simulator/public_executor.ts index a707355b751..4bf333a7ef3 100644 --- a/yarn-project/sequencer-client/src/simulator/public_executor.ts +++ b/yarn-project/sequencer-client/src/simulator/public_executor.ts @@ -166,4 +166,8 @@ export class WorldStateDB implements CommitmentsDB { public async getCommitmentIndex(commitment: Fr): Promise { return await this.db.findLeafIndex(MerkleTreeId.NOTE_HASH_TREE, commitment.toBuffer()); } + + public async getNullifierIndex(nullifier: Fr): Promise { + return await this.db.findLeafIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBuffer()); + } } diff --git a/yarn-project/simulator/src/avm/fixtures/index.ts b/yarn-project/simulator/src/avm/fixtures/index.ts index bc5487099ca..1df07a4418e 100644 --- a/yarn-project/simulator/src/avm/fixtures/index.ts +++ b/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(), mock(), mock()); - 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(), + overrides?.contractsDb || mock(), + overrides?.commitmentsDb || mock(), + ); +} + +/** Creates an empty state manager with mocked storage. */ +export function initMockPersistableStateManager(): AvmPersistableStateManager { + return new AvmPersistableStateManager(initHostStorage()); } /** diff --git a/yarn-project/simulator/src/avm/journal/host_storage.ts b/yarn-project/simulator/src/avm/journal/host_storage.ts index 330f8bf0189..6d249d49e0d 100644 --- a/yarn-project/simulator/src/avm/journal/host_storage.ts +++ b/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, + ) {} } diff --git a/yarn-project/simulator/src/avm/journal/journal.test.ts b/yarn-project/simulator/src/avm/journal/journal.test.ts index 0530fe20c1d..a4af031ebbb 100644 --- a/yarn-project/simulator/src/avm/journal/journal.test.ts +++ b/yarn-project/simulator/src/avm/journal/journal.test.ts @@ -12,8 +12,8 @@ describe('journal', () => { beforeEach(() => { publicDb = mock(); - const commitmentsDb = mock(); const contractsDb = mock(); + const commitmentsDb = mock(); const hostStorage = new HostStorage(publicDb, contractsDb, commitmentsDb); journal = new AvmPersistableStateManager(hostStorage); @@ -60,7 +60,14 @@ 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); @@ -68,14 +75,6 @@ describe('journal', () => { 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,7 +99,7 @@ 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); @@ -108,7 +107,7 @@ describe('journal', () => { childJournal.writeNoteHash(commitmentT1); childJournal.writeLog(logsT1); childJournal.writeL1Message(logsT1); - childJournal.writeNullifier(commitmentT1); + await childJournal.writeNullifier(contractAddress, commitmentT1); journal.acceptNestedCallState(childJournal); @@ -158,7 +157,7 @@ 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); @@ -166,7 +165,7 @@ describe('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); diff --git a/yarn-project/simulator/src/avm/journal/journal.ts b/yarn-project/simulator/src/avm/journal/journal.ts index fab9ec5c914..083e7881923 100644 --- a/yarn-project/simulator/src/avm/journal/journal.ts +++ b/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 { 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[]) { diff --git a/yarn-project/simulator/src/avm/journal/nullifiers.test.ts b/yarn-project/simulator/src/avm/journal/nullifiers.test.ts new file mode 100644 index 00000000000..81725755a6b --- /dev/null +++ b/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; + let nullifiers: Nullifiers; + + beforeEach(() => { + commitmentsDb = mock(); + 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()}.`, + ); + }); + }); +}); diff --git a/yarn-project/simulator/src/avm/journal/nullifiers.ts b/yarn-project/simulator/src/avm/journal/nullifiers.ts new file mode 100644 index 00000000000..f8374f6f8a7 --- /dev/null +++ b/yarn-project/simulator/src/avm/journal/nullifiers.ts @@ -0,0 +1,170 @@ +import { siloNullifier } from '@aztec/circuits.js/hash'; +import { Fr } from '@aztec/foundation/fields'; + +import type { CommitmentsDB } from '../../index.js'; + +/** + * A class to manage new nullifier staging and existence checks during a contract call's AVM simulation. + * Maintains a nullifier cache, and ensures that existence checks fall back to the correct source. + * When a contract call completes, its cached nullifier set can be merged into its parent's. + */ +export class Nullifiers { + /** Cached nullifiers. */ + private 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; + this.cache = new NullifierCache(); + } + + /** + * Get a nullifier's existence status. + * 1. Check cache. + * 2. Check parent's cache. + * 3. Fall back to the host state. + * 4. Not found! Nullifier does not exist. + * + * @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 at all, + * isPending: whether the nullifier was found in a cache, + * leafIndex: the nullifier's leaf index if it exists and is not pending (comes from host state). + */ + public async checkExists( + 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); + } + // 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; + if (!existsAsPending) { + // silo the nullifier before checking for its existence in the host + leafIndex = await this.hostNullifiers.getNullifierIndex(siloNullifier(storageAddress, nullifier)); + } + const exists = existsAsPending || leafIndex !== undefined; + leafIndex = leafIndex === undefined ? BigInt(0) : leafIndex; + return Promise.resolve([exists, existsAsPending, new Fr(leafIndex)]); + } + + /** + * Stage a new nullifier (append it to the cache). + * + * @param storageAddress - the address of the contract that the nullifier is associated with + * @param nullifier - the nullifier to stage + */ + public async append(storageAddress: Fr, nullifier: Fr) { + const [exists, ,] = await this.checkExists(storageAddress, nullifier); + if (exists) { + throw new NullifierCollisionError( + `Nullifier ${nullifier} at contract ${storageAddress} already exists in parent cache or host.`, + ); + } + this.cache.append(storageAddress, nullifier); + } + + /** + * Merges another nullifier cache into this one. + * + * @param incomingNullifiers - the incoming cached nullifiers to merge into this instance's + */ + public acceptAndMerge(incomingNullifiers: Nullifiers) { + this.cache.acceptAndMerge(incomingNullifiers.cache); + } +} + +/** + * A class to cache nullifiers created during a contract call's AVM simulation. + * "append" updates a map, "exists" checks that map. + * An instance of this class can merge another instance's cached nullifiers into its own. + */ +export class NullifierCache { + /** + * Map for staging nullifiers. + * One inner-set per contract storage address, + * each entry being a nullifier. + */ + private cachePerContract: Map> = new Map(); + + /** + * Check whether a nullifier exists in the cache. + * + * @param storageAddress - the address of the contract that the nullifier is associated with + * @param nullifier - the nullifier to check existence of + * @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; + } + + /** + * Stage a new nullifier (append it to the cache). + * + * @param storageAddress - the address of the contract that the nullifier is associated with + * @param nullifier - the nullifier to stage + */ + public append(storageAddress: Fr, nullifier: Fr) { + 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()); + } + + /** + * Merge another cache's nullifiers into this instance's. + * + * Cached nullifiers in "incoming" must not collide with any present in "this". + * + * In practice, "this" is a parent call's pending nullifiers, and "incoming" is a nested call's. + * + * @param incomingNullifiers - the incoming cached nullifiers to merge into this instance's + */ + public acceptAndMerge(incomingNullifiers: NullifierCache) { + // Iterate over all contracts with staged writes in the child. + for (const [incomingAddress, incomingCacheAtContract] of incomingNullifiers.cachePerContract) { + const thisCacheAtContract = this.cachePerContract.get(incomingAddress); + if (!thisCacheAtContract) { + // This contract has no nullifiers cached here + // so just accept incoming cache as-is for this contract. + this.cachePerContract.set(incomingAddress, incomingCacheAtContract); + } else { + // "Incoming" and "this" both have cached nullifiers for this contract. + // Merge in incoming nullifiers, erroring if there are any duplicates. + for (const nullifier of incomingCacheAtContract) { + if (thisCacheAtContract.has(nullifier)) { + throw new NullifierCollisionError( + `Failed to accept child call's nullifiers. Nullifier ${nullifier} already exists at contract ${incomingAddress}.`, + ); + } + thisCacheAtContract.add(nullifier); + } + } + } + } +} + +export class NullifierCollisionError extends Error { + constructor(message: string, ...rest: any[]) { + super(message, ...rest); + this.name = 'NullifierCollisionError'; + } +} diff --git a/yarn-project/simulator/src/avm/opcodes/accrued_substate.test.ts b/yarn-project/simulator/src/avm/opcodes/accrued_substate.test.ts index 34d53d07875..4072c61bd27 100644 --- a/yarn-project/simulator/src/avm/opcodes/accrued_substate.test.ts +++ b/yarn-project/simulator/src/avm/opcodes/accrued_substate.test.ts @@ -1,6 +1,11 @@ +import { mock } from 'jest-mock-extended'; + +import { CommitmentsDB } from '../../index.js'; import { AvmContext } from '../avm_context.js'; import { Field } from '../avm_memory_types.js'; -import { initContext, initExecutionEnvironment } from '../fixtures/index.js'; +import { InstructionExecutionError } from '../errors.js'; +import { initContext, initExecutionEnvironment, initHostStorage } from '../fixtures/index.js'; +import { AvmPersistableStateManager } from '../journal/journal.js'; import { EmitNoteHash, EmitNullifier, EmitUnencryptedLog, SendL2ToL1Message } from './accrued_substate.js'; import { StaticCallStorageAlterError } from './storage.js'; @@ -59,6 +64,40 @@ describe('Accrued Substate', () => { const expected = [value.toFr()]; expect(journalState.newNullifiers).toEqual(expected); }); + + it('Nullifier collision reverts (same nullifier emitted twice)', async () => { + const value = new Field(69n); + context.machineState.memory.set(0, value); + + await new EmitNullifier(/*indirect=*/ 0, /*offset=*/ 0).execute(context); + await expect(new EmitNullifier(/*indirect=*/ 0, /*offset=*/ 0).execute(context)).rejects.toThrowError( + new InstructionExecutionError( + `Attempted to emit duplicate nullifier ${value.toFr()} (storage address: ${ + context.environment.storageAddress + }).`, + ), + ); + }); + + it('Nullifier collision reverts (nullifier exists in host state)', async () => { + const value = new Field(69n); + const storedLeafIndex = BigInt(42); + + // Mock the nullifiers db to return a stored leaf index + const commitmentsDb = mock(); + commitmentsDb.getNullifierIndex.mockResolvedValue(Promise.resolve(storedLeafIndex)); + const hostStorage = initHostStorage({ commitmentsDb }); + context = initContext({ persistableState: new AvmPersistableStateManager(hostStorage) }); + + context.machineState.memory.set(0, value); + await expect(new EmitNullifier(/*indirect=*/ 0, /*offset=*/ 0).execute(context)).rejects.toThrowError( + new InstructionExecutionError( + `Attempted to emit duplicate nullifier ${value.toFr()} (storage address: ${ + context.environment.storageAddress + }).`, + ), + ); + }); }); describe('EmitUnencryptedLog', () => { diff --git a/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts b/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts index c7d1a13e86a..503c4313a51 100644 --- a/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts +++ b/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts @@ -1,4 +1,6 @@ import type { AvmContext } from '../avm_context.js'; +import { InstructionExecutionError } from '../errors.js'; +import { NullifierCollisionError } from '../journal/nullifiers.js'; import { Opcode, OperandType } from '../serialization/instruction_serialization.js'; import { Instruction } from './instruction.js'; import { StaticCallStorageAlterError } from './storage.js'; @@ -41,7 +43,24 @@ export class EmitNullifier extends Instruction { } const nullifier = context.machineState.memory.get(this.nullifierOffset).toFr(); - context.persistableState.writeNullifier(nullifier); + const exists = await context.persistableState.checkNullifierExists(context.environment.storageAddress, nullifier); + if (exists) { + throw new InstructionExecutionError( + `Attempted to emit duplicate nullifier ${nullifier} (storage address: ${context.environment.storageAddress}).`, + ); + } + try { + await context.persistableState.writeNullifier(context.environment.storageAddress, nullifier); + } catch (e) { + if (e instanceof NullifierCollisionError) { + // Error is known/expected, raise as InstructionExecutionError that the will lead the simulator to revert this call + throw new InstructionExecutionError( + `Attempted to emit duplicate nullifier ${nullifier} (storage address: ${context.environment.storageAddress}).`, + ); + } else { + throw e; + } + } context.machineState.incrementPc(); } diff --git a/yarn-project/simulator/src/avm/opcodes/external_calls.test.ts b/yarn-project/simulator/src/avm/opcodes/external_calls.test.ts index a32678a3c80..e557f324278 100644 --- a/yarn-project/simulator/src/avm/opcodes/external_calls.test.ts +++ b/yarn-project/simulator/src/avm/opcodes/external_calls.test.ts @@ -25,7 +25,7 @@ describe('External Calls', () => { const publicStateDb = mock(); const hostStorage = new HostStorage(publicStateDb, contractsDb, commitmentsDb); const journal = new AvmPersistableStateManager(hostStorage); - context = initContext({ worldState: journal }); + context = initContext({ persistableState: journal }); }); describe('Call', () => { diff --git a/yarn-project/simulator/src/avm/opcodes/storage.test.ts b/yarn-project/simulator/src/avm/opcodes/storage.test.ts index 7776b91f1cd..bd53a1d3324 100644 --- a/yarn-project/simulator/src/avm/opcodes/storage.test.ts +++ b/yarn-project/simulator/src/avm/opcodes/storage.test.ts @@ -16,7 +16,10 @@ describe('Storage Instructions', () => { beforeEach(async () => { journal = mock(); - context = initContext({ worldState: journal, env: initExecutionEnvironment({ address, storageAddress: address }) }); + context = initContext({ + persistableState: journal, + env: initExecutionEnvironment({ address, storageAddress: address }), + }); }); describe('SSTORE', () => { @@ -47,7 +50,7 @@ describe('Storage Instructions', () => { it('Should not be able to write to storage in a static call', async () => { context = initContext({ - worldState: journal, + persistableState: journal, env: initExecutionEnvironment({ address, storageAddress: address, isStaticCall: true }), }); diff --git a/yarn-project/simulator/src/public/db.ts b/yarn-project/simulator/src/public/db.ts index 26923cc366e..7701e32ea6e 100644 --- a/yarn-project/simulator/src/public/db.ts +++ b/yarn-project/simulator/src/public/db.ts @@ -82,4 +82,11 @@ export interface CommitmentsDB { * @returns - The index of the commitment. Undefined if it does not exist in the tree. */ getCommitmentIndex(commitment: Fr): Promise; + + /** + * Gets the index of a nullifier in the nullifier tree. + * @param nullifier - The nullifier. + * @returns - The index of the nullifier. Undefined if it does not exist in the tree. + */ + getNullifierIndex(nullifier: Fr): Promise; }