feat(avm): Track gas from memory accesses explicitly (#5563)

Tracks gas usage for all AVM instructions based on memory consumption.
Adds an optional wrapper for TaggedMemory (enabled on test only) that
tracks all memory reads and writes to validate that the number of memory
operations charged match the actual ones.

Replaces existing #5514 and #5518 in favor of a more explicit approach,
at the expense of more duplicated code in each instruction, but
flattening the instruction hierarchy.

Closes #5518 
Closes #5514

yarn-project/foundation/src/types/index.ts

@@ -4,5 +4,11 @@ export type FieldsOf<T> = {
   [P in keyof T as T[P] extends Function ? never : P]: T[P];
 };
 
+/** Extracts methods of a type. */
+export type FunctionsOf<T> = {
+  // eslint-disable-next-line @typescript-eslint/ban-types
+  [P in keyof T as T[P] extends Function ? P : never]: T[P];
+};
+
 /** Marks a set of properties of a type as optional. */
 export type PartialBy<T, K extends keyof T> = Omit<T, K> & Partial<Pick<T, K>>;

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

@@ -6,16 +6,23 @@ 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]],
+    // BASE_GAS(10) * 1 + MEMORY_WRITE(100) = 110
+    [new SetInstruction(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT8, /*value=*/ 1, /*dstOffset=*/ 0), [110, 0, 0]],
+    // BASE_GAS(10) * 1 + MEMORY_WRITE(100) = 110
+    [new SetInstruction(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT32, /*value=*/ 1, /*dstOffset=*/ 0), [110]],
+    // BASE_GAS(10) * 1 + MEMORY_WRITE(100) = 110
+    [new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ TypeTag.UINT8, /*copySize=*/ 1, /*dstOffset=*/ 0), [110]],
+    // BASE_GAS(10) * 5 + MEMORY_WRITE(100) * 5 = 550
+    [new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ TypeTag.UINT8, /*copySize=*/ 5, /*dstOffset=*/ 0), [550]],
+    // BASE_GAS(10) * 1 + MEMORY_READ(10) * 2 + MEMORY_WRITE(100) = 130
+    [new Add(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT8, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [130]],
+    // BASE_GAS(10) * 4 + MEMORY_READ(10) * 2 + MEMORY_WRITE(100) = 160
+    [new Add(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT32, /*aOffset=*/ 1, /*bOffset=*/ 2, /*dstOffset=*/ 3), [160]],
+    // BASE_GAS(10) * 1 + MEMORY_READ(10) * 2 + MEMORY_INDIRECT_READ_PENALTY(10) * 2 + MEMORY_WRITE(100) = 150
+    [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 +34,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

@@ -1,4 +1,5 @@
 import { TypeTag } from './avm_memory_types.js';
+import { InstructionExecutionError } from './errors.js';
 import { Addressing, AddressingMode } from './opcodes/addressing_mode.js';
 import { Opcode } from './serialization/instruction_serialization.js';
 
@@ -20,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 };
 }
 
@@ -36,6 +37,11 @@ export function sumGas(...gases: Partial<Gas>[]) {
   );
 }
 
+/** Multiplies a gas instance by a scalar. */
+export function mulGas(gas: Partial<Gas>, scalar: number) {
+  return { l1Gas: (gas.l1Gas ?? 0) * scalar, l2Gas: (gas.l2Gas ?? 0) * scalar, daGas: (gas.daGas ?? 0) * scalar };
+}
+
 /** Zero gas across all gas dimensions. */
 export const EmptyGas: Gas = {
   l1Gas: 0,
@@ -52,12 +58,12 @@ 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. */
-export const GasCosts = {
-  [Opcode.ADD]: DynamicGasCost,
-  [Opcode.SUB]: DynamicGasCost,
-  [Opcode.MUL]: DynamicGasCost,
-  [Opcode.DIV]: DynamicGasCost,
+/** Base gas costs for each instruction. Additional gas cost may be added on top due to memory or storage accesses, etc. */
+export const GasCosts: Record<Opcode, Gas | typeof DynamicGasCost> = {
+  [Opcode.ADD]: TemporaryDefaultGasCost,
+  [Opcode.SUB]: TemporaryDefaultGasCost,
+  [Opcode.MUL]: TemporaryDefaultGasCost,
+  [Opcode.DIV]: TemporaryDefaultGasCost,
   [Opcode.FDIV]: TemporaryDefaultGasCost,
   [Opcode.EQ]: TemporaryDefaultGasCost,
   [Opcode.LT]: TemporaryDefaultGasCost,
@@ -87,7 +93,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,
@@ -98,7 +104,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
@@ -124,35 +130,42 @@ export const GasCosts = {
   [Opcode.POSEIDON]: TemporaryDefaultGasCost,
   [Opcode.SHA256]: TemporaryDefaultGasCost, // temp - may be removed, but alot of contracts rely on i: TemporaryDefaultGasCost,
   [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; indirect?: number }) {
+  const { reads, writes, indirect } = args;
+  const indirectCount = Addressing.fromWire(indirect ?? 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: Gas) {
+  return mulGas(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;
@@ -168,6 +181,6 @@ export function getGasCostMultiplierFromTypeTag(tag: TypeTag) {
       return 32;
     case TypeTag.INVALID:
     case TypeTag.UNINITIALIZED:
-      throw new Error(`Invalid tag type for gas cost multiplier: ${TypeTag[tag]}`);
+      throw new InstructionExecutionError(`Invalid tag type for gas cost multiplier: ${TypeTag[tag]}`);
   }
 }

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,58 @@ describe('TaggedMemory', () => {
   });
 });
 
+describe('MeteredTaggedMemory', () => {
+  let mem: MeteredTaggedMemory;
+
+  beforeEach(() => {
+    mem = new MeteredTaggedMemory(new TaggedMemory());
+  });
+
+  it(`Counts reads`, () => {
+    mem.get(10);
+    mem.getAs(20);
+    expect(mem.reset()).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.reset();
+
+    mem.getSlice(10, 3);
+    mem.getSliceAs(11, 2);
+    expect(mem.reset()).toEqual({ reads: 5, writes: 0 });
+  });
+
+  it(`Counts writes`, () => {
+    mem.set(10, new Uint8(5));
+    expect(mem.reset()).toEqual({ reads: 0, writes: 1 });
+  });
+
+  it(`Counts writing slices`, () => {
+    mem.setSlice(10, [new Field(5), new Field(6)]);
+    expect(mem.reset()).toEqual({ reads: 0, writes: 2 });
+  });
+
+  it(`Clears stats`, () => {
+    mem.get(10);
+    mem.set(20, new Uint8(5));
+    expect(mem.reset()).toEqual({ reads: 1, writes: 1 });
+    expect(mem.reset()).toEqual({ reads: 0, writes: 0 });
+  });
+
+  it(`Asserts stats`, () => {
+    mem.get(10);
+    mem.set(20, new Uint8(5));
+    expect(() => mem.assert({ reads: 1, writes: 1 })).not.toThrow();
+  });
+
+  it(`Throws on failed stat assertion`, () => {
+    mem.get(10);
+    expect(() => mem.assert({ reads: 1, writes: 1 })).toThrow();
+  });
+});
+
 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}`, () => {