feat(avm): Track gas from memory accesses explicitly

Author: spalladino

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
@@ -123,35 +129,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;
@@ -167,6 +180,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}`, () => {