From b075401e53a6dbe95c413608fc3c30bf19648103 Mon Sep 17 00:00:00 2001 From: David Banks <47112877+dbanks12@users.noreply.github.com> Date: Fri, 23 Feb 2024 15:28:50 -0500 Subject: [PATCH] chore(avm-simulator): pull out public storage caching and merging from the state journal (#4730) --- .../simulator/src/avm/journal/journal.test.ts | 41 ----- .../simulator/src/avm/journal/journal.ts | 97 +++--------- .../src/avm/journal/public_storage.test.ts | 116 ++++++++++++++ .../src/avm/journal/public_storage.ts | 149 ++++++++++++++++++ 4 files changed, 289 insertions(+), 114 deletions(-) create mode 100644 yarn-project/simulator/src/avm/journal/public_storage.test.ts create mode 100644 yarn-project/simulator/src/avm/journal/public_storage.ts diff --git a/yarn-project/simulator/src/avm/journal/journal.test.ts b/yarn-project/simulator/src/avm/journal/journal.test.ts index 0a2b8c4b1b3..ab7e21d0dd2 100644 --- a/yarn-project/simulator/src/avm/journal/journal.test.ts +++ b/yarn-project/simulator/src/avm/journal/journal.test.ts @@ -20,47 +20,6 @@ describe('journal', () => { }); describe('Public Storage', () => { - it('Should cache write to storage', () => { - // When writing to storage we should write to the storage writes map - const contractAddress = new Fr(1); - const key = new Fr(2); - const value = new Fr(3); - - journal.writeStorage(contractAddress, key, value); - - const journalUpdates: JournalData = journal.flush(); - expect(journalUpdates.currentStorageValue.get(contractAddress.toBigInt())?.get(key.toBigInt())).toEqual(value); - }); - - it('When reading from storage, should check the parent first', async () => { - // Store a different value in storage vs the cache, and make sure the cache is returned - const contractAddress = new Fr(1); - const key = new Fr(2); - const storedValue = new Fr(420); - const parentValue = new Fr(69); - const cachedValue = new Fr(1337); - - publicDb.storageRead.mockResolvedValue(Promise.resolve(storedValue)); - - const childJournal = new AvmWorldStateJournal(journal.hostStorage, journal); - - // Get the cache miss - const cacheMissResult = await childJournal.readStorage(contractAddress, key); - expect(cacheMissResult).toEqual(storedValue); - - // Write to storage - journal.writeStorage(contractAddress, key, parentValue); - const parentResult = await childJournal.readStorage(contractAddress, key); - expect(parentResult).toEqual(parentValue); - - // Get the parent value - childJournal.writeStorage(contractAddress, key, cachedValue); - - // Get the storage value - const cachedResult = await childJournal.readStorage(contractAddress, key); - expect(cachedResult).toEqual(cachedValue); - }); - it('When reading from storage, should check the cache first, and be appended to read/write journal', async () => { // Store a different value in storage vs the cache, and make sure the cache is returned const contractAddress = new Fr(1); diff --git a/yarn-project/simulator/src/avm/journal/journal.ts b/yarn-project/simulator/src/avm/journal/journal.ts index cdeb4b4ff92..3c52a334ece 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 { PublicStorage } from './public_storage.js'; /** * Data held within the journal @@ -32,6 +33,9 @@ export class AvmWorldStateJournal { /** Reference to node storage */ public readonly hostStorage: HostStorage; + /** World State's public storage, including cached writes */ + private publicStorage: PublicStorage; + // Reading state - must be tracked for vm execution // contract address -> key -> value[] (array stored in order of reads) private storageReads: Map> = new Map(); @@ -45,14 +49,9 @@ export class AvmWorldStateJournal { private newL1Messages: Fr[][] = []; private newLogs: Fr[][] = []; - // contract address -> key -> value - private currentStorageValue: Map> = new Map(); - - private parentJournal: AvmWorldStateJournal | undefined; - constructor(hostStorage: HostStorage, parentJournal?: AvmWorldStateJournal) { this.hostStorage = hostStorage; - this.parentJournal = parentJournal; + this.publicStorage = new PublicStorage(hostStorage.publicStateDb, parentJournal?.publicStorage); } /** @@ -63,47 +62,29 @@ export class AvmWorldStateJournal { } /** - * Write storage into journal + * Write to public storage, journal/trace the write. * - * @param contractAddress - - * @param key - - * @param value - + * @param storageAddress - the address of the contract whose storage is being written to + * @param slot - the slot in the contract's storage being written to + * @param value - the value being written to the slot */ - public writeStorage(contractAddress: Fr, key: Fr, value: Fr) { - let contractMap = this.currentStorageValue.get(contractAddress.toBigInt()); - if (!contractMap) { - contractMap = new Map(); - this.currentStorageValue.set(contractAddress.toBigInt(), contractMap); - } - contractMap.set(key.toBigInt(), value); - + public writeStorage(storageAddress: Fr, slot: Fr, value: Fr) { + this.publicStorage.write(storageAddress, slot, value); // We want to keep track of all performed writes in the journal - this.journalWrite(contractAddress, key, value); + this.journalWrite(storageAddress, slot, value); } /** - * Read storage from journal - * Read from host storage on cache miss + * Read from public storage, journal/trace the read. * - * @param contractAddress - - * @param key - - * @returns current value + * @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 the latest value written to slot, or 0 if never written to before */ - public async readStorage(contractAddress: Fr, key: Fr): Promise { - // - We first try this journal's storage cache ( if written to before in this call frame ) - // - Then we try the parent journal's storage cache ( if it exists ) ( written to earlier in this block ) - // - Finally we try the host storage ( a trip to the database ) - - // Do not early return as we want to keep track of reads in this.storageReads - let value = this.currentStorageValue.get(contractAddress.toBigInt())?.get(key.toBigInt()); - if (!value && this.parentJournal) { - value = await this.parentJournal?.readStorage(contractAddress, key); - } - if (!value) { - value = await this.hostStorage.publicStateDb.storageRead(contractAddress, key); - } - - this.journalRead(contractAddress, key, value); + 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 in the journal + this.journalRead(storageAddress, slot, value); return Promise.resolve(value); } @@ -158,15 +139,15 @@ export class AvmWorldStateJournal { * - Public state journals (r/w logs), with the accessing being appended in chronological order */ public acceptNestedWorldState(nestedJournal: AvmWorldStateJournal) { + // Merge Public Storage + this.publicStorage.acceptAndMerge(nestedJournal.publicStorage); + // Merge UTXOs this.newNoteHashes = this.newNoteHashes.concat(nestedJournal.newNoteHashes); this.newL1Messages = this.newL1Messages.concat(nestedJournal.newL1Messages); this.newNullifiers = this.newNullifiers.concat(nestedJournal.newNullifiers); this.newLogs = this.newLogs.concat(nestedJournal.newLogs); - // Merge Public State - mergeCurrentValueMaps(this.currentStorageValue, nestedJournal.currentStorageValue); - // Merge storage read and write journals mergeContractJournalMaps(this.storageReads, nestedJournal.storageReads); mergeContractJournalMaps(this.storageWrites, nestedJournal.storageWrites); @@ -195,43 +176,13 @@ export class AvmWorldStateJournal { newNullifiers: this.newNullifiers, newL1Messages: this.newL1Messages, newLogs: this.newLogs, - currentStorageValue: this.currentStorageValue, + currentStorageValue: this.publicStorage.getCache().cachePerContract, storageReads: this.storageReads, storageWrites: this.storageWrites, }; } } -/** - * Merges two contract current value together - * Where childMap keys will take precedent over the hostMap - * The assumption being that the child map is created at a later time - * And thus contains more up to date information - * - * @param hostMap - The map to be merged into - * @param childMap - The map to be merged from - */ -function mergeCurrentValueMaps(hostMap: Map>, childMap: Map>) { - for (const [key, value] of childMap) { - const map1Value = hostMap.get(key); - if (!map1Value) { - hostMap.set(key, value); - } else { - mergeStorageCurrentValueMaps(map1Value, value); - } - } -} - -/** - * @param hostMap - The map to be merge into - * @param childMap - The map to be merged from - */ -function mergeStorageCurrentValueMaps(hostMap: Map, childMap: Map) { - for (const [key, value] of childMap) { - hostMap.set(key, value); - } -} - /** * Merges two contract journalling maps together * For read maps, we just append the childMap arrays into the host map arrays, as the order is important diff --git a/yarn-project/simulator/src/avm/journal/public_storage.test.ts b/yarn-project/simulator/src/avm/journal/public_storage.test.ts new file mode 100644 index 00000000000..d8b79fc79c0 --- /dev/null +++ b/yarn-project/simulator/src/avm/journal/public_storage.test.ts @@ -0,0 +1,116 @@ +import { Fr } from '@aztec/foundation/fields'; + +import { MockProxy, mock } from 'jest-mock-extended'; + +import { PublicStateDB } from '../../index.js'; +import { PublicStorage } from './public_storage.js'; + +describe('avm public storage', () => { + let publicDb: MockProxy; + let publicStorage: PublicStorage; + + beforeEach(() => { + publicDb = mock(); + publicStorage = new PublicStorage(publicDb); + }); + + describe('AVM Public Storage', () => { + it('Reading an unwritten slot works (gets zero & DNE)', async () => { + const contractAddress = new Fr(1); + const slot = new Fr(2); + // never written! + const [exists, gotValue] = await publicStorage.read(contractAddress, slot); + // doesn't exist, value is zero + expect(exists).toEqual(false); + expect(gotValue).toEqual(Fr.ZERO); + }); + 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); + // exists because it was previously written + expect(exists).toEqual(true); + expect(gotValue).toEqual(value); + }); + 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); + // it exists in the host, so it must've been written before + expect(exists).toEqual(true); + expect(gotValue).toEqual(storedValue); + }); + 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); + // exists because it was previously written! + expect(exists).toEqual(true); + expect(gotValue).toEqual(value); + }); + 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); + const slot = new Fr(2); + const storedValue = new Fr(420); + const parentValue = new Fr(69); + const cachedValue = new Fr(1337); + + publicDb.storageRead.mockResolvedValue(Promise.resolve(storedValue)); + const childStorage = new PublicStorage(publicDb, publicStorage); + + // Cache miss falls back to host + const [, cacheMissResult] = await childStorage.read(contractAddress, slot); + expect(cacheMissResult).toEqual(storedValue); + + // Write to storage + publicStorage.write(contractAddress, slot, parentValue); + // Reading from child should give value written in parent + const [, valueFromParent] = await childStorage.read(contractAddress, slot); + expect(valueFromParent).toEqual(parentValue); + + // 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); + expect(cachedResult).toEqual(cachedValue); + }); + }); + + it('Should be able to merge two public storages together', async () => { + // Checking that child's writes take precedence on marge + const contractAddress = new Fr(1); + const slot = new Fr(2); + // value written initially in parent + const value = new Fr(1); + // value overwritten in child and later merged into parent + const valueT1 = new Fr(2); + + // Write initial value to parent + publicStorage.write(contractAddress, slot, value); + + const childStorage = new PublicStorage(publicDb, publicStorage); + // Write valueT1 to child + childStorage.write(contractAddress, slot, valueT1); + + // Parent accepts child's staged writes + publicStorage.acceptAndMerge(childStorage); + + // Read from parent gives latest value written in child before merge (valueT1) + const [exists, result] = await publicStorage.read(contractAddress, slot); + expect(exists).toEqual(true); + expect(result).toEqual(valueT1); + }); +}); diff --git a/yarn-project/simulator/src/avm/journal/public_storage.ts b/yarn-project/simulator/src/avm/journal/public_storage.ts new file mode 100644 index 00000000000..f6b3b7a35d8 --- /dev/null +++ b/yarn-project/simulator/src/avm/journal/public_storage.ts @@ -0,0 +1,149 @@ +import { Fr } from '@aztec/foundation/fields'; + +import type { PublicStateDB } from '../../index.js'; + +/** + * 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. + * When a contract call completes, its storage cache can be merged into its parent's. + */ +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; + this.cache = new PublicStorageCache(); + } + + /** + * Get the pending storage. + */ + public getCache() { + return this.cache; + } + + /** + * Read a value from storage. + * 1. Check cache. + * 2. Check parent's 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. + * + * @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 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]> { + // 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); + } + // Finally try the host's Aztec state (a trip to the database) + if (!value) { + value = await this.hostPublicStorage.storageRead(storageAddress, slot); + } + // 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]); + } + + /** + * Stage a storage write. + * + * @param storageAddress - the address of the contract whose storage is being written to + * @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); + } + + /** + * Merges another PublicStorage's cache (pending writes) into this one. + * + * @param incomingPublicStorage - the incoming public storage to merge into this instance's + */ + public acceptAndMerge(incomingPublicStorage: PublicStorage) { + this.cache.acceptAndMerge(incomingPublicStorage.cache); + } +} + +/** + * A class to cache writes to public storage during a contract call's AVM simulation. + * "Writes" update a map, "reads" check that map or return undefined. + * An instance of this class can merge another instance's staged writes into its own. + */ +class PublicStorageCache { + /** + * Map for staging storage writes. + * One inner-map per contract storage address, + * mapping storage slot to latest staged write value. + */ + public cachePerContract: Map> = new Map(); + // FIXME: storage ^ should be private, but its value is used in tests for "currentStorageValue" + + /** + * Read a staged value from storage, if it has been previously written to. + * + * @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 the latest value written to slot, or undefined if no value has been written + */ + public read(storageAddress: Fr, slot: Fr): Fr | undefined { + return this.cachePerContract.get(storageAddress.toBigInt())?.get(slot.toBigInt()); + } + + /** + * Stage a storage write. + * + * @param storageAddress - the address of the contract whose storage is being written to + * @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, slot: Fr, value: Fr) { + let cacheAtContract = this.cachePerContract.get(storageAddress.toBigInt()); + if (!cacheAtContract) { + // If this contract's storage has no staged modifications, create a new inner map to store them + cacheAtContract = new Map(); + this.cachePerContract.set(storageAddress.toBigInt(), cacheAtContract); + } + cacheAtContract.set(slot.toBigInt(), value); + } + + /** + * Merges another cache's staged writes into this instance's cache. + * + * Staged modifications in "incoming" take precedence over those + * present in "this" as they are assumed to occur after this' writes. + * + * In practice, "this" is a parent call's storage cache, and "incoming" is a nested call's. + * + * @param incomingStorageCache - the incoming storage write cache to merge into this instance's + */ + public acceptAndMerge(incomingStorageCache: PublicStorageCache) { + // Iterate over all incoming contracts with staged writes. + for (const [incomingAddress, incomingCacheAtContract] of incomingStorageCache.cachePerContract) { + const thisCacheAtContract = this.cachePerContract.get(incomingAddress); + if (!thisCacheAtContract) { + // The contract has no storage writes staged here + // so just accept the incoming cache as-is for this contract. + this.cachePerContract.set(incomingAddress, incomingCacheAtContract); + } else { + // "Incoming" and "this" both have staged writes for this contract. + // Merge in incoming staged writes, giving them precedence over this'. + for (const [slot, value] of incomingCacheAtContract) { + thisCacheAtContract.set(slot, value); + } + } + } + } +}