feat(avm): Track gas usage based on memory accesses

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

@@ -6,16 +6,16 @@ import { encodeToBytecode } from './serialization/bytecode_serialization.js';
 
 describe('AVM simulator: dynamic gas costs per instruction', () => {
   it.each([
-    [new SetInstruction(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT8, /*value=*/ 1, /*dstOffset=*/ 0), [100, 0, 0]],
-    [new SetInstruction(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT32, /*value=*/ 1, /*dstOffset=*/ 0), [400, 0, 0]],
-    [new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ TypeTag.UINT8, /*copySize=*/ 1, /*dstOffset=*/ 0), [10, 0, 0]],
-    [new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ TypeTag.UINT8, /*copySize=*/ 5, /*dstOffset=*/ 0), [50, 0, 0]],
-    [new Add(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [10, 0, 0]],
-    [new Add(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT32, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [40, 0, 0]],
-    [new Add(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [20, 0, 0]],
-    [new Sub(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [20, 0, 0]],
-    [new Mul(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [20, 0, 0]],
-    [new Div(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [20, 0, 0]],
+    [new SetInstruction(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT8, /*value=*/ 1, /*dstOffset=*/ 0), [110, 0, 0]],
+    [new SetInstruction(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT32, /*value=*/ 1, /*dstOffset=*/ 0), [110]],
+    [new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ TypeTag.UINT8, /*copySize=*/ 1, /*dstOffset=*/ 0), [110]],
+    [new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ TypeTag.UINT8, /*copySize=*/ 5, /*dstOffset=*/ 0), [550]],
+    [new Add(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [130]],
+    [new Add(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT32, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [160]],
+    [new Add(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [150]],
+    [new Sub(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [150]],
+    [new Mul(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [150]],
+    [new Div(/*indirect=*/ 3, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [150]],
   ] as const)('computes gas cost for %s', async (instruction, [l2GasCost, l1GasCost, daGasCost]) => {
     const bytecode = encodeToBytecode([instruction]);
     const context = initContext();
@@ -27,8 +27,8 @@ describe('AVM simulator: dynamic gas costs per instruction', () => {
 
     await new AvmSimulator(context).executeBytecode(bytecode);
 
-    expect(initialL2GasLeft - context.machineState.l2GasLeft).toEqual(l2GasCost);
-    expect(initialL1GasLeft - context.machineState.l1GasLeft).toEqual(l1GasCost);
-    expect(initialDaGasLeft - context.machineState.daGasLeft).toEqual(daGasCost);
+    expect(initialL2GasLeft - context.machineState.l2GasLeft).toEqual(l2GasCost ?? 0);
+    expect(initialL1GasLeft - context.machineState.l1GasLeft).toEqual(l1GasCost ?? 0);
+    expect(initialDaGasLeft - context.machineState.daGasLeft).toEqual(daGasCost ?? 0);
   });
 });

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

@@ -21,7 +21,7 @@ export function gasLeftToGas(gasLeft: { l1GasLeft: number; l2GasLeft: number; da
 }
 
 /** Creates a new instance with all values set to zero except the ones set. */
-export function makeGasCost(gasCost: Partial<Gas>) {
+export function makeGas(gasCost: Partial<Gas>) {
   return { ...EmptyGas, ...gasCost };
 }
 
@@ -53,13 +53,13 @@ export const DynamicGasCost = Symbol('DynamicGasCost');
 /** Temporary default gas cost. We should eventually remove all usage of this variable in favor of actual gas for each opcode. */
 const TemporaryDefaultGasCost = { l1Gas: 0, l2Gas: 10, daGas: 0 };
 
-/** Gas costs for each instruction. */
+/** Base gas costs for each instruction. Additional gas cost may be added on top due to memory or storage accesses, etc. */
 export const GasCosts = {
   [Opcode.ADD]: DynamicGasCost,
   [Opcode.SUB]: DynamicGasCost,
   [Opcode.MUL]: DynamicGasCost,
   [Opcode.DIV]: DynamicGasCost,
-  [Opcode.FDIV]: TemporaryDefaultGasCost,
+  [Opcode.FDIV]: DynamicGasCost,
   [Opcode.EQ]: TemporaryDefaultGasCost,
   [Opcode.LT]: TemporaryDefaultGasCost,
   [Opcode.LTE]: TemporaryDefaultGasCost,
@@ -88,7 +88,7 @@ export const GasCosts = {
   [Opcode.BLOCKL1GASLIMIT]: TemporaryDefaultGasCost,
   [Opcode.BLOCKL2GASLIMIT]: TemporaryDefaultGasCost,
   [Opcode.BLOCKDAGASLIMIT]: TemporaryDefaultGasCost,
-  [Opcode.CALLDATACOPY]: DynamicGasCost,
+  [Opcode.CALLDATACOPY]: TemporaryDefaultGasCost,
   // Gas
   [Opcode.L1GASLEFT]: TemporaryDefaultGasCost,
   [Opcode.L2GASLEFT]: TemporaryDefaultGasCost,
@@ -99,7 +99,7 @@ export const GasCosts = {
   [Opcode.INTERNALCALL]: TemporaryDefaultGasCost,
   [Opcode.INTERNALRETURN]: TemporaryDefaultGasCost,
   // Memory
-  [Opcode.SET]: DynamicGasCost,
+  [Opcode.SET]: TemporaryDefaultGasCost,
   [Opcode.MOV]: TemporaryDefaultGasCost,
   [Opcode.CMOV]: TemporaryDefaultGasCost,
   // World state
@@ -126,33 +126,41 @@ 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 {
+/** Returns the fixed base gas cost for a given opcode, or throws if set to dynamic. */
+export function getBaseGasCost(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 };
+/** Returns the gas cost associated with the memory operations performed. */
+export function getMemoryGasCost(args: { reads?: number; writes?: number; indirectFlags?: number }) {
+  const { reads, writes, indirectFlags } = args;
+  const indirectCount = Addressing.fromWire(indirectFlags ?? 0).count(AddressingMode.INDIRECT);
+  const l2MemoryGasCost =
+    (reads ?? 0) * GasCostConstants.MEMORY_READ +
+    (writes ?? 0) * GasCostConstants.MEMORY_WRITE +
+    indirectCount * GasCostConstants.MEMORY_INDIRECT_READ_PENALTY;
+  return makeGas({ l2Gas: l2MemoryGasCost });
 }
 
 /** Constants used in base cost calculations. */
 export const GasCostConstants = {
-  SET_COST_PER_BYTE: 100,
-  CALLDATACOPY_COST_PER_BYTE: 10,
   ARITHMETIC_COST_PER_BYTE: 10,
-  COST_PER_INDIRECT_ACCESS: 5,
+  MEMORY_READ: 10,
+  MEMORY_INDIRECT_READ_PENALTY: 10,
+  MEMORY_WRITE: 100,
 };
 
+/** Returns gas cost for an operation on a given type tag based on the base cost per byte. */
+export function getGasCostForTypeTag(tag: TypeTag, baseCost: number) {
+  return baseCost * getGasCostMultiplierFromTypeTag(tag);
+}
+
 /** Returns a multiplier based on the size of the type represented by the tag. Throws on uninitialized or invalid. */
-export function getGasCostMultiplierFromTypeTag(tag: TypeTag) {
+function getGasCostMultiplierFromTypeTag(tag: TypeTag) {
   switch (tag) {
     case TypeTag.UINT8:
       return 1;

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

@@ -1,7 +1,7 @@
 import { Fr } from '@aztec/circuits.js';
 
 import { Gas, GasDimensions } from './avm_gas.js';
-import { TaggedMemory } from './avm_memory_types.js';
+import { MeteredTaggedMemory } from './avm_memory_types.js';
 import { AvmContractCallResults } from './avm_message_call_result.js';
 import { OutOfGasError } from './errors.js';
 
@@ -32,7 +32,7 @@ export class AvmMachineState {
   public internalCallStack: number[] = [];
 
   /** Memory accessible to user code */
-  public readonly memory: TaggedMemory = new TaggedMemory();
+  public readonly memory: MeteredTaggedMemory = new MeteredTaggedMemory();
 
   /**
    * Signals that execution should end.

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

@@ -1,4 +1,13 @@
-import { Field, TaggedMemory, Uint8, Uint16, Uint32, Uint64, Uint128 } from './avm_memory_types.js';
+import {
+  Field,
+  MeteredTaggedMemory,
+  TaggedMemory,
+  Uint8,
+  Uint16,
+  Uint32,
+  Uint64,
+  Uint128,
+} from './avm_memory_types.js';
 
 describe('TaggedMemory', () => {
   it('Elements should be undefined after construction', () => {
@@ -37,6 +46,48 @@ describe('TaggedMemory', () => {
   });
 });
 
+describe('MeteredTaggedMemory', () => {
+  let mem: MeteredTaggedMemory;
+
+  beforeEach(() => {
+    mem = new MeteredTaggedMemory();
+  });
+
+  it(`Counts reads`, () => {
+    mem.get(10);
+    mem.getAs(20);
+    expect(mem.getStats()).toEqual({ reads: 2, writes: 0 });
+  });
+
+  it(`Counts reading slices`, () => {
+    const val = [new Field(5), new Field(6), new Field(7)];
+    mem.setSlice(10, val);
+    mem.clearStats();
+
+    mem.getSlice(10, 3);
+    mem.getSliceAs(11, 2);
+    expect(mem.getStats()).toEqual({ reads: 5, writes: 0 });
+  });
+
+  it(`Counts writes`, () => {
+    mem.set(10, new Uint8(5));
+    expect(mem.getStats()).toEqual({ reads: 0, writes: 1 });
+  });
+
+  it(`Counts writing slices`, () => {
+    mem.setSlice(10, [new Field(5), new Field(6)]);
+    expect(mem.getStats()).toEqual({ reads: 0, writes: 2 });
+  });
+
+  it(`Clears stats`, () => {
+    mem.get(10);
+    mem.set(20, new Uint8(5));
+    expect(mem.getStats()).toEqual({ reads: 1, writes: 1 });
+    mem.clearStats();
+    expect(mem.getStats()).toEqual({ reads: 0, writes: 0 });
+  });
+});
+
 type IntegralClass = typeof Uint8 | typeof Uint16 | typeof Uint32 | typeof Uint64 | typeof Uint128;
 describe.each([Uint8, Uint16, Uint32, Uint64, Uint128])('Integral Types', (clsValue: IntegralClass) => {
   describe(`${clsValue.name}`, () => {

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

@@ -4,7 +4,8 @@ import { DebugLogger, createDebugLogger } from '@aztec/foundation/log';
 
 import { strict as assert } from 'assert';
 
-import { TagCheckError } from './errors.js';
+import { InstructionExecutionError, TagCheckError } from './errors.js';
+import { Addressing, AddressingMode } from './opcodes/addressing_mode.js';
 
 /** MemoryValue gathers the common operations for all memory types. */
 export abstract class MemoryValue {
@@ -373,3 +374,65 @@ export class TaggedMemory {
     }
   }
 }
+
+/** Tagged memory with metering for each memory read and write. */
+export class MeteredTaggedMemory extends TaggedMemory {
+  private reads: number = 0;
+  private writes: number = 0;
+
+  public getStats(): MemoryOperations {
+    return { reads: this.reads, writes: this.writes };
+  }
+
+  public clearStats(): MemoryOperations {
+    const stats = this.getStats();
+    this.reads = 0;
+    this.writes = 0;
+    return stats;
+  }
+
+  public assertStats(operations: MemoryOperations & { indirectFlags: number }, type = 'instruction') {
+    const { reads: expectedReads, writes: expectedWrites, indirectFlags } = operations;
+
+    const totalExpectedReads = expectedReads + Addressing.fromWire(indirectFlags).count(AddressingMode.INDIRECT);
+    const { reads: actualReads, writes: actualWrites } = this.clearStats();
+    if (actualReads !== totalExpectedReads) {
+      throw new InstructionExecutionError(
+        `Incorrect number of memory reads for ${type}: expected ${totalExpectedReads} but executed ${actualReads}`,
+      );
+    }
+    if (actualWrites !== expectedWrites) {
+      throw new InstructionExecutionError(
+        `Incorrect number of memory writes for ${type}: expected ${expectedWrites} but executed ${actualWrites}`,
+      );
+    }
+  }
+
+  public getAs<T>(offset: number): T {
+    this.reads++;
+    return super.getAs(offset);
+  }
+
+  public getSlice(offset: number, size: number): MemoryValue[] {
+    this.reads += size;
+    return super.getSlice(offset, size);
+  }
+
+  public set(offset: number, v: MemoryValue): void {
+    this.writes++;
+    super.set(offset, v);
+  }
+
+  public setSlice(offset: number, vs: MemoryValue[]): void {
+    this.writes += vs.length;
+    super.setSlice(offset, vs);
+  }
+}
+
+/** Tracks number of memory reads and writes. */
+export type MemoryOperations = {
+  /** How many total reads are performed. Slice reads are count as one per element. */
+  reads: number;
+  /** How many total writes are performed. Slice writes are count as one per element. */
+  writes: number;
+};

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

@@ -44,7 +44,7 @@ describe('AVM simulator: injected bytecode', () => {
 
     expect(results.reverted).toBe(false);
     expect(results.output).toEqual([new Fr(3)]);
-    expect(context.machineState.l2GasLeft).toEqual(initialL2GasLeft - 350);
+    expect(context.machineState.l2GasLeft).toEqual(initialL2GasLeft - 680);
   });
 
   it('Should halt if runs out of gas', async () => {

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

@@ -42,6 +42,10 @@ export class NoteHashExists extends FixedGasInstruction {
 
     context.machineState.incrementPc();
   }
+
+  protected memoryOperations() {
+    return { reads: 2, writes: 1 };
+  }
 }
 
 export class EmitNoteHash extends FixedGasInstruction {
@@ -64,6 +68,10 @@ export class EmitNoteHash extends FixedGasInstruction {
 
     context.machineState.incrementPc();
   }
+
+  protected memoryOperations() {
+    return { reads: 1 };
+  }
 }
 
 export class NullifierExists extends FixedGasInstruction {
@@ -84,6 +92,10 @@ export class NullifierExists extends FixedGasInstruction {
 
     context.machineState.incrementPc();
   }
+
+  protected memoryOperations() {
+    return { reads: 1, writes: 1 };
+  }
 }
 
 export class EmitNullifier extends FixedGasInstruction {
@@ -117,6 +129,10 @@ export class EmitNullifier extends FixedGasInstruction {
 
     context.machineState.incrementPc();
   }
+
+  protected memoryOperations() {
+    return { reads: 1 };
+  }
 }
 
 export class L1ToL2MessageExists extends FixedGasInstruction {
@@ -148,6 +164,10 @@ export class L1ToL2MessageExists extends FixedGasInstruction {
 
     context.machineState.incrementPc();
   }
+
+  protected memoryOperations() {
+    return { reads: 2, writes: 1 };
+  }
 }
 
 export class EmitUnencryptedLog extends FixedGasInstruction {
@@ -188,6 +208,10 @@ export class EmitUnencryptedLog extends FixedGasInstruction {
 
     context.machineState.incrementPc();
   }
+
+  protected memoryOperations() {
+    return { reads: 1 + this.logSize };
+  }
 }
 
 export class SendL2ToL1Message extends FixedGasInstruction {
@@ -211,4 +235,8 @@ export class SendL2ToL1Message extends FixedGasInstruction {
 
     context.machineState.incrementPc();
   }
+
+  protected memoryOperations() {
+    return { reads: 2 };
+  }
 }

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

@@ -12,7 +12,7 @@ export enum AddressingMode {
 export class Addressing {
   public constructor(
     /** The addressing mode for each operand. The length of this array is the number of operands of the instruction. */
-    public readonly modePerOperand: AddressingMode[],
+    private readonly modePerOperand: AddressingMode[],
   ) {
     assert(modePerOperand.length <= 8, 'At most 8 operands are supported');
   }
@@ -39,6 +39,11 @@ export class Addressing {
     return wire;
   }
 
+  /** Returns how many operands use the given addressing mode. */
+  public count(mode: AddressingMode): number {
+    return this.modePerOperand.filter(m => m === mode).length;
+  }
+
   /**
    * Resolves the offsets using the addressing mode.
    * @param offsets The offsets to resolve.