chore(avm-simulator): create a dedicated component just for tracing w…

…orld state accesses (#4733)

yarn-project/simulator/src/avm/avm_context.test.ts

@@ -26,7 +26,7 @@ describe('Avm Context', () => {
     );
 
     // We stringify to remove circular references (parentJournal)
-    expect(JSON.stringify(newContext.worldState)).toEqual(JSON.stringify(context.worldState.fork()));
+    expect(JSON.stringify(newContext.persistableState)).toEqual(JSON.stringify(context.persistableState.fork()));
   });
 
   it('New static call should fork context correctly', () => {
@@ -52,6 +52,6 @@ describe('Avm Context', () => {
     );
 
     // We stringify to remove circular references (parentJournal)
-    expect(JSON.stringify(newContext.worldState)).toEqual(JSON.stringify(context.worldState.fork()));
+    expect(JSON.stringify(newContext.persistableState)).toEqual(JSON.stringify(context.persistableState.fork()));
   });
 });

yarn-project/simulator/src/avm/avm_context.ts

@@ -3,7 +3,7 @@ import { Fr } from '@aztec/foundation/fields';
 
 import { AvmExecutionEnvironment } from './avm_execution_environment.js';
 import { AvmMachineState } from './avm_machine_state.js';
-import { AvmWorldStateJournal } from './journal/journal.js';
+import { AvmPersistableStateManager } from './journal/journal.js';
 
 /**
  * An execution context includes the information necessary to initiate AVM
@@ -12,13 +12,13 @@ import { AvmWorldStateJournal } from './journal/journal.js';
 export class AvmContext {
   /**
    * Create a new AVM context
-   * @param worldState - Manages mutable state during execution - (caching, fetching)
+   * @param persistableState - Manages world state and accrued substate during execution - (caching, fetching, tracing)
    * @param environment - Contains constant variables provided by the kernel
    * @param machineState - VM state that is modified on an instruction-by-instruction basis
    * @returns new AvmContext instance
    */
   constructor(
-    public worldState: AvmWorldStateJournal,
+    public persistableState: AvmPersistableStateManager,
     public environment: AvmExecutionEnvironment,
     public machineState: AvmMachineState,
   ) {}
@@ -37,7 +37,7 @@ export class AvmContext {
    */
   public createNestedContractCallContext(address: AztecAddress, calldata: Fr[]): AvmContext {
     const newExecutionEnvironment = this.environment.deriveEnvironmentForNestedCall(address, calldata);
-    const forkedWorldState = this.worldState.fork();
+    const forkedWorldState = this.persistableState.fork();
     const machineState = AvmMachineState.fromState(this.machineState);
     return new AvmContext(forkedWorldState, newExecutionEnvironment, machineState);
   }
@@ -56,7 +56,7 @@ export class AvmContext {
    */
   public createNestedContractStaticCallContext(address: AztecAddress, calldata: Fr[]): AvmContext {
     const newExecutionEnvironment = this.environment.deriveEnvironmentForNestedStaticCall(address, calldata);
-    const forkedWorldState = this.worldState.fork();
+    const forkedWorldState = this.persistableState.fork();
     const machineState = AvmMachineState.fromState(this.machineState);
     return new AvmContext(forkedWorldState, newExecutionEnvironment, machineState);
   }

yarn-project/simulator/src/avm/avm_simulator.test.ts

@@ -24,7 +24,9 @@ describe('AVM simulator', () => {
     ]);
 
     const context = initContext({ env: initExecutionEnvironment({ calldata }) });
-    jest.spyOn(context.worldState.hostStorage.contractsDb, 'getBytecode').mockReturnValue(Promise.resolve(bytecode));
+    jest
+      .spyOn(context.persistableState.hostStorage.contractsDb, 'getBytecode')
+      .mockReturnValue(Promise.resolve(bytecode));
 
     const results = await new AvmSimulator(context).execute();
 
@@ -46,7 +48,9 @@ describe('AVM simulator', () => {
       const bytecode = Buffer.from(addArtifact.bytecode, 'base64');
 
       const context = initContext({ env: initExecutionEnvironment({ calldata }) });
-      jest.spyOn(context.worldState.hostStorage.contractsDb, 'getBytecode').mockReturnValue(Promise.resolve(bytecode));
+      jest
+        .spyOn(context.persistableState.hostStorage.contractsDb, 'getBytecode')
+        .mockReturnValue(Promise.resolve(bytecode));
 
       const results = await new AvmSimulator(context).execute();
 
@@ -71,7 +75,7 @@ describe('AVM simulator', () => {
 
         const context = initContext();
         jest
-          .spyOn(context.worldState.hostStorage.contractsDb, 'getBytecode')
+          .spyOn(context.persistableState.hostStorage.contractsDb, 'getBytecode')
           .mockReturnValue(Promise.resolve(bytecode));
 
         const results = await new AvmSimulator(context).execute();
@@ -99,7 +103,7 @@ describe('AVM simulator', () => {
 
         const context = initContext({ env: initExecutionEnvironment({ calldata }) });
         jest
-          .spyOn(context.worldState.hostStorage.contractsDb, 'getBytecode')
+          .spyOn(context.persistableState.hostStorage.contractsDb, 'getBytecode')
           .mockReturnValue(Promise.resolve(bytecode));
 
         const results = await new AvmSimulator(context).execute();
@@ -131,7 +135,7 @@ describe('AVM simulator', () => {
 
         const context = initContext({ env: initExecutionEnvironment({ calldata }) });
         jest
-          .spyOn(context.worldState.hostStorage.contractsDb, 'getBytecode')
+          .spyOn(context.persistableState.hostStorage.contractsDb, 'getBytecode')
           .mockReturnValue(Promise.resolve(bytecode));
 
         const results = await new AvmSimulator(context).execute();
@@ -160,7 +164,7 @@ describe('AVM simulator', () => {
         // Decode bytecode into instructions
         const bytecode = Buffer.from(getterArtifact.bytecode, 'base64');
         jest
-          .spyOn(context.worldState.hostStorage.contractsDb, 'getBytecode')
+          .spyOn(context.persistableState.hostStorage.contractsDb, 'getBytecode')
           .mockReturnValue(Promise.resolve(bytecode));
         // Execute
 

yarn-project/simulator/src/avm/avm_simulator.ts

@@ -73,7 +73,7 @@ export class AvmSimulator {
     // NOTE: the following is mocked as getPublicBytecode does not exist yet
 
     const selector = this.context.environment.temporaryFunctionSelector;
-    const bytecode = await this.context.worldState.hostStorage.contractsDb.getBytecode(
+    const bytecode = await this.context.persistableState.hostStorage.contractsDb.getBytecode(
       this.context.environment.address,
       selector,
     );

yarn-project/simulator/src/avm/fixtures/index.ts

@@ -12,13 +12,13 @@ import { AvmContext } from '../avm_context.js';
 import { AvmExecutionEnvironment } from '../avm_execution_environment.js';
 import { AvmMachineState } from '../avm_machine_state.js';
 import { HostStorage } from '../journal/host_storage.js';
-import { AvmWorldStateJournal } from '../journal/journal.js';
+import { AvmPersistableStateManager } from '../journal/journal.js';
 
 /**
  * Create a new AVM context with default values.
  */
 export function initContext(overrides?: {
-  worldState?: AvmWorldStateJournal;
+  worldState?: AvmPersistableStateManager;
   env?: AvmExecutionEnvironment;
   machineState?: AvmMachineState;
 }): AvmContext {
@@ -30,9 +30,9 @@ export function initContext(overrides?: {
 }
 
 /** Creates an empty world state with mocked storage. */
-export function initMockWorldStateJournal(): AvmWorldStateJournal {
+export function initMockWorldStateJournal(): AvmPersistableStateManager {
   const hostStorage = new HostStorage(mock<PublicStateDB>(), mock<PublicContractsDB>(), mock<CommitmentsDB>());
-  return new AvmWorldStateJournal(hostStorage);
+  return new AvmPersistableStateManager(hostStorage);
 }
 
 /**

yarn-project/simulator/src/avm/journal/journal.test.ts

@@ -4,19 +4,19 @@ import { MockProxy, mock } from 'jest-mock-extended';
 
 import { CommitmentsDB, PublicContractsDB, PublicStateDB } from '../../index.js';
 import { HostStorage } from './host_storage.js';
-import { AvmWorldStateJournal, JournalData } from './journal.js';
+import { AvmPersistableStateManager, JournalData } from './journal.js';
 
 describe('journal', () => {
   let publicDb: MockProxy<PublicStateDB>;
-  let journal: AvmWorldStateJournal;
+  let journal: AvmPersistableStateManager;
 
   beforeEach(() => {
     publicDb = mock<PublicStateDB>();
     const commitmentsDb = mock<CommitmentsDB>();
     const contractsDb = mock<PublicContractsDB>();
 
     const hostStorage = new HostStorage(publicDb, contractsDb, commitmentsDb);
-    journal = new AvmWorldStateJournal(hostStorage);
+    journal = new AvmPersistableStateManager(hostStorage);
   });
 
   describe('Public Storage', () => {
@@ -102,15 +102,15 @@ describe('journal', () => {
     journal.writeL1Message(logs);
     journal.writeNullifier(commitment);
 
-    const childJournal = new AvmWorldStateJournal(journal.hostStorage, journal);
+    const childJournal = new AvmPersistableStateManager(journal.hostStorage, journal);
     childJournal.writeStorage(contractAddress, key, valueT1);
     await childJournal.readStorage(contractAddress, key);
     childJournal.writeNoteHash(commitmentT1);
     childJournal.writeLog(logsT1);
     childJournal.writeL1Message(logsT1);
     childJournal.writeNullifier(commitmentT1);
 
-    journal.acceptNestedWorldState(childJournal);
+    journal.acceptNestedCallState(childJournal);
 
     const result = await journal.readStorage(contractAddress, key);
     expect(result).toEqual(valueT1);
@@ -158,25 +158,24 @@ describe('journal', () => {
     journal.writeStorage(contractAddress, key, value);
     await journal.readStorage(contractAddress, key);
     journal.writeNoteHash(commitment);
+    journal.writeNullifier(commitment);
     journal.writeLog(logs);
     journal.writeL1Message(logs);
-    journal.writeNullifier(commitment);
 
-    const childJournal = new AvmWorldStateJournal(journal.hostStorage, journal);
+    const childJournal = new AvmPersistableStateManager(journal.hostStorage, journal);
     childJournal.writeStorage(contractAddress, key, valueT1);
     await childJournal.readStorage(contractAddress, key);
     childJournal.writeNoteHash(commitmentT1);
+    childJournal.writeNullifier(commitmentT1);
     childJournal.writeLog(logsT1);
     childJournal.writeL1Message(logsT1);
-    childJournal.writeNullifier(commitmentT1);
 
-    journal.rejectNestedWorldState(childJournal);
+    journal.rejectNestedCallState(childJournal);
 
     // 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
@@ -191,17 +190,20 @@ describe('journal', () => {
     const slotWrites = contractWrites?.get(key.toBigInt());
     expect(slotWrites).toEqual([value, valueT1]);
 
-    expect(journalUpdates.newNoteHashes).toEqual([commitment]);
+    // Check that the UTXOs _traces_ are merged even on rejection
+    expect(journalUpdates.newNoteHashes).toEqual([commitment, commitmentT1]);
+    expect(journalUpdates.newNullifiers).toEqual([commitment, commitmentT1]);
+
+    // Check that rejected Accrued Substate is absent
     expect(journalUpdates.newLogs).toEqual([logs]);
     expect(journalUpdates.newL1Messages).toEqual([logs]);
-    expect(journalUpdates.newNullifiers).toEqual([commitment]);
   });
 
   it('Can fork and merge journals', () => {
-    const rootJournal = new AvmWorldStateJournal(journal.hostStorage);
+    const rootJournal = new AvmPersistableStateManager(journal.hostStorage);
     const childJournal = rootJournal.fork();
 
-    expect(() => rootJournal.acceptNestedWorldState(childJournal));
-    expect(() => rootJournal.rejectNestedWorldState(childJournal));
+    expect(() => rootJournal.acceptNestedCallState(childJournal));
+    expect(() => rootJournal.rejectNestedCallState(childJournal));
   });
 });