From 56584683340cd29b26adbcc60d3cd58fe889e8ad Mon Sep 17 00:00:00 2001 From: Maddiaa <47148561+Maddiaa0@users.noreply.github.com> Date: Thu, 18 Jan 2024 13:04:19 +0000 Subject: [PATCH] feat(public-vm): avm journal (#3945) --- .../src/avm/avm_state_manager.ts | 5 +- .../acir-simulator/src/avm/journal/errors.ts | 8 + .../src/avm/journal/host_storage.ts | 13 +- .../src/avm/journal/journal.test.ts | 157 +++++++++++++++- .../acir-simulator/src/avm/journal/journal.ts | 168 ++++++++++++++---- 5 files changed, 302 insertions(+), 49 deletions(-) create mode 100644 yarn-project/acir-simulator/src/avm/journal/errors.ts diff --git a/yarn-project/acir-simulator/src/avm/avm_state_manager.ts b/yarn-project/acir-simulator/src/avm/avm_state_manager.ts index 250a7d8f122..b5396accc7a 100644 --- a/yarn-project/acir-simulator/src/avm/avm_state_manager.ts +++ b/yarn-project/acir-simulator/src/avm/avm_state_manager.ts @@ -29,7 +29,7 @@ export class AvmStateManager { * @returns Avm State Manager */ public static rootStateManager(blockHeader: BlockHeader, hostStorage: HostStorage): AvmStateManager { - const journal = new AvmJournal(hostStorage); + const journal = AvmJournal.rootJournal(hostStorage); return new AvmStateManager(blockHeader, journal); } @@ -39,6 +39,7 @@ export class AvmStateManager { * @returns */ public static forkStateManager(parent: AvmStateManager): AvmStateManager { - return new AvmStateManager(parent.blockHeader, parent.journal); + const journal = AvmJournal.branchParent(parent.journal); + return new AvmStateManager(parent.blockHeader, journal); } } diff --git a/yarn-project/acir-simulator/src/avm/journal/errors.ts b/yarn-project/acir-simulator/src/avm/journal/errors.ts new file mode 100644 index 00000000000..17696b1de68 --- /dev/null +++ b/yarn-project/acir-simulator/src/avm/journal/errors.ts @@ -0,0 +1,8 @@ +/** + * Error thrown when a base journal is attempted to be merged. + */ +export class RootJournalCannotBeMerged extends Error { + constructor() { + super('Root journal cannot be merged'); + } +} diff --git a/yarn-project/acir-simulator/src/avm/journal/host_storage.ts b/yarn-project/acir-simulator/src/avm/journal/host_storage.ts index ccd31891205..fcd9d27f722 100644 --- a/yarn-project/acir-simulator/src/avm/journal/host_storage.ts +++ b/yarn-project/acir-simulator/src/avm/journal/host_storage.ts @@ -1,17 +1,20 @@ import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../index.js'; -/** - */ +/** + * Host storage + * + * A wrapper around the node dbs + */ export class HostStorage { /** - */ - public readonly stateDb: PublicStateDB; + public readonly publicStateDb: PublicStateDB; /** - */ public readonly contractsDb: PublicContractsDB; - /** - */ public readonly commitmentsDb: CommitmentsDB; - constructor(stateDb: PublicStateDB, contractsDb: PublicContractsDB, commitmentsDb: CommitmentsDB) { - this.stateDb = stateDb; + constructor(publicStateDb: PublicStateDB, contractsDb: PublicContractsDB, commitmentsDb: CommitmentsDB) { + this.publicStateDb = publicStateDb; this.contractsDb = contractsDb; this.commitmentsDb = commitmentsDb; } diff --git a/yarn-project/acir-simulator/src/avm/journal/journal.test.ts b/yarn-project/acir-simulator/src/avm/journal/journal.test.ts index c356db72386..a11a3b38743 100644 --- a/yarn-project/acir-simulator/src/avm/journal/journal.test.ts +++ b/yarn-project/acir-simulator/src/avm/journal/journal.test.ts @@ -1,7 +1,158 @@ +import { Fr } from '@aztec/foundation/fields'; + +import { MockProxy, mock } from 'jest-mock-extended'; + +import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../index.js'; +import { HostStorage } from './host_storage.js'; +import { AvmJournal, JournalData } from './journal.js'; + describe('journal', () => { - it('Should write to storage', () => {}); + let publicDb: MockProxy; + let journal: AvmJournal; + + beforeEach(() => { + publicDb = mock(); + const commitmentsDb = mock(); + const contractsDb = mock(); + + const hostStorage = new HostStorage(publicDb, contractsDb, commitmentsDb); + journal = new AvmJournal(hostStorage); + }); + + 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.storageWrites.get(contractAddress)?.get(key)).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 AvmJournal(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', 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 cachedValue = new Fr(69); + + publicDb.storageRead.mockResolvedValue(Promise.resolve(storedValue)); + + // Get the cache first + const cacheMissResult = await journal.readStorage(contractAddress, key); + expect(cacheMissResult).toEqual(storedValue); + + // Write to storage + journal.writeStorage(contractAddress, key, cachedValue); + + // Get the storage value + const cachedResult = await journal.readStorage(contractAddress, key); + expect(cachedResult).toEqual(cachedValue); + }); + }); + + describe('UTXOs', () => { + it('Should maintain commitments', () => { + const utxo = new Fr(1); + journal.writeCommitment(utxo); + + const journalUpdates = journal.flush(); + expect(journalUpdates.newCommitments).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 journals together', async () => { + // Fundamentally checking that insert ordering of public storage is preserved upon journal merge + // time | journal | op | value + // t0 -> journal0 -> write | 1 + // t1 -> journal1 -> write | 2 + // merge journals + // t2 -> journal0 -> read | 2 + + const contractAddress = new Fr(1); + const key = new Fr(2); + const value = new Fr(1); + const valueT1 = new Fr(2); + const commitment = new Fr(10); + const commitmentT1 = new Fr(20); + + journal.writeStorage(contractAddress, key, value); + journal.writeCommitment(commitment); + journal.writeL1Message(commitment); + journal.writeNullifier(commitment); + + const journal1 = new AvmJournal(journal.hostStorage, journal); + journal.writeStorage(contractAddress, key, valueT1); + journal.writeCommitment(commitmentT1); + journal.writeL1Message(commitmentT1); + journal.writeNullifier(commitmentT1); + + journal1.mergeWithParent(); + + // Check that the storage is merged by reading from the journal + const result = await journal.readStorage(contractAddress, key); + expect(result).toEqual(valueT1); + + // Check that the UTXOs are merged + const journalUpdates: JournalData = journal.flush(); + expect(journalUpdates.newCommitments).toEqual([commitment, commitmentT1]); + expect(journalUpdates.newL1Messages).toEqual([commitment, commitmentT1]); + expect(journalUpdates.newNullifiers).toEqual([commitment, commitmentT1]); + }); - it('Should read from storage', () => {}); + it('Cannot merge a root journal, but can merge a child journal', () => { + const rootJournal = AvmJournal.rootJournal(journal.hostStorage); + const childJournal = AvmJournal.branchParent(rootJournal); - it('Should merge two journals together', () => {}); + expect(() => rootJournal.mergeWithParent()).toThrow(); + expect(() => childJournal.mergeWithParent()); + }); }); diff --git a/yarn-project/acir-simulator/src/avm/journal/journal.ts b/yarn-project/acir-simulator/src/avm/journal/journal.ts index 67563fe6370..cb69b3c33ef 100644 --- a/yarn-project/acir-simulator/src/avm/journal/journal.ts +++ b/yarn-project/acir-simulator/src/avm/journal/journal.ts @@ -1,101 +1,191 @@ import { Fr } from '@aztec/foundation/fields'; +import { RootJournalCannotBeMerged } from './errors.js'; import { HostStorage } from './host_storage.js'; -// TODO: all of the data that comes out of the avm ready for write should be in this format -/** - */ +/** + * Data held within the journal + */ export type JournalData = { /** - */ newCommitments: Fr[]; /** - */ - newL1Message: Fr[]; + newL1Messages: Fr[]; /** - */ - storageWrites: { [key: string]: { [key: string]: Fr } }; + newNullifiers: Fr[]; + /** contract address -\> key -\> value */ + storageWrites: Map>; }; -// This persists for an entire block -// Each transaction should have its own journal that gets appended to this one upon success -/** - */ +/** + * A cache of the current state of the AVM + * The interpreter should make any state queries through this object + * + * When a sub context's call succeeds, it's journal is merge into the parent + * When a a call fails, it's journal is discarded and the parent is used from this point forward + * When a call succeeds's we can merge a child into its parent + */ export class AvmJournal { - // TODO: should we make private? - /** - */ + /** Reference to node storage */ public readonly hostStorage: HostStorage; - // We need to keep track of the following - // - State reads - // - State updates - // - New Commitments - // - Commitment reads + // Reading state - must be tracked for vm execution + // contract address -> key -> value + // TODO(https://github.com/AztecProtocol/aztec-packages/issues/3999) + private storageReads: Map> = new Map(); + // New written state private newCommitments: Fr[] = []; + private newNullifiers: Fr[] = []; private newL1Message: Fr[] = []; - // TODO: type this structure -> contract address -> key -> value - private storageWrites: { [key: string]: { [key: string]: Fr } } = {}; + // New Substrate + private newLogs: Fr[][] = []; + + // contract address -> key -> value + private storageWrites: Map> = new Map(); - constructor(hostStorage: HostStorage) { + private parentJournal: AvmJournal | undefined; + + constructor(hostStorage: HostStorage, parentJournal?: AvmJournal) { this.hostStorage = hostStorage; + this.parentJournal = parentJournal; + } + + /** + * Create a new root journal, without a parent + * @param hostStorage - + */ + public static rootJournal(hostStorage: HostStorage) { + return new AvmJournal(hostStorage); } - // TODO: work on the typing /** - * - + * Create a new journal from a parent + * @param parentJournal - + */ + public static branchParent(parentJournal: AvmJournal) { + return new AvmJournal(parentJournal.hostStorage, parentJournal); + } + + /** + * Write storage into journal + * * @param contractAddress - * @param key - * @param value - */ public writeStorage(contractAddress: Fr, key: Fr, value: Fr) { - // TODO: do we want this map to be ordered -> is there performance upside to this? - this.storageWrites[contractAddress.toString()][key.toString()] = value; + let contractMap = this.storageWrites.get(contractAddress); + if (!contractMap) { + contractMap = new Map(); + this.storageWrites.set(contractAddress, contractMap); + } + contractMap.set(key, value); } /** - * - + * Read storage from journal + * Read from host storage on cache miss + * * @param contractAddress - * @param key - + * @returns current value */ - public readStorage(contractAddress: Fr, key: Fr) { - const cachedValue = this.storageWrites[contractAddress.toString()][key.toString()]; + public readStorage(contractAddress: Fr, key: Fr): Promise { + const cachedValue = this.storageWrites.get(contractAddress)?.get(key); if (cachedValue) { - return cachedValue; + return Promise.resolve(cachedValue); } - return this.hostStorage.stateDb.storageRead(contractAddress, key); + if (this.parentJournal) { + return this.parentJournal?.readStorage(contractAddress, key); + } + return this.hostStorage.publicStateDb.storageRead(contractAddress, key); } - /** - * - + /** - * @param commitment - */ public writeCommitment(commitment: Fr) { this.newCommitments.push(commitment); } - /** - * - + /** - * @param message - */ public writeL1Message(message: Fr) { this.newL1Message.push(message); } - // TODO: This function will merge two journals together -> the new head of the chain - /** - * - - * @param journal - + /** - + * @param nullifier - */ - public mergeJournal(journal: AvmJournal) { - // TODO: This function will - void journal; + public writeNullifier(nullifier: Fr) { + this.newNullifiers.push(nullifier); } /** - * - + * Merge Journal into parent + * - Utxo objects are concatenated + * - Public state is merged, with the value in the incoming journal taking precedent + */ + public mergeWithParent() { + if (!this.parentJournal) { + throw new RootJournalCannotBeMerged(); + } + + const incomingFlush = this.flush(); + + // Merge UTXOs + this.parentJournal.newCommitments = this.parentJournal.newCommitments.concat(incomingFlush.newCommitments); + this.parentJournal.newL1Message = this.parentJournal.newL1Message.concat(incomingFlush.newL1Messages); + this.parentJournal.newNullifiers = this.parentJournal.newNullifiers.concat(incomingFlush.newNullifiers); + + // Merge Public State + mergeContractMaps(this.parentJournal.storageWrites, incomingFlush.storageWrites); + } + + /** Access the current state of the journal + * + * @returns a JournalData object that can be used to write to the storage */ public flush(): JournalData { return { newCommitments: this.newCommitments, - newL1Message: this.newL1Message, + newL1Messages: this.newL1Message, + newNullifiers: this.newNullifiers, storageWrites: this.storageWrites, }; } } + +/** + * Merges two contract maps 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 mergeContractMaps(hostMap: Map>, childMap: Map>) { + for (const [key, value] of childMap) { + const map1Value = hostMap.get(key); + if (!map1Value) { + hostMap.set(key, value); + } else { + mergeStorageMaps(map1Value, value); + } + } +} + +/** + * + * @param hostMap - The map to be merge into + * @param childMap - The map to be merged from + */ +function mergeStorageMaps(hostMap: Map, childMap: Map) { + for (const [key, value] of childMap) { + hostMap.set(key, value); + } +}