feat(avm): initial external calls (#4194)

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

@@ -1,7 +1,10 @@
+import { AztecAddress, FunctionSelector } from '@aztec/circuits.js';
+import { Fr } from '@aztec/foundation/fields';
+
 import { AvmExecutionEnvironment } from './avm_execution_environment.js';
 import { AvmMachineState } from './avm_machine_state.js';
 import { AvmMessageCallResult } from './avm_message_call_result.js';
-import { AvmInterpreter } from './interpreter/index.js';
+import { AvmInterpreter, AvmInterpreterError } from './interpreter/index.js';
 import { AvmJournal } from './journal/journal.js';
 import { decodeBytecode } from './opcodes/decode_bytecode.js';
 import { Instruction } from './opcodes/index.js';
@@ -30,10 +33,19 @@ export class AvmContext {
    * - We run the interpreter
    *
    */
-  public call(): AvmMessageCallResult {
+  async call(): Promise<AvmMessageCallResult> {
     // NOTE: the following is mocked as getPublicBytecode does not exist yet
-    // const bytecode = journal.journal.hostStorage.contractsDb.getBytecode(this.executionEnvironment.address);
-    const bytecode = Buffer.from('0x01000100020003');
+    const selector = new FunctionSelector(0);
+    const bytecode = await this.journal.hostStorage.contractsDb.getBytecode(
+      this.executionEnvironment.address,
+      selector,
+    );
+
+    // This assumes that we will not be able to send messages to accounts without code
+    // Pending classes and instances impl details
+    if (!bytecode) {
+      throw new NoBytecodeFoundInterpreterError(this.executionEnvironment.address);
+    }
 
     const instructions: Instruction[] = decodeBytecode(bytecode);
 
@@ -42,4 +54,79 @@ export class AvmContext {
 
     return interpreter.run();
   }
+
+  /**
+   * Create a new forked avm context - for internal calls
+   */
+  public newWithForkedState(): AvmContext {
+    const forkedState = AvmJournal.branchParent(this.journal);
+    return new AvmContext(this.executionEnvironment, forkedState);
+  }
+
+  /**
+   * Create a new forked avm context - for external calls
+   */
+  public static newWithForkedState(executionEnvironment: AvmExecutionEnvironment, journal: AvmJournal): AvmContext {
+    const forkedState = AvmJournal.branchParent(journal);
+    return new AvmContext(executionEnvironment, forkedState);
+  }
+
+  /**
+   * Prepare a new AVM context that will be ready for an external call
+   * - It will fork the journal
+   * - It will set the correct execution Environment Variables for a call
+   *    - Alter both address and storageAddress
+   *
+   * @param address - The contract to call
+   * @param executionEnvironment - The current execution environment
+   * @param journal - The current journal
+   * @returns new AvmContext instance
+   */
+  public static prepExternalCallContext(
+    address: AztecAddress,
+    calldata: Fr[],
+    executionEnvironment: AvmExecutionEnvironment,
+    journal: AvmJournal,
+  ): AvmContext {
+    const newExecutionEnvironment = executionEnvironment.newCall(address, calldata);
+    const forkedState = AvmJournal.branchParent(journal);
+    return new AvmContext(newExecutionEnvironment, forkedState);
+  }
+
+  /**
+   * Prepare a new AVM context that will be ready for an external static call
+   * - It will fork the journal
+   * - It will set the correct execution Environment Variables for a call
+   *    - Alter both address and storageAddress
+   *
+   * @param address - The contract to call
+   * @param executionEnvironment - The current execution environment
+   * @param journal - The current journal
+   * @returns new AvmContext instance
+   */
+  public static prepExternalStaticCallContext(
+    address: AztecAddress,
+    calldata: Fr[],
+    executionEnvironment: AvmExecutionEnvironment,
+    journal: AvmJournal,
+  ): AvmContext {
+    const newExecutionEnvironment = executionEnvironment.newStaticCall(address, calldata);
+    const forkedState = AvmJournal.branchParent(journal);
+    return new AvmContext(newExecutionEnvironment, forkedState);
+  }
+
+  /**
+   * 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();
+  }
+}
+
+class NoBytecodeFoundInterpreterError extends AvmInterpreterError {
+  constructor(contractAddress: AztecAddress) {
+    super(`No bytecode found at: ${contractAddress}`);
+    this.name = 'NoBytecodeFoundInterpreterError';
+  }
 }

yarn-project/acir-simulator/src/avm/avm_execution_environment.test.ts

@@ -0,0 +1,55 @@
+import { Fr } from '@aztec/foundation/fields';
+
+import { initExecutionEnvironment } from './fixtures/index.js';
+
+describe('Execution Environment', () => {
+  const newAddress = new Fr(123456n);
+  const calldata = [new Fr(1n), new Fr(2n), new Fr(3n)];
+
+  it('New call should fork execution environment correctly', () => {
+    const executionEnvironment = initExecutionEnvironment();
+    const newExecutionEnvironment = executionEnvironment.newCall(newAddress, calldata);
+
+    allTheSameExcept(executionEnvironment, newExecutionEnvironment, {
+      address: newAddress,
+      storageAddress: newAddress,
+      calldata,
+    });
+  });
+
+  it('New delegate call should fork execution environment correctly', () => {
+    const executionEnvironment = initExecutionEnvironment();
+    const newExecutionEnvironment = executionEnvironment.newDelegateCall(newAddress, calldata);
+
+    allTheSameExcept(executionEnvironment, newExecutionEnvironment, {
+      address: newAddress,
+      isDelegateCall: true,
+      calldata,
+    });
+  });
+
+  it('New static call call should fork execution environment correctly', () => {
+    const executionEnvironment = initExecutionEnvironment();
+    const newExecutionEnvironment = executionEnvironment.newStaticCall(newAddress, calldata);
+
+    allTheSameExcept(executionEnvironment, newExecutionEnvironment, {
+      address: newAddress,
+      storageAddress: newAddress,
+      isStaticCall: true,
+      calldata,
+    });
+  });
+});
+
+/**
+ * Check all properties of one object are the same, except for the specified differentProperties
+ */
+function allTheSameExcept(referenceObject: any, comparingObject: any, differentProperties: Record<string, any>): void {
+  for (const key in referenceObject) {
+    if (Object.keys(differentProperties).includes(key)) {
+      expect(comparingObject[key]).toEqual(differentProperties[key]);
+    } else {
+      expect(comparingObject[key]).toEqual(referenceObject[key]);
+    }
+  }
+}

yarn-project/acir-simulator/src/avm/avm_execution_environment.ts

@@ -7,6 +7,7 @@ import { Fr } from '@aztec/foundation/fields';
  * Contains variables that remain constant during AVM execution
  * These variables are provided by the public kernel circuit
  */
+// TODO(https://github.com/AztecProtocol/aztec-packages/issues/3992): gas not implemented
 export class AvmExecutionEnvironment {
   constructor(
     /** - */
@@ -36,4 +37,58 @@ export class AvmExecutionEnvironment {
     /** - */
     public readonly calldata: Fr[],
   ) {}
+
+  public newCall(address: AztecAddress, calldata: Fr[]): AvmExecutionEnvironment {
+    return new AvmExecutionEnvironment(
+      /*address=*/ address,
+      /*storageAddress=*/ address,
+      this.origin,
+      this.sender,
+      this.portal,
+      this.feePerL1Gas,
+      this.feePerL2Gas,
+      this.feePerDaGas,
+      this.contractCallDepth,
+      this.globals,
+      this.isStaticCall,
+      this.isDelegateCall,
+      /*calldata=*/ calldata,
+    );
+  }
+
+  public newStaticCall(address: AztecAddress, calldata: Fr[]): AvmExecutionEnvironment {
+    return new AvmExecutionEnvironment(
+      /*address=*/ address,
+      /*storageAddress=*/ address,
+      this.origin,
+      this.sender,
+      this.portal,
+      this.feePerL1Gas,
+      this.feePerL2Gas,
+      this.feePerDaGas,
+      this.contractCallDepth,
+      this.globals,
+      /*isStaticCall=*/ true,
+      this.isDelegateCall,
+      /*calldata=*/ calldata,
+    );
+  }
+
+  public newDelegateCall(address: AztecAddress, calldata: Fr[]): AvmExecutionEnvironment {
+    return new AvmExecutionEnvironment(
+      /*address=*/ address,
+      this.storageAddress,
+      this.origin,
+      this.sender,
+      this.portal,
+      this.feePerL1Gas,
+      this.feePerL2Gas,
+      this.feePerDaGas,
+      this.contractCallDepth,
+      this.globals,
+      this.isStaticCall,
+      /*isDelegateCall=*/ true,
+      /*calldata=*/ calldata,
+    );
+  }
 }

yarn-project/acir-simulator/src/avm/index.test.ts

@@ -11,7 +11,7 @@ import { encodeToBytecode } from './opcodes/encode_to_bytecode.js';
 import { Opcode } from './opcodes/opcodes.js';
 
 describe('avm', () => {
-  it('Should execute bytecode', () => {
+  it('Should execute bytecode', async () => {
     const calldata: Fr[] = [new Fr(1), new Fr(2)];
     const journal = mock<AvmJournal>();
 
@@ -31,7 +31,7 @@ describe('avm', () => {
     // Execute instructions
     const context = new AvmMachineState(initExecutionEnvironment({ calldata }));
     const interpreter = new AvmInterpreter(context, journal, instructions);
-    const avmReturnData = interpreter.run();
+    const avmReturnData = await interpreter.run();
 
     expect(avmReturnData.reverted).toBe(false);
 

yarn-project/acir-simulator/src/avm/interpreter/interpreter.test.ts

@@ -18,7 +18,7 @@ describe('interpreter', () => {
     journal = mock<AvmJournal>();
   });
 
-  it('Should execute a series of instructions', () => {
+  it('Should execute a series of instructions', async () => {
     const calldata: Fr[] = [new Fr(1), new Fr(2)];
 
     const instructions: Instruction[] = [
@@ -27,26 +27,26 @@ describe('interpreter', () => {
       new Return(/*returnOffset=*/ 2, /*copySize=*/ 1),
     ];
 
-    const context = new AvmMachineState(initExecutionEnvironment({ calldata }));
-    const interpreter = new AvmInterpreter(context, journal, instructions);
-    const avmReturnData = interpreter.run();
+    const machineState = new AvmMachineState(initExecutionEnvironment({ calldata }));
+    const interpreter = new AvmInterpreter(machineState, journal, instructions);
+    const avmReturnData = await interpreter.run();
 
     expect(avmReturnData.reverted).toBe(false);
     expect(avmReturnData.revertReason).toBeUndefined();
     expect(avmReturnData.output).toEqual([new Fr(3)]);
   });
 
-  it('Should revert with an invalid jump', () => {
+  it('Should revert with an invalid jump', async () => {
     const calldata: Fr[] = [];
 
     const invalidJumpDestination = 22;
 
     const instructions: Instruction[] = [new Jump(invalidJumpDestination)];
 
-    const context = new AvmMachineState(initExecutionEnvironment({ calldata }));
-    const interpreter = new AvmInterpreter(context, journal, instructions);
+    const machineState = new AvmMachineState(initExecutionEnvironment({ calldata }));
+    const interpreter = new AvmInterpreter(machineState, journal, instructions);
 
-    const avmReturnData = interpreter.run();
+    const avmReturnData = await interpreter.run();
 
     expect(avmReturnData.reverted).toBe(true);
     expect(avmReturnData.revertReason).toBeInstanceOf(InvalidProgramCounterError);

yarn-project/acir-simulator/src/avm/interpreter/interpreter.ts

@@ -29,15 +29,15 @@ export class AvmInterpreter {
    *               - reverted execution will return false
    *               - any other panic will throw
    */
-  run(): AvmMessageCallResult {
+  async run(): Promise<AvmMessageCallResult> {
     assert(this.instructions.length > 0);
 
     try {
       while (!this.machineState.halted) {
         const instruction = this.instructions[this.machineState.pc];
         assert(!!instruction); // This should never happen
 
-        instruction.execute(this.machineState, this.journal);
+        await instruction.execute(this.machineState, this.journal);
 
         if (this.machineState.pc >= this.instructions.length) {
           throw new InvalidProgramCounterError(this.machineState.pc, /*max=*/ this.instructions.length);

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

@@ -29,7 +29,7 @@ describe('journal', () => {
       journal.writeStorage(contractAddress, key, value);
 
       const journalUpdates: JournalData = journal.flush();
-      expect(journalUpdates.storageWrites.get(contractAddress)?.get(key)).toEqual(value);
+      expect(journalUpdates.storageWrites.get(contractAddress.toBigInt())?.get(key.toBigInt())).toEqual(value);
     });
 
     it('When reading from storage, should check the parent first', async () => {

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

@@ -14,7 +14,7 @@ export type JournalData = {
   /** - */
   newNullifiers: Fr[];
   /** contract address -\> key -\> value */
-  storageWrites: Map<Fr, Map<Fr, Fr>>;
+  storageWrites: Map<bigint, Map<bigint, Fr>>;
 };
 
 /**
@@ -32,7 +32,7 @@ export class AvmJournal {
   // Reading state - must be tracked for vm execution
   // contract address -> key -> value
   // TODO(https://github.com/AztecProtocol/aztec-packages/issues/3999)
-  private storageReads: Map<Fr, Map<Fr, Fr>> = new Map();
+  private storageReads: Map<bigint, Map<bigint, Fr>> = new Map();
 
   // New written state
   private newCommitments: Fr[] = [];
@@ -43,7 +43,7 @@ export class AvmJournal {
   private newLogs: Fr[][] = [];
 
   // contract address -> key -> value
-  private storageWrites: Map<Fr, Map<Fr, Fr>> = new Map();
+  private storageWrites: Map<bigint, Map<bigint, Fr>> = new Map();
 
   private parentJournal: AvmJournal | undefined;
 
@@ -76,12 +76,12 @@ export class AvmJournal {
    * @param value -
    */
   public writeStorage(contractAddress: Fr, key: Fr, value: Fr) {
-    let contractMap = this.storageWrites.get(contractAddress);
+    let contractMap = this.storageWrites.get(contractAddress.toBigInt());
     if (!contractMap) {
       contractMap = new Map();
-      this.storageWrites.set(contractAddress, contractMap);
+      this.storageWrites.set(contractAddress.toBigInt(), contractMap);
     }
-    contractMap.set(key, value);
+    contractMap.set(key.toBigInt(), value);
   }
 
   /**
@@ -93,7 +93,7 @@ export class AvmJournal {
    * @returns current value
    */
   public readStorage(contractAddress: Fr, key: Fr): Promise<Fr> {
-    const cachedValue = this.storageWrites.get(contractAddress)?.get(key);
+    const cachedValue = this.storageWrites.get(contractAddress.toBigInt())?.get(key.toBigInt());
     if (cachedValue) {
       return Promise.resolve(cachedValue);
     }
@@ -168,7 +168,7 @@ export class AvmJournal {
  * @param hostMap - The map to be merged into
  * @param childMap - The map to be merged from
  */
-function mergeContractMaps(hostMap: Map<Fr, Map<Fr, Fr>>, childMap: Map<Fr, Map<Fr, Fr>>) {
+function mergeContractMaps(hostMap: Map<bigint, Map<bigint, Fr>>, childMap: Map<bigint, Map<bigint, Fr>>) {
   for (const [key, value] of childMap) {
     const map1Value = hostMap.get(key);
     if (!map1Value) {
@@ -184,7 +184,7 @@ function mergeContractMaps(hostMap: Map<Fr, Map<Fr, Fr>>, childMap: Map<Fr, Map<
  * @param hostMap - The map to be merge into
  * @param childMap - The map to be merged from
  */
-function mergeStorageMaps(hostMap: Map<Fr, Fr>, childMap: Map<Fr, Fr>) {
+function mergeStorageMaps(hostMap: Map<bigint, Fr>, childMap: Map<bigint, Fr>) {
   for (const [key, value] of childMap) {
     hostMap.set(key, value);
   }