Per block gas accounting

assembly/api-debugger.ts

@@ -11,13 +11,21 @@ import { decodeSpi } from "./spi";
 
 let interpreter: Interpreter | null = null;
 
-export function resetJAM(program: u8[], pc: u32, initialGas: Gas, args: u8[], hasMetadata: boolean = false): void {
+export function resetJAM(
+  program: u8[],
+  pc: u32,
+  initialGas: Gas,
+  args: u8[],
+  hasMetadata: boolean = false,
+  useBlockGas: boolean = false,
+): void {
   const code = hasMetadata ? extractCodeAndMetadata(liftBytes(program)).code : liftBytes(program);
 
   const p = decodeSpi(code, liftBytes(args), 128);
   const int = new Interpreter(p.program, p.registers, p.memory);
   int.nextPc = <u32>pc;
   int.gas.set(initialGas);
+  int.useBlockGas = useBlockGas;
 
   if (interpreter !== null) {
     (<Interpreter>interpreter).memory.free();
@@ -26,14 +34,21 @@ export function resetJAM(program: u8[], pc: u32, initialGas: Gas, args: u8[], ha
   interpreter = int;
 }
 
-export function resetGeneric(program: u8[], flatRegisters: u8[], initialGas: Gas, hasMetadata: boolean = false): void {
+export function resetGeneric(
+  program: u8[],
+  flatRegisters: u8[],
+  initialGas: Gas,
+  hasMetadata: boolean = false,
+  useBlockGas: boolean = false,
+): void {
   const code = hasMetadata ? extractCodeAndMetadata(liftBytes(program)).code : liftBytes(program);
 
   const p = deblob(code);
   const registers: Registers = newRegisters();
   fillRegisters(registers, flatRegisters);
   const int = new Interpreter(p, registers);
   int.gas.set(initialGas);
+  int.useBlockGas = useBlockGas;
 
   if (interpreter !== null) {
     (<Interpreter>interpreter).memory.free();
@@ -49,6 +64,7 @@ export function resetGenericWithMemory(
   chunks: Uint8Array,
   initialGas: Gas,
   hasMetadata: boolean = false,
+  useBlockGas: boolean = false,
 ): void {
   const code = hasMetadata ? extractCodeAndMetadata(liftBytes(program)).code : liftBytes(program);
 
@@ -61,6 +77,7 @@ export function resetGenericWithMemory(
 
   const int = new Interpreter(p, registers, memory);
   int.gas.set(initialGas);
+  int.useBlockGas = useBlockGas;
 
   interpreter = int;
 }

assembly/api-internal.ts

@@ -71,17 +71,18 @@ export function buildMemory(builder: MemoryBuilder, pages: InitialPage[], chunks
 }
 
 /** Initialize new VM for execution. */
-export function vmInit(input: VmInput, useSbrkGas: boolean = false): Interpreter {
+export function vmInit(input: VmInput, useSbrkGas: boolean = false, useBlockGas: boolean = false): Interpreter {
   const int = new Interpreter(input.program, input.registers, input.memory);
   int.useSbrkGas = useSbrkGas;
+  int.useBlockGas = useBlockGas;
   int.nextPc = input.pc;
   int.gas.set(input.gas);
   return int;
 }
 
 /** Initialize & run & destroy a VM in a single go. */
 export function vmRunOnce(input: VmInput, options: VmRunOptions): VmOutput {
-  const int = vmInit(input, options.useSbrkGas);
+  const int = vmInit(input, options.useSbrkGas, options.useBlockGas);
   vmExecute(int, options.logs);
   return vmDestroy(int, options.dumpMemory);
 }

assembly/api-types.ts

@@ -19,6 +19,7 @@ export class InitialChunk {
 
 export class VmRunOptions {
   useSbrkGas: boolean = false;
+  useBlockGas: boolean = false;
   logs: boolean = false;
   dumpMemory: boolean = false;
 }

assembly/api-utils.ts

@@ -103,6 +103,7 @@ export function runProgram(
   logs: boolean = false,
   useSbrkGas: boolean = false,
   dumpMemory: boolean = false,
+  useBlockGas: boolean = false,
 ): VmOutput {
   const vmInput = new VmInput(program.program, program.memory, program.registers);
   vmInput.gas = i64(initialGas);
@@ -111,6 +112,7 @@ export function runProgram(
   const vmOptions = new VmRunOptions();
   vmOptions.logs = logs;
   vmOptions.useSbrkGas = useSbrkGas;
+  vmOptions.useBlockGas = useBlockGas;
   vmOptions.dumpMemory = dumpMemory;
 
   return vmRunOnce(vmInput, vmOptions);
@@ -126,11 +128,11 @@ const pvms = new Map<u32, Interpreter>();
  *
  * NOTE: the PVM MUST be de-allocated using `pvmDestroy`.
  */
-export function pvmStart(program: StandardProgram, useSbrkGas: boolean = false): u32 {
+export function pvmStart(program: StandardProgram, useSbrkGas: boolean = false, useBlockGas: boolean = false): u32 {
   const vmInput = new VmInput(program.program, program.memory, program.registers);
 
   nextPvmId += 1;
-  pvms.set(nextPvmId, vmInit(vmInput, useSbrkGas));
+  pvms.set(nextPvmId, vmInit(vmInput, useSbrkGas, useBlockGas));
   return nextPvmId;
 }
 

assembly/interpreter.ts

@@ -48,6 +48,7 @@ export class Interpreter {
   public exitCode: u32;
   public nextPc: u32;
   public useSbrkGas: boolean;
+  public useBlockGas: boolean;
 
   private djumpRes: DjumpResult = new DjumpResult();
   private argsRes: Args = new Args();
@@ -64,6 +65,7 @@ export class Interpreter {
     this.exitCode = 0;
     this.nextPc = 0;
     this.useSbrkGas = false;
+    this.useBlockGas = false;
   }
 
   nextSteps(nSteps: u32 = 1): boolean {
@@ -93,6 +95,8 @@ export class Interpreter {
     const mask = program.mask;
     const argsRes = this.argsRes;
     const outcomeRes = this.outcomeRes;
+    const useBlockGas = this.useBlockGas;
+    const blockGasCosts = program.blockGasCosts;
 
     for (let i: u32 = 0; i < nSteps; i++) {
       // reset some stuff at start
@@ -115,10 +119,20 @@ export class Interpreter {
       const instruction = unchecked(code[pc]);
       const iData = <i32>instruction < INSTRUCTIONS.length ? unchecked(INSTRUCTIONS[instruction]) : MISSING_INSTRUCTION;
 
-      // check gas (might be done for each block instead).
-      if (this.gas.sub(iData.gas)) {
-        this.status = Status.OOG;
-        return false;
+      // check gas: per-block or per-instruction
+      if (useBlockGas) {
+        const blockGas = portable.staticArrayAt(blockGasCosts, pc);
+        if (blockGas > 0) {
+          if (this.gas.sub(blockGas)) {
+            this.status = Status.OOG;
+            return false;
+          }
+        }
+      } else {
+        if (this.gas.sub(iData.gas)) {
+          this.status = Status.OOG;
+          return false;
+        }
       }
 
       if (iData === MISSING_INSTRUCTION) {

assembly/program.ts

@@ -229,12 +229,55 @@ export class JumpTable {
 }
 
 export class Program {
+  /**
+   * Pre-computed gas cost per basic block, indexed by PC.
+   * At block-start PCs, stores the total gas cost for the entire block; 0 elsewhere.
+   */
+  public readonly blockGasCosts: StaticArray<u64>;
+
   constructor(
     public readonly code: u8[],
     public readonly mask: Mask,
     public readonly jumpTable: JumpTable,
     public readonly basicBlocks: BasicBlocks,
-  ) {}
+  ) {
+    const len = code.length;
+    const blockGasCosts = new StaticArray<u64>(len);
+
+    // Two-pass approach: first accumulate gas per block, then store at block-start PCs
+    let blockStartPc: i32 = -1;
+    let blockGas: u64 = u64(0);
+
+    for (let i = 0; i < len; i++) {
+      if (!mask.isInstruction(i)) {
+        continue;
+      }
+
+      const instruction = code[i];
+      const iData = <i32>instruction < INSTRUCTIONS.length ? INSTRUCTIONS[instruction] : MISSING_INSTRUCTION;
+
+      if (basicBlocks.isStart(i)) {
+        // Store previous block's gas
+        if (blockStartPc >= 0) {
+          blockGasCosts[blockStartPc] = blockGas;
+        }
+        blockStartPc = i;
+        blockGas = iData.gas;
+      } else {
+        blockGas = portable.u64_add(blockGas, iData.gas);
+      }
+
+      // Skip argument bytes
+      i += mask.skipBytesToNextInstruction(i);
+    }
+
+    // Store the final block's gas
+    if (blockStartPc >= 0) {
+      blockGasCosts[blockStartPc] = blockGas;
+    }
+
+    this.blockGasCosts = blockGasCosts;
+  }
 
   toString(): string {
     return `Program { code: ${this.code}, mask: ${this.mask}, jumpTable: ${this.jumpTable}, basicBlocks: ${this.basicBlocks} }`;

bench/run.ts

@@ -28,6 +28,7 @@ const { values } = parseArgs({
     iterations: { type: "string", default: "5" },
     warmup: { type: "string", default: "1" },
     output: { type: "string", default: "" },
+    "block-gas": { type: "boolean", default: false },
     help: { type: "boolean", short: "h", default: false },
   },
 });
@@ -116,6 +117,9 @@ function benchRun(name: string, fn: () => void): BenchResult {
   const max = Math.max(...samples);
   const p95 = percentile(samples, 95);
 
+  return { name, medianMs: med, minMs: min, maxMs: max, p95Ms: p95, samples };
+}
+
 function formatResult(r: BenchResult): string {
   return `  ${r.name.padEnd(40)} median=${r.medianMs.toFixed(1)}ms  min=${r.minMs.toFixed(1)}ms  max=${r.maxMs.toFixed(1)}ms  p95=${r.p95Ms.toFixed(1)}ms`;
 }
@@ -145,6 +149,7 @@ function benchTraces(dir: string): BenchResult[] {
         hasMetadata: HasMetadata.Yes,
         verify: false,
         tracer: new NoOpTracer(),
+        useBlockGas: values["block-gas"],
       });
     });
 
@@ -216,7 +221,7 @@ function benchW3f(dir: string): BenchResult | null {
         [],
         16,
       );
-      runProgram(exe, gas, pc, false, false, false);
+      runProgram(exe, gas, pc, false, false, false, values["block-gas"]);
     }
   });
 

bin/src/trace-replay.ts

@@ -32,6 +32,7 @@ type ReplayOptions = {
   verify: boolean;
   logHostCall?: boolean;
   tracer?: Tracer;
+  useBlockGas?: boolean;
 };
 
 export function replayTraceFile(filePath: string, options: ReplayOptions): TraceSummary {
@@ -58,7 +59,7 @@ export function replayTraceFile(filePath: string, options: ReplayOptions): Trace
         128,
       );
 
-  const id = pvmStart(preparedProgram, true);
+  const id = pvmStart(preparedProgram, true, options.useBlockGas ?? false);
   const initialEcalliCount = ecalliEntries.length;
   const tracer = options.tracer ?? new ConsoleTracer();