feat(avm): Gas usage for nested calls (#5495)

Adds gas metering for CALLs and STATICCALLs in the AVM. Both opcodes
have a fixed gas cost, an extra cost for indirect accesses to memory,
and consume whatever gas they pass onto the nested call. The unused gas
from the nested call gets refunded when the call returns, as specified
on the yellow paper.

noir-projects/aztec-nr/aztec/src/context/avm_context.nr

@@ -123,7 +123,7 @@ impl PublicContextInterface for AvmContext {
         temporary_function_selector: FunctionSelector,
         args: [Field; ARGS_COUNT]
     ) -> [Field; RETURN_VALUES_LENGTH] {
-        let gas = [/*l1_gas*/42, /*l2_gas*/24, /*da_gas*/420];
+        let gas = [/*l1_gas*/10000, /*l2_gas*/10000, /*da_gas*/10000];
 
         let results = call(
             gas,
@@ -144,7 +144,7 @@ impl PublicContextInterface for AvmContext {
         temporary_function_selector: FunctionSelector,
         args: [Field; ARGS_COUNT]
     ) -> [Field; RETURN_VALUES_LENGTH] {
-        let gas = [/*l1_gas*/42, /*l2_gas*/24, /*da_gas*/420];
+        let gas = [/*l1_gas*/10000, /*l2_gas*/10000, /*da_gas*/10000];
 
         let (data_to_return, success): ([Field; RETURN_VALUES_LENGTH], u8) = call_static(
             gas,

noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr

@@ -339,7 +339,7 @@ contract AvmTest {
     #[aztec(public-vm)]
     fn raw_nested_call_to_add(arg_a: Field, arg_b: Field) -> pub Field {
         let selector = FunctionSelector::from_signature("add_args_return(Field,Field)").to_field();
-        let gas = [/*l1_gas*/42, /*l2_gas*/24, /*da_gas*/420];
+        let gas = [/*l1_gas*/10000, /*l2_gas*/10000, /*da_gas*/10000];
 
         // Nested call
         let results = context.call_public_function_raw(gas, context.this_address(), selector, [arg_a, arg_b]);
@@ -372,7 +372,7 @@ contract AvmTest {
     #[aztec(public-vm)]
     fn raw_nested_static_call_to_add(arg_a: Field, arg_b: Field) -> pub (Field, u8) {
         let selector = FunctionSelector::from_signature("add_args_return(Field,Field)").to_field();
-        let gas = [/*l1_gas*/42, /*l2_gas*/24, /*da_gas*/420];
+        let gas = [/*l1_gas*/10000, /*l2_gas*/10000, /*da_gas*/10000];
 
         let (result_data, success): ([Field; 1], u8) = context.static_call_public_function_raw(gas, context.this_address(), selector, [arg_a, arg_b]);
 
@@ -383,7 +383,7 @@ contract AvmTest {
     #[aztec(public-vm)]
     fn raw_nested_static_call_to_set_storage() -> pub u8 {
         let selector = FunctionSelector::from_signature("set_storage_single(Field)").to_field();
-        let gas = [/*l1_gas*/42, /*l2_gas*/24, /*da_gas*/420];
+        let gas = [/*l1_gas*/10000, /*l2_gas*/10000, /*da_gas*/10000];
         let calldata: [Field; 1] = [20];
 
         let (_data_to_return, success): ([Field; 0], u8) = context.static_call_public_function_raw(gas, context.this_address(), selector, calldata);

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

@@ -9,7 +9,8 @@ describe('Avm Context', () => {
 
     const newAddress = AztecAddress.random();
     const newCalldata = [new Fr(1), new Fr(2)];
-    const newContext = context.createNestedContractCallContext(newAddress, newCalldata);
+    const allocatedGas = { l1Gas: 1, l2Gas: 2, daGas: 3 }; // How much of the current call gas we pass to the nested call
+    const newContext = context.createNestedContractCallContext(newAddress, newCalldata, allocatedGas, 'CALL');
 
     expect(newContext.environment).toEqual(
       allSameExcept(context.environment, {
@@ -23,6 +24,9 @@ describe('Avm Context', () => {
     expect(newContext.machineState).toEqual(
       allSameExcept(context.machineState, {
         pc: 0,
+        l1GasLeft: 1,
+        l2GasLeft: 2,
+        daGasLeft: 3,
       }),
     );
 
@@ -36,7 +40,8 @@ describe('Avm Context', () => {
 
     const newAddress = AztecAddress.random();
     const newCalldata = [new Fr(1), new Fr(2)];
-    const newContext = context.createNestedContractStaticCallContext(newAddress, newCalldata);
+    const allocatedGas = { l1Gas: 1, l2Gas: 2, daGas: 3 };
+    const newContext = context.createNestedContractCallContext(newAddress, newCalldata, allocatedGas, 'STATICCALL');
 
     expect(newContext.environment).toEqual(
       allSameExcept(context.environment, {
@@ -50,6 +55,9 @@ describe('Avm Context', () => {
     expect(newContext.machineState).toEqual(
       allSameExcept(context.machineState, {
         pc: 0,
+        l1GasLeft: 1,
+        l2GasLeft: 2,
+        daGasLeft: 3,
       }),
     );
 

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

@@ -2,6 +2,7 @@ import { type AztecAddress, FunctionSelector } from '@aztec/circuits.js';
 import { type Fr } from '@aztec/foundation/fields';
 
 import { type AvmExecutionEnvironment } from './avm_execution_environment.js';
+import { type Gas, gasToGasLeft } from './avm_gas.js';
 import { AvmMachineState } from './avm_machine_state.js';
 import { type AvmPersistableStateManager } from './journal/journal.js';
 
@@ -33,47 +34,24 @@ export class AvmContext {
    *
    * @param address - The contract instance to initialize a context for
    * @param calldata - Data/arguments for nested call
+   * @param allocatedGas - Gas allocated for the nested call
+   * @param callType - Type of call (CALL or STATICCALL)
    * @returns new AvmContext instance
    */
   public createNestedContractCallContext(
     address: AztecAddress,
     calldata: Fr[],
+    allocatedGas: Gas,
+    callType: 'CALL' | 'STATICCALL',
     temporaryFunctionSelector: FunctionSelector = FunctionSelector.empty(),
   ): AvmContext {
-    const newExecutionEnvironment = this.environment.deriveEnvironmentForNestedCall(
-      address,
-      calldata,
-      temporaryFunctionSelector,
-    );
+    const deriveFn =
+      callType === 'CALL'
+        ? this.environment.deriveEnvironmentForNestedCall
+        : this.environment.deriveEnvironmentForNestedStaticCall;
+    const newExecutionEnvironment = deriveFn.call(this.environment, address, calldata, temporaryFunctionSelector);
     const forkedWorldState = this.persistableState.fork();
-    const machineState = AvmMachineState.fromState(this.machineState);
-    return new AvmContext(forkedWorldState, newExecutionEnvironment, machineState);
-  }
-
-  /**
-   * Prepare a new AVM context that will be ready for an external/nested static call
-   * - Fork the world state journal
-   * - Derive a machine state from the current state
-   *   - E.g., gas metering is preserved but pc is reset
-   * - Derive an execution environment from the caller/parent
-   *   - Alter both address and storageAddress
-   *
-   * @param address - The contract instance to initialize a context for
-   * @param calldata - Data/arguments for nested call
-   * @returns new AvmContext instance
-   */
-  public createNestedContractStaticCallContext(
-    address: AztecAddress,
-    calldata: Fr[],
-    temporaryFunctionSelector: FunctionSelector = FunctionSelector.empty(),
-  ): AvmContext {
-    const newExecutionEnvironment = this.environment.deriveEnvironmentForNestedStaticCall(
-      address,
-      calldata,
-      temporaryFunctionSelector,
-    );
-    const forkedWorldState = this.persistableState.fork();
-    const machineState = AvmMachineState.fromState(this.machineState);
+    const machineState = AvmMachineState.fromState(gasToGasLeft(allocatedGas));
     return new AvmContext(forkedWorldState, newExecutionEnvironment, machineState);
   }
 }

yarn-project/simulator/src/avm/avm_gas_cost.ts → yarn-project/simulator/src/avm/avm_gas.ts

@@ -1,20 +1,43 @@
 import { TypeTag } from './avm_memory_types.js';
+import { Addressing, AddressingMode } from './opcodes/addressing_mode.js';
 import { Opcode } from './serialization/instruction_serialization.js';
 
-/** Gas cost in L1, L2, and DA for a given instruction. */
-export type GasCost = {
+/** Gas counters in L1, L2, and DA. */
+export type Gas = {
   l1Gas: number;
   l2Gas: number;
   daGas: number;
 };
 
+/** Maps a Gas struct to gasLeft properties. */
+export function gasToGasLeft(gas: Gas) {
+  return { l1GasLeft: gas.l1Gas, l2GasLeft: gas.l2Gas, daGasLeft: gas.daGas };
+}
+
+/** Maps gasLeft properties to a gas struct. */
+export function gasLeftToGas(gasLeft: { l1GasLeft: number; l2GasLeft: number; daGasLeft: number }) {
+  return { l1Gas: gasLeft.l1GasLeft, l2Gas: gasLeft.l2GasLeft, daGas: gasLeft.daGasLeft };
+}
+
 /** Creates a new instance with all values set to zero except the ones set. */
-export function makeGasCost(gasCost: Partial<GasCost>) {
-  return { ...EmptyGasCost, ...gasCost };
+export function makeGasCost(gasCost: Partial<Gas>) {
+  return { ...EmptyGas, ...gasCost };
+}
+
+/** Sums together multiple instances of Gas. */
+export function sumGas(...gases: Partial<Gas>[]) {
+  return gases.reduce(
+    (acc: Gas, gas) => ({
+      l1Gas: acc.l1Gas + (gas.l1Gas ?? 0),
+      l2Gas: acc.l2Gas + (gas.l2Gas ?? 0),
+      daGas: acc.daGas + (gas.daGas ?? 0),
+    }),
+    EmptyGas,
+  );
 }
 
-/** Gas cost of zero across all gas dimensions. */
-export const EmptyGasCost = {
+/** Zero gas across all gas dimensions. */
+export const EmptyGas: Gas = {
   l1Gas: 0,
   l2Gas: 0,
   daGas: 0,
@@ -103,12 +126,29 @@ export const GasCosts = {
   [Opcode.PEDERSEN]: TemporaryDefaultGasCost, // temp - may be removed, but alot of contracts rely on i: TemporaryDefaultGasCost,t
 } as const;
 
+/** Returns the fixed gas cost for a given opcode, or throws if set to dynamic. */
+export function getFixedGasCost(opcode: Opcode): Gas {
+  const cost = GasCosts[opcode];
+  if (cost === DynamicGasCost) {
+    throw new Error(`Opcode ${Opcode[opcode]} has dynamic gas cost`);
+  }
+  return cost;
+}
+
+/** Returns the additional cost from indirect accesses to memory. */
+export function getCostFromIndirectAccess(indirect: number): Partial<Gas> {
+  const indirectCount = Addressing.fromWire(indirect).modePerOperand.filter(
+    mode => mode === AddressingMode.INDIRECT,
+  ).length;
+  return { l2Gas: indirectCount * GasCostConstants.COST_PER_INDIRECT_ACCESS };
+}
+
 /** Constants used in base cost calculations. */
 export const GasCostConstants = {
   SET_COST_PER_BYTE: 100,
   CALLDATACOPY_COST_PER_BYTE: 10,
   ARITHMETIC_COST_PER_BYTE: 10,
-  ARITHMETIC_COST_PER_INDIRECT_ACCESS: 5,
+  COST_PER_INDIRECT_ACCESS: 5,
 };
 
 /** Returns a multiplier based on the size of the type represented by the tag. Throws on uninitialized or invalid. */

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

@@ -1,6 +1,6 @@
 import { type Fr } from '@aztec/circuits.js';
 
-import { type GasCost, GasDimensions } from './avm_gas_cost.js';
+import { type Gas, GasDimensions } from './avm_gas.js';
 import { TaggedMemory } from './avm_memory_types.js';
 import { AvmContractCallResults } from './avm_message_call_result.js';
 import { OutOfGasError } from './errors.js';
@@ -59,7 +59,7 @@ export class AvmMachineState {
    * Should any of the gas dimensions get depleted, it sets all gas left to zero and triggers
    * an exceptional halt by throwing an OutOfGasError.
    */
-  public consumeGas(gasCost: Partial<GasCost>) {
+  public consumeGas(gasCost: Partial<Gas>) {
     // Assert there is enough gas on every dimension.
     const outOfGasDimensions = GasDimensions.filter(
       dimension => this[`${dimension}Left`] - (gasCost[dimension] ?? 0) < 0,
@@ -76,6 +76,13 @@ export class AvmMachineState {
     }
   }
 
+  /** Increases the gas left by the amounts specified. */
+  public refundGas(gasRefund: Partial<Gas>) {
+    for (const dimension of GasDimensions) {
+      this[`${dimension}Left`] += gasRefund[dimension] ?? 0;
+    }
+  }
+
   /**
    * Most instructions just increment PC before they complete
    */

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

@@ -30,6 +30,11 @@ export abstract class MemoryValue {
     return new Fr(this.toBigInt());
   }
 
+  // To number. Throws if exceeds max safe int.
+  public toNumber(): number {
+    return this.toFr().toNumber();
+  }
+
   public toString(): string {
     return `${this.constructor.name}(0x${this.toBigInt().toString(16)})`;
   }

yarn-project/simulator/src/avm/opcodes/arithmetic.ts

@@ -1,8 +1,13 @@
 import type { AvmContext } from '../avm_context.js';
-import { type GasCost, GasCostConstants, getGasCostMultiplierFromTypeTag, makeGasCost } from '../avm_gas_cost.js';
+import {
+  type Gas,
+  GasCostConstants,
+  getCostFromIndirectAccess,
+  getGasCostMultiplierFromTypeTag,
+  sumGas,
+} from '../avm_gas.js';
 import { type Field, type MemoryValue, TypeTag } from '../avm_memory_types.js';
 import { Opcode, OperandType } from '../serialization/instruction_serialization.js';
-import { Addressing, AddressingMode } from './addressing_mode.js';
 import { Instruction } from './instruction.js';
 import { ThreeOperandInstruction } from './instruction_impl.js';
 
@@ -19,15 +24,12 @@ export abstract class ThreeOperandArithmeticInstruction extends ThreeOperandInst
     context.machineState.incrementPc();
   }
 
-  protected gasCost(): GasCost {
-    const indirectCount = Addressing.fromWire(this.indirect).modePerOperand.filter(
-      mode => mode === AddressingMode.INDIRECT,
-    ).length;
-
-    const l2Gas =
-      indirectCount * GasCostConstants.ARITHMETIC_COST_PER_INDIRECT_ACCESS +
-      getGasCostMultiplierFromTypeTag(this.inTag) * GasCostConstants.ARITHMETIC_COST_PER_BYTE;
-    return makeGasCost({ l2Gas });
+  protected gasCost(): Gas {
+    const arithmeticCost = {
+      l2Gas: getGasCostMultiplierFromTypeTag(this.inTag) * GasCostConstants.ARITHMETIC_COST_PER_BYTE,
+    };
+    const indirectCost = getCostFromIndirectAccess(this.indirect);
+    return sumGas(arithmeticCost, indirectCost);
   }
 
   protected abstract compute(a: MemoryValue, b: MemoryValue): MemoryValue;