From 161580312a753fb2f0507c11fb10d895b5073e3e Mon Sep 17 00:00:00 2001 From: Maddiaa <47148561+Maddiaa0@users.noreply.github.com> Date: Thu, 1 Feb 2024 21:30:37 +0000 Subject: [PATCH] feat(avm): add revert tracking to the journal (#4349) fixes: https://github.com/AztecProtocol/aztec-packages/issues/3998 --- .../acir-simulator/src/avm/avm_context.ts | 12 ++- .../src/avm/journal/journal.test.ts | 79 +++++++++++++++++-- .../acir-simulator/src/avm/journal/journal.ts | 25 +++++- .../src/avm/opcodes/external_calls.ts | 9 ++- 4 files changed, 109 insertions(+), 16 deletions(-) diff --git a/yarn-project/acir-simulator/src/avm/avm_context.ts b/yarn-project/acir-simulator/src/avm/avm_context.ts index 9816fda32d8..b6226c818e9 100644 --- a/yarn-project/acir-simulator/src/avm/avm_context.ts +++ b/yarn-project/acir-simulator/src/avm/avm_context.ts @@ -119,8 +119,16 @@ export class AvmContext { * Merge the journal of this call with it's parent * NOTE: this should never be called on a root context - only from within a nested call */ - public mergeJournal() { - this.journal.mergeWithParent(); + public mergeJournalSuccess() { + this.journal.mergeSuccessWithParent(); + } + + /** + * Merge the journal of this call with it's parent + * For when the child call fails ( we still must track state accesses ) + */ + public mergeJournalFailure() { + this.journal.mergeFailureWithParent(); } } 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 1588712e1c7..2a3aafd6da6 100644 --- a/yarn-project/acir-simulator/src/avm/journal/journal.test.ts +++ b/yarn-project/acir-simulator/src/avm/journal/journal.test.ts @@ -3,6 +3,7 @@ import { Fr } from '@aztec/foundation/fields'; import { MockProxy, mock } from 'jest-mock-extended'; import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../index.js'; +import { RootJournalCannotBeMerged } from './errors.js'; import { HostStorage } from './host_storage.js'; import { AvmJournal, JournalData } from './journal.js'; @@ -119,7 +120,7 @@ describe('journal', () => { }); }); - it('Should merge two journals together', async () => { + it('Should merge two successful journals together', async () => { // Fundamentally checking that insert ordering of public storage is preserved upon journal merge // time | journal | op | value // t0 -> journal0 -> write | 1 @@ -151,12 +152,12 @@ describe('journal', () => { journal1.writeL1Message(logsT1); journal1.writeNullifier(commitmentT1); - journal1.mergeWithParent(); + journal1.mergeSuccessWithParent(); - // 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 storage is merged by reading from the journal // Check that the UTXOs are merged const journalUpdates: JournalData = journal.flush(); @@ -164,10 +165,10 @@ describe('journal', () => { // We first read value from t0, then value from t1 const contractReads = journalUpdates.storageReads.get(contractAddress.toBigInt()); const slotReads = contractReads?.get(key.toBigInt()); - expect(slotReads).toEqual([value, valueT1]); + expect(slotReads).toEqual([value, valueT1, valueT1]); // Read a third time to check storage // We first write value from t0, then value from t1 - const contractWrites = journalUpdates.storageReads.get(contractAddress.toBigInt()); + const contractWrites = journalUpdates.storageWrites.get(contractAddress.toBigInt()); const slotWrites = contractWrites?.get(key.toBigInt()); expect(slotWrites).toEqual([value, valueT1]); @@ -177,11 +178,75 @@ describe('journal', () => { expect(journalUpdates.newNullifiers).toEqual([commitment, commitmentT1]); }); + it('Should merge failed journals together', async () => { + // Checking public storage update journals are preserved upon journal merge, + // But the latest state is not + + // time | journal | op | value + // t0 -> journal0 -> write | 1 + // t1 -> journal1 -> write | 2 + // merge journals + // t2 -> journal0 -> read | 1 + + 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); + const logs = [new Fr(1), new Fr(2)]; + const logsT1 = [new Fr(3), new Fr(4)]; + + journal.writeStorage(contractAddress, key, value); + await journal.readStorage(contractAddress, key); + journal.writeNoteHash(commitment); + journal.writeLog(logs); + journal.writeL1Message(logs); + journal.writeNullifier(commitment); + + const journal1 = new AvmJournal(journal.hostStorage, journal); + journal1.writeStorage(contractAddress, key, valueT1); + await journal1.readStorage(contractAddress, key); + journal1.writeNoteHash(commitmentT1); + journal1.writeLog(logsT1); + journal1.writeL1Message(logsT1); + journal1.writeNullifier(commitmentT1); + + journal1.mergeFailureWithParent(); + + // Check that the storage is reverted by reading from the journal + const result = await journal.readStorage(contractAddress, key); + expect(result).toEqual(value); // rather than valueT1 + + // Check that the UTXOs are merged + const journalUpdates: JournalData = journal.flush(); + + // Reads and writes should be preserved + // Check storage reads order is preserved upon merge + // We first read value from t0, then value from t1 + const contractReads = journalUpdates.storageReads.get(contractAddress.toBigInt()); + const slotReads = contractReads?.get(key.toBigInt()); + expect(slotReads).toEqual([value, valueT1, value]); // Read a third time to check storage above + + // We first write value from t0, then value from t1 + const contractWrites = journalUpdates.storageWrites.get(contractAddress.toBigInt()); + const slotWrites = contractWrites?.get(key.toBigInt()); + expect(slotWrites).toEqual([value, valueT1]); + + expect(journalUpdates.newNoteHashes).toEqual([commitment]); + expect(journalUpdates.newLogs).toEqual([logs]); + expect(journalUpdates.newL1Messages).toEqual([logs]); + expect(journalUpdates.newNullifiers).toEqual([commitment]); + }); + it('Cannot merge a root journal, but can merge a child journal', () => { const rootJournal = AvmJournal.rootJournal(journal.hostStorage); const childJournal = AvmJournal.branchParent(rootJournal); - expect(() => rootJournal.mergeWithParent()).toThrow(); - expect(() => childJournal.mergeWithParent()); + expect(() => rootJournal.mergeSuccessWithParent()).toThrow(RootJournalCannotBeMerged); + expect(() => rootJournal.mergeFailureWithParent()).toThrow(RootJournalCannotBeMerged); + + expect(() => childJournal.mergeSuccessWithParent()); + expect(() => childJournal.mergeFailureWithParent()); }); }); diff --git a/yarn-project/acir-simulator/src/avm/journal/journal.ts b/yarn-project/acir-simulator/src/avm/journal/journal.ts index 699d6b93f00..3b2445e7dc6 100644 --- a/yarn-project/acir-simulator/src/avm/journal/journal.ts +++ b/yarn-project/acir-simulator/src/avm/journal/journal.ts @@ -162,11 +162,12 @@ export class AvmJournal { } /** - * Merge Journal into parent + * Merge Journal from successful call into parent * - Utxo objects are concatenated - * - Public state is merged, with the value in the incoming journal taking precedent + * - Public state changes are merged, with the value in the incoming journal taking precedent + * - Public state journals (r/w logs), with the accessing being appended in chronological order */ - public mergeWithParent() { + public mergeSuccessWithParent() { if (!this.parentJournal) { throw new RootJournalCannotBeMerged(); } @@ -185,6 +186,22 @@ export class AvmJournal { mergeContractJournalMaps(this.parentJournal.storageWrites, this.storageWrites); } + /** + * Merge Journal for failed call into parent + * - Utxo objects are concatenated + * - Public state changes are dropped + * - Public state journals (r/w logs) are maintained, with the accessing being appended in chronological order + */ + public mergeFailureWithParent() { + if (!this.parentJournal) { + throw new RootJournalCannotBeMerged(); + } + + // Merge storage read and write journals + mergeContractJournalMaps(this.parentJournal.storageReads, this.storageReads); + mergeContractJournalMaps(this.parentJournal.storageWrites, this.storageWrites); + } + /** * Access the current state of the journal * @@ -261,7 +278,7 @@ function mergeStorageJournalMaps(hostMap: Map, childMap: Map