feat: Return gas usage per phase from node tx simulation (#6255)

Allows clients set their total and teardown gas limits accordingly
before submitting a tx to the network.

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -662,6 +662,7 @@ export class AztecNodeService implements AztecNode {
     const processor = await publicProcessorFactory.create(prevHeader, newGlobalVariables);
     // REFACTOR: Consider merging ProcessReturnValues into ProcessedTx
     const [processedTxs, failedTxs, returns] = await processor.process([tx]);
+    // REFACTOR: Consider returning the error/revert rather than throwing
     if (failedTxs.length) {
       this.log.warn(`Simulated tx ${tx.getTxHash()} fails: ${failedTxs[0].error}`);
       throw failedTxs[0].error;
@@ -680,6 +681,7 @@ export class AztecNodeService implements AztecNode {
       end: processedTx.data.end,
       revertReason: processedTx.revertReason,
       publicReturnValues: returns[0],
+      gasUsed: processedTx.gasUsed,
     };
   }
 

yarn-project/circuit-types/src/mocks.ts

@@ -136,6 +136,7 @@ export const mockSimulatedTx = (seed = 1, hasLogs = true) => {
     end: makeCombinedAccumulatedData(),
     revertReason: undefined,
     publicReturnValues: dec,
+    gasUsed: {},
   };
   return new SimulatedTx(tx, dec, output);
 };

yarn-project/circuit-types/src/tx/processed_tx.ts

@@ -9,6 +9,7 @@ import {
 } from '@aztec/circuit-types';
 import {
   Fr,
+  type Gas,
   type Header,
   KernelCircuitPublicInputs,
   type Proof,
@@ -68,6 +69,11 @@ export type ProcessedTx = Pick<Tx, 'proof' | 'encryptedLogs' | 'unencryptedLogs'
    * The collection of public kernel circuit inputs for simulation/proving
    */
   publicKernelRequests: PublicKernelRequest[];
+  /**
+   * Gas usage per public execution phase.
+   * Doesn't account for any base costs nor DA gas used in private execution.
+   */
+  gasUsed: Partial<Record<PublicKernelType, Gas>>;
 };
 
 export type RevertedTx = ProcessedTx & {
@@ -122,6 +128,7 @@ export function makeProcessedTx(
   proof: Proof,
   publicKernelRequests: PublicKernelRequest[],
   revertReason?: SimulationError,
+  gasUsed: ProcessedTx['gasUsed'] = {},
 ): ProcessedTx {
   return {
     hash: tx.getTxHash(),
@@ -132,6 +139,7 @@ export function makeProcessedTx(
     isEmpty: false,
     revertReason,
     publicKernelRequests,
+    gasUsed,
   };
 }
 
@@ -156,6 +164,7 @@ export function makeEmptyProcessedTx(header: Header, chainId: Fr, version: Fr):
     isEmpty: true,
     revertReason: undefined,
     publicKernelRequests: [],
+    gasUsed: {},
   };
 }
 

yarn-project/circuit-types/src/tx/simulated_tx.test.ts

@@ -1,16 +1,59 @@
+import { Gas } from '@aztec/circuits.js';
+
 import { mockSimulatedTx } from '../mocks.js';
+import { PublicKernelType } from './processed_tx.js';
 import { SimulatedTx } from './simulated_tx.js';
 
 describe('simulated_tx', () => {
-  it('convert to and from json', () => {
-    const simulatedTx = mockSimulatedTx();
-    expect(SimulatedTx.fromJSON(simulatedTx.toJSON())).toEqual(simulatedTx);
+  let simulatedTx: SimulatedTx;
+
+  beforeEach(() => {
+    simulatedTx = mockSimulatedTx();
+  });
+
+  describe('json', () => {
+    it('convert to and from json', () => {
+      expect(SimulatedTx.fromJSON(simulatedTx.toJSON())).toEqual(simulatedTx);
+    });
+
+    it('convert undefined effects to and from json', () => {
+      simulatedTx.privateReturnValues = undefined;
+      simulatedTx.publicOutput = undefined;
+      expect(SimulatedTx.fromJSON(simulatedTx.toJSON())).toEqual(simulatedTx);
+    });
   });
 
-  it('convert undefined effects to and from json', () => {
-    const simulatedTx = mockSimulatedTx();
-    simulatedTx.privateReturnValues = undefined;
-    simulatedTx.publicOutput = undefined;
-    expect(SimulatedTx.fromJSON(simulatedTx.toJSON())).toEqual(simulatedTx);
+  describe('getGasLimits', () => {
+    beforeEach(() => {
+      simulatedTx.tx.data.publicInputs.end.gasUsed = Gas.from({ daGas: 100, l2Gas: 200 });
+      simulatedTx.publicOutput!.gasUsed = {
+        [PublicKernelType.SETUP]: Gas.from({ daGas: 10, l2Gas: 20 }),
+        [PublicKernelType.APP_LOGIC]: Gas.from({ daGas: 20, l2Gas: 40 }),
+        [PublicKernelType.TEARDOWN]: Gas.from({ daGas: 10, l2Gas: 20 }),
+      };
+    });
+
+    it('returns gas limits from private gas usage only', () => {
+      simulatedTx.publicOutput = undefined;
+      // Should be 110 and 220 but oh floating point
+      expect(simulatedTx.getGasLimits()).toEqual({
+        totalGas: Gas.from({ daGas: 111, l2Gas: 221 }),
+        teardownGas: Gas.empty(),
+      });
+    });
+
+    it('returns gas limits for private and public', () => {
+      expect(simulatedTx.getGasLimits()).toEqual({
+        totalGas: Gas.from({ daGas: 154, l2Gas: 308 }),
+        teardownGas: Gas.from({ daGas: 11, l2Gas: 22 }),
+      });
+    });
+
+    it('pads gas limits', () => {
+      expect(simulatedTx.getGasLimits(1)).toEqual({
+        totalGas: Gas.from({ daGas: 280, l2Gas: 560 }),
+        teardownGas: Gas.from({ daGas: 20, l2Gas: 40 }),
+      });
+    });
   });
 });

yarn-project/circuit-types/src/tx/simulated_tx.ts

@@ -1,7 +1,8 @@
-import { CombinedAccumulatedData, CombinedConstantData, Fr } from '@aztec/circuits.js';
+import { CombinedAccumulatedData, CombinedConstantData, Fr, Gas } from '@aztec/circuits.js';
+import { mapValues } from '@aztec/foundation/collection';
 
 import { EncryptedTxL2Logs, UnencryptedTxL2Logs } from '../logs/index.js';
-import { type ProcessedTx } from './processed_tx.js';
+import { type ProcessedTx, PublicKernelType } from './processed_tx.js';
 import { Tx } from './tx.js';
 
 /** Return values of simulating a circuit. */
@@ -11,7 +12,7 @@ export type ProcessReturnValues = Fr[] | undefined;
  * Outputs of processing the public component of a transaction.
  * REFACTOR: Rename.
  */
-export type ProcessOutput = Pick<ProcessedTx, 'encryptedLogs' | 'unencryptedLogs' | 'revertReason'> &
+export type ProcessOutput = Pick<ProcessedTx, 'encryptedLogs' | 'unencryptedLogs' | 'revertReason' | 'gasUsed'> &
   Pick<ProcessedTx['data'], 'constants' | 'end'> & { publicReturnValues: ProcessReturnValues };
 
 function processOutputToJSON(output: ProcessOutput) {
@@ -22,6 +23,7 @@ function processOutputToJSON(output: ProcessOutput) {
     constants: output.constants.toBuffer().toString('hex'),
     end: output.end.toBuffer().toString('hex'),
     publicReturnValues: output.publicReturnValues?.map(fr => fr.toString()),
+    gasUsed: mapValues(output.gasUsed, gas => gas?.toJSON()),
   };
 }
 
@@ -33,6 +35,7 @@ function processOutputFromJSON(json: any): ProcessOutput {
     constants: CombinedConstantData.fromBuffer(Buffer.from(json.constants, 'hex')),
     end: CombinedAccumulatedData.fromBuffer(Buffer.from(json.end, 'hex')),
     publicReturnValues: json.publicReturnValues?.map(Fr.fromString),
+    gasUsed: mapValues(json.gasUsed, gas => (gas ? Gas.fromJSON(gas) : undefined)),
   };
 }
 
@@ -45,6 +48,29 @@ function processOutputFromJSON(json: any): ProcessOutput {
 export class SimulatedTx {
   constructor(public tx: Tx, public privateReturnValues?: ProcessReturnValues, public publicOutput?: ProcessOutput) {}
 
+  /**
+   * Returns suggested total and teardown gas limits for the simulated tx.
+   * Note that public gas usage is only accounted for if the publicOutput is present.
+   * @param pad - Percentage to pad the suggested gas limits by, defaults to 10%.
+   */
+  public getGasLimits(pad = 0.1) {
+    const privateGasUsed = this.tx.data.publicInputs.end.gasUsed;
+    if (this.publicOutput) {
+      const publicGasUsed = Object.values(this.publicOutput.gasUsed).reduce(
+        (total, current) => total.add(current),
+        Gas.empty(),
+      );
+      const teardownGas = this.publicOutput.gasUsed[PublicKernelType.TEARDOWN] ?? Gas.empty();
+
+      return {
+        totalGas: privateGasUsed.add(publicGasUsed).mul(1 + pad),
+        teardownGas: teardownGas.mul(1 + pad),
+      };
+    }
+
+    return { totalGas: privateGasUsed.mul(1 + pad), teardownGas: Gas.empty() };
+  }
+
   /**
    * Convert a SimulatedTx class object to a plain JSON object.
    * @returns A plain object with SimulatedTx properties.

yarn-project/circuits.js/src/structs/gas.ts

@@ -83,4 +83,12 @@ export class Gas {
     const reader = FieldReader.asReader(fields);
     return new Gas(reader.readU32(), reader.readU32());
   }
+
+  toJSON() {
+    return { daGas: this.daGas, l2Gas: this.l2Gas };
+  }
+
+  static fromJSON(json: any) {
+    return new Gas(json.daGas, json.l2Gas);
+  }
 }

yarn-project/circuits.js/src/structs/kernel/private_kernel_tail_circuit_public_inputs.ts

@@ -124,6 +124,10 @@ export class PrivateKernelTailCircuitPublicInputs {
     }
   }
 
+  get publicInputs(): PartialPrivateTailPublicInputsForPublic | PartialPrivateTailPublicInputsForRollup {
+    return (this.forPublic ?? this.forRollup)!;
+  }
+
   toPublicKernelCircuitPublicInputs() {
     if (!this.forPublic) {
       throw new Error('Private tail public inputs is not for public circuit.');

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

@@ -1 +1,2 @@
 export * from './array.js';
+export * from './object.js';

yarn-project/foundation/src/collection/object.test.ts

@@ -0,0 +1,30 @@
+import { mapValues } from './object.js';
+
+describe('mapValues', () => {
+  it('should return a new object with mapped values', () => {
+    const obj = { a: 1, b: 2, c: 3 };
+    const fn = (value: number) => value * 2;
+
+    const result = mapValues(obj, fn);
+
+    expect(result).toEqual({ a: 2, b: 4, c: 6 });
+  });
+
+  it('should handle an empty object', () => {
+    const obj = {};
+    const fn = (value: number) => value * 2;
+
+    const result = mapValues(obj, fn);
+
+    expect(result).toEqual({});
+  });
+
+  it('should handle different value types', () => {
+    const obj = { a: 'hello', b: true, c: [1, 2, 3] };
+    const fn = (value: any) => typeof value;
+
+    const result = mapValues(obj, fn);
+
+    expect(result).toEqual({ a: 'string', b: 'boolean', c: 'object' });
+  });
+});

yarn-project/foundation/src/collection/object.ts

@@ -0,0 +1,19 @@
+/** Returns a new object with the same keys and where each value has been passed through the mapping function. */
+export function mapValues<K extends string | number | symbol, T, U>(
+  obj: Record<K, T>,
+  fn: (value: T) => U,
+): Record<K, U>;
+export function mapValues<K extends string | number | symbol, T, U>(
+  obj: Partial<Record<K, T>>,
+  fn: (value: T) => U,
+): Partial<Record<K, U>>;
+export function mapValues<K extends string | number | symbol, T, U>(
+  obj: Record<K, T>,
+  fn: (value: T) => U,
+): Record<K, U> {
+  const result: Record<K, U> = {} as Record<K, U>;
+  for (const key in obj) {
+    result[key] = fn(obj[key]);
+  }
+  return result;
+}

yarn-project/simulator/src/public/abstract_phase_manager.ts

@@ -2,6 +2,7 @@ import {
   MerkleTreeId,
   type ProcessReturnValues,
   type PublicKernelRequest,
+  PublicKernelType,
   type SimulationError,
   type Tx,
   type UnencryptedFunctionL2Logs,
@@ -81,6 +82,20 @@ export const PhaseIsRevertible: Record<PublicKernelPhase, boolean> = {
   [PublicKernelPhase.TAIL]: false,
 };
 
+// REFACTOR: Unify both enums and move to types or circuit-types.
+export function publicKernelPhaseToKernelType(phase: PublicKernelPhase): PublicKernelType {
+  switch (phase) {
+    case PublicKernelPhase.SETUP:
+      return PublicKernelType.SETUP;
+    case PublicKernelPhase.APP_LOGIC:
+      return PublicKernelType.APP_LOGIC;
+    case PublicKernelPhase.TEARDOWN:
+      return PublicKernelType.TEARDOWN;
+    case PublicKernelPhase.TAIL:
+      return PublicKernelType.TAIL;
+  }
+}
+
 export abstract class AbstractPhaseManager {
   protected hintsBuilder: HintsBuilder;
   protected log: DebugLogger;
@@ -127,6 +142,8 @@ export abstract class AbstractPhaseManager {
      */
     revertReason: SimulationError | undefined;
     returnValues: ProcessReturnValues;
+    /** Gas used during the execution this particular phase. */
+    gasUsed: Gas | undefined;
   }>;
 
   public static extractEnqueuedPublicCallsByPhase(
@@ -202,6 +219,7 @@ export abstract class AbstractPhaseManager {
     return calls;
   }
 
+  // REFACTOR: Do not return an array and instead return a struct with similar shape to that returned by `handle`
   protected async processEnqueuedPublicCalls(
     tx: Tx,
     previousPublicKernelOutput: PublicKernelCircuitPublicInputs,
@@ -214,6 +232,7 @@ export abstract class AbstractPhaseManager {
       UnencryptedFunctionL2Logs[],
       SimulationError | undefined,
       ProcessReturnValues,
+      Gas,
     ]
   > {
     let kernelOutput = previousPublicKernelOutput;
@@ -223,7 +242,7 @@ export abstract class AbstractPhaseManager {
     const enqueuedCalls = this.extractEnqueuedPublicCalls(tx);
 
     if (!enqueuedCalls || !enqueuedCalls.length) {
-      return [[], kernelOutput, kernelProof, [], undefined, undefined];
+      return [[], kernelOutput, kernelProof, [], undefined, undefined, Gas.empty()];
     }
 
     const newUnencryptedFunctionLogs: UnencryptedFunctionL2Logs[] = [];
@@ -236,6 +255,7 @@ export abstract class AbstractPhaseManager {
     // and submitted separately to the base rollup?
 
     let returns: ProcessReturnValues = undefined;
+    let gasUsed = Gas.empty();
 
     for (const enqueuedCall of enqueuedCalls) {
       const executionStack: (PublicExecution | PublicExecutionResult)[] = [enqueuedCall];
@@ -263,6 +283,9 @@ export abstract class AbstractPhaseManager {
             )
           : current;
 
+        // Accumulate gas used in this execution
+        gasUsed = gasUsed.add(Gas.from(result.startGasLeft).sub(Gas.from(result.endGasLeft)));
+
         const functionSelector = result.execution.functionData.selector.toString();
         if (result.reverted && !PhaseIsRevertible[this.phase]) {
           this.log.debug(
@@ -306,7 +329,8 @@ export abstract class AbstractPhaseManager {
               result.revertReason
             }`,
           );
-          return [[], kernelOutput, kernelProof, [], result.revertReason, undefined];
+          // TODO(@spalladino): Check gasUsed is correct. The AVM should take care of setting gasLeft to zero upon a revert.
+          return [[], kernelOutput, kernelProof, [], result.revertReason, undefined, gasUsed];
         }
 
         if (!enqueuedExecutionResult) {
@@ -322,7 +346,7 @@ export abstract class AbstractPhaseManager {
     // TODO(#3675): This should be done in a public kernel circuit
     removeRedundantPublicDataWrites(kernelOutput, this.phase);
 
-    return [publicKernelInputs, kernelOutput, kernelProof, newUnencryptedFunctionLogs, undefined, returns];
+    return [publicKernelInputs, kernelOutput, kernelProof, newUnencryptedFunctionLogs, undefined, returns, gasUsed];
   }
 
   /** Returns all pending private and public nullifiers.  */

yarn-project/simulator/src/public/app_logic_phase_manager.ts

@@ -47,6 +47,7 @@ export class AppLogicPhaseManager extends AbstractPhaseManager {
       newUnencryptedFunctionLogs,
       revertReason,
       returnValues,
+      gasUsed,
     ] = await this.processEnqueuedPublicCalls(tx, previousPublicKernelOutput, previousPublicKernelProof).catch(
       // if we throw for any reason other than simulation, we need to rollback and drop the TX
       async err => {
@@ -71,6 +72,6 @@ export class AppLogicPhaseManager extends AbstractPhaseManager {
       };
       return request;
     });
-    return { kernelRequests, publicKernelOutput, publicKernelProof, revertReason, returnValues };
+    return { kernelRequests, publicKernelOutput, publicKernelProof, revertReason, returnValues, gasUsed };
   }
 }