chore: refactor AVM simulator's side-effect tracing

AztecProtocol · Jun 25, 2024 · 9701c1f · 9701c1f
1 parent 8d4245c
commit 9701c1f
Show file tree

Hide file tree

Showing 36 changed files with 2,128 additions and 2,336 deletions.
diff --git a/noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr b/noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr
@@ -363,19 +363,19 @@ contract AvmTest {
     // Use the standard context interface to check for a nullifier
     #[aztec(public)]
     fn nullifier_exists(nullifier: Field) -> bool {
-        context.nullifier_exists(nullifier, context.this_address())
+        context.nullifier_exists(nullifier, context.storage_address())
     }
 
     #[aztec(public)]
     fn assert_nullifier_exists(nullifier: Field) {
-        assert(context.nullifier_exists(nullifier, context.this_address()), "Nullifier doesn't exist!");
+        assert(context.nullifier_exists(nullifier, context.storage_address()), "Nullifier doesn't exist!");
     }
 
     // Use the standard context interface to emit a new nullifier
     #[aztec(public)]
     fn emit_nullifier_and_check(nullifier: Field) {
         context.push_new_nullifier(nullifier, 0);
-        let exists = context.nullifier_exists(nullifier, context.this_address());
+        let exists = context.nullifier_exists(nullifier, context.storage_address());
         assert(exists, "Nullifier was just created, but its existence wasn't detected!");
     }
 

diff --git a/yarn-project/bb-prover/src/avm_proving.test.ts b/yarn-project/bb-prover/src/avm_proving.test.ts
@@ -35,6 +35,7 @@ import {
   initContext,
   initExecutionEnvironment,
   initHostStorage,
+  initPersistableStateManager,
 } from '@aztec/simulator/avm/fixtures';
 
 import { jest } from '@jest/globals';
@@ -43,11 +44,7 @@ import fs from 'node:fs/promises';
 import { tmpdir } from 'node:os';
 import path from 'path';
 
-import { AvmPersistableStateManager } from '../../simulator/src/avm/journal/journal.js';
-import {
-  convertAvmResultsToPxResult,
-  createPublicExecution,
-} from '../../simulator/src/public/transitional_adaptors.js';
+import { PublicSideEffectTrace } from '../../simulator/src/public/side_effect_trace.js';
 import { SerializableContractInstance } from '../../types/src/contracts/contract_instance.js';
 import { type BBSuccess, BB_RESULT, generateAvmProof, verifyAvmProof } from './bb/execute.js';
 import { extractVkData } from './verification_key/verification_key_data.js';
@@ -224,15 +221,13 @@ const proveAndVerifyAvmTestContract = async (
   storageDb.storageRead.mockResolvedValue(Promise.resolve(storageValue));
 
   const hostStorage = initHostStorage({ contractsDb });
-  const persistableState = new AvmPersistableStateManager(hostStorage);
+  const trace = new PublicSideEffectTrace(startSideEffectCounter);
+  const persistableState = initPersistableStateManager({ hostStorage, trace });
   const context = initContext({ env: environment, persistableState });
   const nestedCallBytecode = getAvmTestContractBytecode('add_args_return');
-  jest
-    .spyOn(context.persistableState.hostStorage.contractsDb, 'getBytecode')
-    .mockReturnValue(Promise.resolve(nestedCallBytecode));
+  jest.spyOn(hostStorage.contractsDb, 'getBytecode').mockResolvedValue(nestedCallBytecode);
 
   const startGas = new Gas(context.machineState.gasLeft.daGas, context.machineState.gasLeft.l2Gas);
-  const oldPublicExecution = createPublicExecution(startSideEffectCounter, environment, calldata);
 
   const internalLogger = createDebugLogger('aztec:avm-proving-test');
   const logger = (msg: string, _data?: any) => internalLogger.verbose(msg);
@@ -255,25 +250,21 @@ const proveAndVerifyAvmTestContract = async (
     expect(avmResult.revertReason?.message).toContain(assertionErrString);
   }
 
-  const pxResult = convertAvmResultsToPxResult(
-    avmResult,
-    startSideEffectCounter,
-    oldPublicExecution,
+  const pxResult = trace.toPublicExecutionResult(
+    environment,
     startGas,
-    context,
-    simulator.getBytecode(),
+    /*endGasLeft=*/ Gas.from(context.machineState.gasLeft),
+    /*bytecode=*/ simulator.getBytecode()!,
+    avmResult,
     functionName,
   );
-  // TODO(dbanks12): public inputs should not be empty.... Need to construct them from AvmContext?
-  const uncompressedBytecode = simulator.getBytecode()!;
-  const publicInputs = getPublicInputs(pxResult);
 
   const avmCircuitInputs = new AvmCircuitInputs(
     functionName,
-    uncompressedBytecode,
-    context.environment.calldata,
-    publicInputs,
-    pxResult.avmHints,
+    /*bytecode=*/ simulator.getBytecode()!, // uncompressed bytecode
+    /*calldata=*/ context.environment.calldata,
+    /*publicInputs=*/ getPublicInputs(pxResult),
+    /*avmHints=*/ pxResult.avmCircuitHints,
   );
 
   // Then we prove.

diff --git a/yarn-project/circuit-types/src/mocks.ts b/yarn-project/circuit-types/src/mocks.ts
@@ -203,8 +203,10 @@ export const randomContractArtifact = (): ContractArtifact => ({
   notes: {},
 });
 
-export const randomContractInstanceWithAddress = (opts: { contractClassId?: Fr } = {}): ContractInstanceWithAddress =>
-  SerializableContractInstance.random(opts).withAddress(AztecAddress.random());
+export const randomContractInstanceWithAddress = (
+  opts: { contractClassId?: Fr } = {},
+  address: AztecAddress = AztecAddress.random(),
+): ContractInstanceWithAddress => SerializableContractInstance.random(opts).withAddress(address);
 
 export const randomDeployedContract = () => {
   const artifact = randomContractArtifact();

diff --git a/yarn-project/circuits.js/src/structs/contract_storage_read.ts b/yarn-project/circuits.js/src/structs/contract_storage_read.ts
@@ -30,12 +30,7 @@ export class ContractStorageRead {
     public contractAddress?: AztecAddress, // TODO: Should not be optional. This is a temporary hack to silo the storage slot with the correct address for nested executions.
   ) {}
 
-  static from(args: {
-    storageSlot: Fr;
-    currentValue: Fr;
-    counter: number;
-    contractAddress?: AztecAddress;
-  }) {
+  static from(args: { storageSlot: Fr; currentValue: Fr; counter: number; contractAddress?: AztecAddress }) {
     return new ContractStorageRead(args.storageSlot, args.currentValue, args.counter, args.contractAddress);
   }
 

diff --git a/yarn-project/circuits.js/src/structs/contract_storage_update_request.ts b/yarn-project/circuits.js/src/structs/contract_storage_update_request.ts
@@ -22,14 +22,17 @@ export class ContractStorageUpdateRequest {
      */
     public readonly newValue: Fr,
     /**
-     * Optional side effect counter tracking position of this event in tx execution.
+     * Side effect counter tracking position of this event in tx execution.
+     */
+    public readonly counter: number,
+    /**
+     * Contract address whose storage is being read.
      */
-    public readonly sideEffectCounter: number,
     public contractAddress?: AztecAddress, // TODO: Should not be optional. This is a temporary hack to silo the storage slot with the correct address for nested executions.
   ) {}
 
   toBuffer() {
-    return serializeToBuffer(this.storageSlot, this.newValue, this.sideEffectCounter);
+    return serializeToBuffer(this.storageSlot, this.newValue, this.counter);
   }
 
   static fromBuffer(buffer: Buffer | BufferReader) {
@@ -52,7 +55,7 @@ export class ContractStorageUpdateRequest {
    * @returns The array.
    */
   static getFields(fields: FieldsOf<ContractStorageUpdateRequest>) {
-    return [fields.storageSlot, fields.newValue, fields.sideEffectCounter, fields.contractAddress] as const;
+    return [fields.storageSlot, fields.newValue, fields.counter, fields.contractAddress] as const;
   }
 
   static empty() {
@@ -65,12 +68,12 @@ export class ContractStorageUpdateRequest {
 
   toFriendlyJSON() {
     return `Slot=${this.storageSlot.toFriendlyJSON()}: ${this.newValue.toFriendlyJSON()}, sideEffectCounter=${
-      this.sideEffectCounter
+      this.counter
     }`;
   }
 
   toFields(): Fr[] {
-    const fields = [this.storageSlot, this.newValue, new Fr(this.sideEffectCounter)];
+    const fields = [this.storageSlot, this.newValue, new Fr(this.counter)];
     if (fields.length !== CONTRACT_STORAGE_UPDATE_REQUEST_LENGTH) {
       throw new Error(
         `Invalid number of fields for ContractStorageUpdateRequest. Expected ${CONTRACT_STORAGE_UPDATE_REQUEST_LENGTH}, got ${fields.length}`,

diff --git a/yarn-project/noir-protocol-circuits-types/src/type_conversion.ts b/yarn-project/noir-protocol-circuits-types/src/type_conversion.ts
@@ -1826,7 +1826,7 @@ export function mapStorageUpdateRequestToNoir(
   return {
     storage_slot: mapFieldToNoir(storageUpdateRequest.storageSlot),
     new_value: mapFieldToNoir(storageUpdateRequest.newValue),
-    counter: mapNumberToNoir(storageUpdateRequest.sideEffectCounter),
+    counter: mapNumberToNoir(storageUpdateRequest.counter),
   };
 }
 /**
@@ -1855,7 +1855,7 @@ export function mapStorageReadToNoir(storageRead: ContractStorageRead): StorageR
   return {
     storage_slot: mapFieldToNoir(storageRead.storageSlot),
     current_value: mapFieldToNoir(storageRead.currentValue),
-    counter: mapNumberToNoir(storageRead.sideEffectCounter),
+    counter: mapNumberToNoir(storageRead.counter),
   };
 }
 /**

diff --git a/yarn-project/simulator/src/avm/avm_context.test.ts b/yarn-project/simulator/src/avm/avm_context.test.ts
@@ -16,6 +16,7 @@ describe('Avm Context', () => {
       allSameExcept(context.environment, {
         address: newAddress,
         storageAddress: newAddress,
+        contractCallDepth: Fr.ONE,
         // Calldata also includes AvmContextInputs
         calldata: anyAvmContextInputs().concat(newCalldata),
         isStaticCall: false,
@@ -46,6 +47,7 @@ describe('Avm Context', () => {
       allSameExcept(context.environment, {
         address: newAddress,
         storageAddress: newAddress,
+        contractCallDepth: Fr.ONE,
         // Calldata also includes AvmContextInputs
         calldata: anyAvmContextInputs().concat(newCalldata),
         isStaticCall: true,

diff --git a/yarn-project/simulator/src/avm/avm_execution_environment.test.ts b/yarn-project/simulator/src/avm/avm_execution_environment.test.ts
@@ -16,6 +16,7 @@ describe('Execution Environment', () => {
       allSameExcept(executionEnvironment, {
         address: newAddress,
         storageAddress: newAddress,
+        contractCallDepth: Fr.ONE,
         // Calldata also includes AvmContextInputs
         calldata: anyAvmContextInputs().concat(calldata),
       }),
@@ -30,6 +31,7 @@ describe('Execution Environment', () => {
     expect(newExecutionEnvironment).toEqual(
       allSameExcept(executionEnvironment, {
         address: newAddress,
+        contractCallDepth: Fr.ONE,
         isDelegateCall: true,
         // Calldata also includes AvmContextInputs
         calldata: anyAvmContextInputs().concat(calldata),
@@ -49,6 +51,7 @@ describe('Execution Environment', () => {
       allSameExcept(executionEnvironment, {
         address: newAddress,
         storageAddress: newAddress,
+        contractCallDepth: Fr.ONE,
         isStaticCall: true,
         // Calldata also includes AvmContextInputs
         calldata: anyAvmContextInputs().concat(calldata),

diff --git a/yarn-project/simulator/src/avm/avm_execution_environment.ts b/yarn-project/simulator/src/avm/avm_execution_environment.ts
@@ -62,7 +62,7 @@ export class AvmExecutionEnvironment {
       /*sender=*/ this.address,
       this.feePerL2Gas,
       this.feePerDaGas,
-      this.contractCallDepth,
+      this.contractCallDepth.add(Fr.ONE),
       this.header,
       this.globals,
       isStaticCall,