feat(avm): Gas usage for nested calls

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

@@ -315,7 +315,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]);
@@ -348,7 +348,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]);
 
@@ -359,7 +359,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 };
+    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 { AztecAddress, FunctionSelector } from '@aztec/circuits.js';
 import { Fr } from '@aztec/foundation/fields';
 
 import { AvmExecutionEnvironment } from './avm_execution_environment.js';
+import { Gas, gasToGasLeft } from './avm_gas.js';
 import { AvmMachineState } from './avm_machine_state.js';
 import { 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,40 @@
 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 };
+}
+
+/** Adds multiple instances of Gas. */
+export function addGas(...gases: Partial<Gas>[]) {
+  return {
+    l1Gas: gases.reduce((acc, gas) => acc + (gas.l1Gas ?? 0), 0),
+    l2Gas: gases.reduce((acc, gas) => acc + (gas.l2Gas ?? 0), 0),
+    daGas: gases.reduce((acc, gas) => acc + (gas.daGas ?? 0), 0),
+  };
 }
 
-/** Gas cost of zero across all gas dimensions. */
-export const EmptyGasCost = {
+/** Zero gas across all gas dimensions. */
+export const EmptyGas = {
   l1Gas: 0,
   l2Gas: 0,
   daGas: 0,
@@ -102,12 +122,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 { Fr } from '@aztec/circuits.js';
 
-import { GasCost, GasDimensions } from './avm_gas_cost.js';
+import { 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 { GasCost, GasCostConstants, getGasCostMultiplierFromTypeTag, makeGasCost } from '../avm_gas_cost.js';
+import {
+  Gas,
+  GasCostConstants,
+  addGas,
+  getCostFromIndirectAccess,
+  getGasCostMultiplierFromTypeTag,
+} from '../avm_gas.js';
 import { Field, 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 addGas(arithmeticCost, indirectCost);
   }
 
   protected abstract compute(a: MemoryValue, b: MemoryValue): MemoryValue;