Tweak aztec-node public simulation so we can get gas-used out of it

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

@@ -13,6 +13,7 @@ import {
   LogType,
   MerkleTreeId,
   NullifierMembershipWitness,
+  type ProcessOutput,
   type ProverClient,
   PublicDataWitness,
   type SequencerConfig,
@@ -634,7 +635,7 @@ export class AztecNodeService implements AztecNode {
    * Simulates the public part of a transaction with the current state.
    * @param tx - The transaction to simulate.
    **/
-  public async simulatePublicCalls(tx: Tx) {
+  public async simulatePublicCalls(tx: Tx): Promise<ProcessOutput> {
     this.log.info(`Simulating tx ${tx.getTxHash()}`);
     const blockNumber = (await this.blockSource.getBlockNumber()) + 1;
 
@@ -660,6 +661,7 @@ export class AztecNodeService implements AztecNode {
       new WASMSimulator(),
     );
     const processor = await publicProcessorFactory.create(prevHeader, newGlobalVariables);
+    // REFACTOR: Consider merging ProcessReturnValues into ProcessedTx
     const [processedTxs, failedTxs, returns] = await processor.process([tx]);
     if (failedTxs.length) {
       this.log.warn(`Simulated tx ${tx.getTxHash()} fails: ${failedTxs[0].error}`);
@@ -671,7 +673,15 @@ export class AztecNodeService implements AztecNode {
       throw reverted[0].revertReason;
     }
     this.log.info(`Simulated tx ${tx.getTxHash()} succeeds`);
-    return returns[0];
+    const [processedTx] = processedTxs;
+    return {
+      constants: processedTx.data.constants,
+      encryptedLogs: processedTx.encryptedLogs,
+      unencryptedLogs: processedTx.unencryptedLogs,
+      end: processedTx.data.end,
+      revertReason: processedTx.revertReason,
+      publicReturnValues: returns[0],
+    };
   }
 
   public setConfig(config: Partial<SequencerConfig>): Promise<void> {

yarn-project/aztec.js/src/contract/contract_function_interaction.ts

@@ -105,7 +105,7 @@ export class ContractFunctionInteraction extends BaseContractInteraction {
       const txRequest = await this.create();
       const simulatedTx = await this.pxe.simulateTx(txRequest, true);
       this.txRequest = undefined;
-      const flattened = simulatedTx.publicReturnValues;
+      const flattened = simulatedTx.publicOutput?.publicReturnValues;
       return flattened ? decodeReturnValues(this.functionDao, flattened) : [];
     }
   }

yarn-project/circuit-types/src/interfaces/aztec-node.ts

@@ -21,7 +21,7 @@ import {
 } from '../logs/index.js';
 import { type MerkleTreeId } from '../merkle_tree_id.js';
 import { type SiblingPath } from '../sibling_path/index.js';
-import { type ProcessReturnValues, type Tx, type TxHash, type TxReceipt } from '../tx/index.js';
+import { type ProcessOutput, type Tx, type TxHash, type TxReceipt } from '../tx/index.js';
 import { type TxEffect } from '../tx_effect.js';
 import { type SequencerConfig } from './configs.js';
 import { type L2BlockNumber } from './l2_block_number.js';
@@ -282,7 +282,7 @@ export interface AztecNode {
    * This currently just checks that the transaction execution succeeds.
    * @param tx - The transaction to simulate.
    **/
-  simulatePublicCalls(tx: Tx): Promise<ProcessReturnValues>;
+  simulatePublicCalls(tx: Tx): Promise<ProcessOutput>;
 
   /**
    * Updates the configuration of this node.

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

@@ -12,7 +12,11 @@ import {
   computeContractClassId,
   getContractClassFromArtifact,
 } from '@aztec/circuits.js';
-import { makePublicCallRequest } from '@aztec/circuits.js/testing';
+import {
+  makeCombinedAccumulatedData,
+  makeCombinedConstantData,
+  makePublicCallRequest,
+} from '@aztec/circuits.js/testing';
 import { type ContractArtifact } from '@aztec/foundation/abi';
 import { makeTuple } from '@aztec/foundation/array';
 import { times } from '@aztec/foundation/collection';
@@ -23,7 +27,7 @@ import { type ContractInstanceWithAddress, SerializableContractInstance } from '
 import { EncryptedL2Log } from './logs/encrypted_l2_log.js';
 import { EncryptedFunctionL2Logs, EncryptedTxL2Logs, Note, UnencryptedTxL2Logs } from './logs/index.js';
 import { ExtendedNote } from './notes/index.js';
-import { type ProcessReturnValues, SimulatedTx, Tx, TxHash } from './tx/index.js';
+import { type ProcessOutput, type ProcessReturnValues, SimulatedTx, Tx, TxHash } from './tx/index.js';
 
 /**
  * Testing utility to create empty logs composed from a single empty log.
@@ -114,7 +118,15 @@ export const mockTxForRollup = (seed = 1, { hasLogs = false }: { hasLogs?: boole
 export const mockSimulatedTx = (seed = 1, hasLogs = true) => {
   const tx = mockTx(seed, { hasLogs });
   const dec: ProcessReturnValues = [new Fr(1n), new Fr(2n), new Fr(3n), new Fr(4n)];
-  return new SimulatedTx(tx, dec, dec);
+  const output: ProcessOutput = {
+    constants: makeCombinedConstantData(),
+    encryptedLogs: tx.encryptedLogs,
+    unencryptedLogs: tx.unencryptedLogs,
+    end: makeCombinedAccumulatedData(),
+    revertReason: undefined,
+    publicReturnValues: dec,
+  };
+  return new SimulatedTx(tx, dec, output);
 };
 
 export const randomContractArtifact = (): ContractArtifact => ({

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

@@ -10,7 +10,7 @@ describe('simulated_tx', () => {
   it('convert undefined effects to and from json', () => {
     const simulatedTx = mockSimulatedTx();
     simulatedTx.privateReturnValues = undefined;
-    simulatedTx.publicReturnValues = undefined;
+    simulatedTx.publicOutput = undefined;
     expect(SimulatedTx.fromJSON(simulatedTx.toJSON())).toEqual(simulatedTx);
   });
 });

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

@@ -1,32 +1,59 @@
-import { Fr } from '@aztec/circuits.js';
+import { CombinedAccumulatedData, CombinedConstantData, Fr } from '@aztec/circuits.js';
 
+import { EncryptedTxL2Logs, UnencryptedTxL2Logs } from '../logs/index.js';
+import { type ProcessedTx } from './processed_tx.js';
 import { Tx } from './tx.js';
 
+/** Return values of simulating a circuit. */
 export type ProcessReturnValues = Fr[] | undefined;
 
+/**
+ * Outputs of processing the public component of a transaction.
+ * REFACTOR: Rename.
+ */
+export type ProcessOutput = Pick<ProcessedTx, 'encryptedLogs' | 'unencryptedLogs' | 'revertReason'> &
+  Pick<ProcessedTx['data'], 'constants' | 'end'> & { publicReturnValues: ProcessReturnValues };
+
+function processOutputToJSON(output: ProcessOutput) {
+  return {
+    encryptedLogs: output.encryptedLogs.toJSON(),
+    unencryptedLogs: output.unencryptedLogs.toJSON(),
+    revertReason: output.revertReason,
+    constants: output.constants.toBuffer().toString('hex'),
+    end: output.end.toBuffer().toString('hex'),
+    publicReturnValues: JSON.stringify(output.publicReturnValues?.map(fr => fr.toString())),
+  };
+}
+
+function processOutputFromJSON(json: any): ProcessOutput {
+  return {
+    encryptedLogs: EncryptedTxL2Logs.fromJSON(json.encryptedLogs),
+    unencryptedLogs: UnencryptedTxL2Logs.fromJSON(json.unencryptedLogs),
+    revertReason: json.revertReason,
+    constants: CombinedConstantData.fromBuffer(Buffer.from(json.constants, 'hex')),
+    end: CombinedAccumulatedData.fromBuffer(Buffer.from(json.end, 'hex')),
+    publicReturnValues: json.publicReturnValues.map(Fr.fromString),
+  };
+}
+
+// REFACTOR: Review what we need to expose to the user when running a simulation.
+// Eg tx already has encrypted and unencrypted logs, but those cover only the ones
+// emitted during private. We need the ones from ProcessOutput to include the public
+// ones as well. However, those would only be present if the user chooses to simulate
+// the public side of things. This also points at this class needing to be split into
+// two: one with just private simulation, and one that also includes public simulation.
 export class SimulatedTx {
-  constructor(
-    public tx: Tx,
-    public privateReturnValues?: ProcessReturnValues,
-    public publicReturnValues?: ProcessReturnValues,
-  ) {}
+  constructor(public tx: Tx, public privateReturnValues?: ProcessReturnValues, public publicOutput?: ProcessOutput) {}
 
   /**
    * Convert a SimulatedTx class object to a plain JSON object.
    * @returns A plain object with SimulatedTx properties.
    */
   public toJSON() {
-    const returnToJson = (data: ProcessReturnValues | undefined): string => {
-      if (data === undefined) {
-        return JSON.stringify(data);
-      }
-      return JSON.stringify(data.map(fr => fr.toString()));
-    };
-
     return {
       tx: this.tx.toJSON(),
-      privateReturnValues: returnToJson(this.privateReturnValues),
-      publicReturnValues: returnToJson(this.publicReturnValues),
+      privateReturnValues: this.privateReturnValues?.map(fr => fr.toString()),
+      publicOutput: this.publicOutput && processOutputToJSON(this.publicOutput),
     };
   }
 
@@ -36,17 +63,10 @@ export class SimulatedTx {
    * @returns A Tx class object.
    */
   public static fromJSON(obj: any) {
-    const returnFromJson = (json: string): ProcessReturnValues | undefined => {
-      if (json === undefined) {
-        return json;
-      }
-      return JSON.parse(json).map(Fr.fromString);
-    };
-
     const tx = Tx.fromJSON(obj.tx);
-    const privateReturnValues = returnFromJson(obj.privateReturnValues);
-    const publicReturnValues = returnFromJson(obj.publicReturnValues);
+    const publicOutput = obj.publicOutput ? processOutputFromJSON(obj.publicOutput) : undefined;
+    const privateReturnValues = obj.privateReturnValues?.map(Fr.fromString);
 
-    return new SimulatedTx(tx, privateReturnValues, publicReturnValues);
+    return new SimulatedTx(tx, privateReturnValues, publicOutput);
   }
 }

yarn-project/circuits.js/src/tests/factories.ts

@@ -292,6 +292,10 @@ export function makeRollupValidationRequests(seed = 1) {
   return new RollupValidationRequests(new MaxBlockNumber(true, new Fr(seed + 0x31415)));
 }
 
+export function makeCombinedConstantData(seed = 1): CombinedConstantData {
+  return new CombinedConstantData(makeHeader(seed), makeTxContext(seed + 0x100));
+}
+
 /**
  * Creates arbitrary accumulated data.
  * @param seed - The seed to use for generating the accumulated data.

yarn-project/pxe/src/pxe_service/pxe_service.ts

@@ -420,7 +420,7 @@ export class PXEService implements PXE {
     txRequest: TxExecutionRequest,
     simulatePublic: boolean,
     msgSender: AztecAddress | undefined = undefined,
-  ) {
+  ): Promise<SimulatedTx> {
     if (!txRequest.functionData.isPrivate) {
       throw new Error(`Public entrypoints are not allowed`);
     }
@@ -441,7 +441,7 @@ export class PXEService implements PXE {
       }
 
       if (simulatePublic) {
-        simulatedTx.publicReturnValues = await this.#simulatePublicCalls(simulatedTx.tx);
+        simulatedTx.publicOutput = await this.#simulatePublicCalls(simulatedTx.tx);
       }
 
       if (!msgSender) {