feat(avm): add revert tracking to the journal (AztecProtocol#4349)

fixes: AztecProtocol#3998

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();
   }
 }
 

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,23 +152,23 @@ 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();
 
     // 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]);
+    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());
   });
 });

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<bigint, Fr[]>, childMap: Map<bigin
     if (!readArr) {
       hostMap.set(key, value);
     } else {
-      readArr?.concat(value);
+      hostMap.set(key, readArr?.concat(...value));
     }
   }
 }

yarn-project/acir-simulator/src/avm/opcodes/external_calls.ts

@@ -60,9 +60,10 @@ export class Call extends Instruction {
     machineState.memory.setSlice(this.retOffset, convertedReturnData);
 
     if (success) {
-      avmContext.mergeJournal();
+      avmContext.mergeJournalSuccess();
+    } else {
+      avmContext.mergeJournalFailure();
     }
-
     this.incrementPc(machineState);
   }
 }
@@ -119,7 +120,9 @@ export class StaticCall extends Instruction {
     machineState.memory.setSlice(this.retOffset, convertedReturnData);
 
     if (success) {
-      avmContext.mergeJournal();
+      avmContext.mergeJournalSuccess();
+    } else {
+      avmContext.mergeJournalFailure();
     }
 
     this.incrementPc(machineState);