@aztec/simulator 0.67.1 → 0.68.1

package/src/public/public_processor.ts

@@ -4,7 +4,6 @@ import {
   type MerkleTreeWriteOperations,
   NestedProcessReturnValues,
   type ProcessedTx,
-  type ProcessedTxHandler,
   Tx,
   TxExecutionPhase,
   type TxValidator,
@@ -16,6 +15,7 @@ import {
   type BlockHeader,
   type ContractDataSource,
   Fr,
+  Gas,
   type GlobalVariables,
   MAX_NOTE_HASHES_PER_TX,
   MAX_NULLIFIERS_PER_TX,
@@ -24,9 +24,10 @@ import {
 } from '@aztec/circuits.js';
 import { padArrayEnd } from '@aztec/foundation/collection';
 import { createLogger } from '@aztec/foundation/log';
-import { Timer } from '@aztec/foundation/timer';
-import { ContractClassRegisteredEvent, ProtocolContractAddress } from '@aztec/protocol-contracts';
-import { Attributes, type TelemetryClient, type Tracer, trackSpan } from '@aztec/telemetry-client';
+import { type DateProvider, Timer, elapsed, executeTimeout } from '@aztec/foundation/timer';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import { ContractClassRegisteredEvent } from '@aztec/protocol-contracts/class-registerer';
+import { Attributes, type TelemetryClient, type Traceable, type Tracer, trackSpan } from '@aztec/telemetry-client';
 
 import { computeFeePayerBalanceLeafSlot, computeFeePayerBalanceStorageSlot } from './fee_payment.js';
 import { WorldStateDB } from './public_db_sources.js';
@@ -37,28 +38,35 @@ import { PublicTxSimulator } from './public_tx_simulator.js';
  * Creates new instances of PublicProcessor given the provided merkle tree db and contract data source.
  */
 export class PublicProcessorFactory {
-  constructor(private contractDataSource: ContractDataSource, private telemetryClient: TelemetryClient) {}
+  constructor(
+    private contractDataSource: ContractDataSource,
+    private dateProvider: DateProvider,
+    private telemetryClient: TelemetryClient,
+  ) {}
 
   /**
    * Creates a new instance of a PublicProcessor.
    * @param historicalHeader - The header of a block previous to the one in which the tx is included.
    * @param globalVariables - The global variables for the block being processed.
+   * @param enforceFeePayment - Allows disabling balance checks for fee estimations.
    * @returns A new instance of a PublicProcessor.
    */
   public create(
     merkleTree: MerkleTreeWriteOperations,
     maybeHistoricalHeader: BlockHeader | undefined,
     globalVariables: GlobalVariables,
+    enforceFeePayment: boolean,
   ): PublicProcessor {
     const historicalHeader = maybeHistoricalHeader ?? merkleTree.getInitialHeader();
 
     const worldStateDB = new WorldStateDB(merkleTree, this.contractDataSource);
-    const publicTxSimulator = new PublicTxSimulator(
+    const publicTxSimulator = this.createPublicTxSimulator(
       merkleTree,
       worldStateDB,
       this.telemetryClient,
       globalVariables,
       /*doMerkleOperations=*/ true,
+      enforceFeePayment,
     );
 
     return new PublicProcessor(
@@ -67,16 +75,42 @@ export class PublicProcessorFactory {
       historicalHeader,
       worldStateDB,
       publicTxSimulator,
+      this.dateProvider,
       this.telemetryClient,
     );
   }
+
+  protected createPublicTxSimulator(
+    db: MerkleTreeWriteOperations,
+    worldStateDB: WorldStateDB,
+    telemetryClient: TelemetryClient,
+    globalVariables: GlobalVariables,
+    doMerkleOperations: boolean,
+    enforceFeePayment: boolean,
+  ) {
+    return new PublicTxSimulator(
+      db,
+      worldStateDB,
+      telemetryClient,
+      globalVariables,
+      doMerkleOperations,
+      enforceFeePayment,
+    );
+  }
+}
+
+class PublicProcessorTimeoutError extends Error {
+  constructor(message: string = 'Timed out while processing tx') {
+    super(message);
+    this.name = 'PublicProcessorTimeoutError';
+  }
 }
 
 /**
  * Converts Txs lifted from the P2P module into ProcessedTx objects by executing
  * any public function calls in them. Txs with private calls only are unaffected.
  */
-export class PublicProcessor {
+export class PublicProcessor implements Traceable {
   private metrics: PublicProcessorMetrics;
   constructor(
     protected db: MerkleTreeWriteOperations,
@@ -84,6 +118,7 @@ export class PublicProcessor {
     protected historicalHeader: BlockHeader,
     protected worldStateDB: WorldStateDB,
     protected publicTxSimulator: PublicTxSimulator,
+    private dateProvider: DateProvider,
     telemetryClient: TelemetryClient,
     private log = createLogger('simulator:public-processor'),
   ) {
@@ -103,14 +138,16 @@ export class PublicProcessor {
   public async process(
     txs: Tx[],
     maxTransactions = txs.length,
-    processedTxHandler?: ProcessedTxHandler,
     txValidator?: TxValidator<ProcessedTx>,
+    deadline?: Date,
   ): Promise<[ProcessedTx[], FailedTx[], NestedProcessReturnValues[]]> {
     // The processor modifies the tx objects in place, so we need to clone them.
     txs = txs.map(tx => Tx.clone(tx));
     const result: ProcessedTx[] = [];
     const failed: FailedTx[] = [];
     let returns: NestedProcessReturnValues[] = [];
+    let totalGas = new Gas(0, 0);
+    const timer = new Timer();
 
     for (const tx of txs) {
       // only process up to the limit of the block
@@ -118,95 +155,141 @@ export class PublicProcessor {
         break;
       }
       try {
-        const [processedTx, returnValues] = !tx.hasPublicCalls()
-          ? await this.processPrivateOnlyTx(tx)
-          : await this.processTxWithPublicCalls(tx);
-
-        this.log.verbose(
-          !tx.hasPublicCalls()
-            ? `Processed tx ${processedTx.hash} with no public calls`
-            : `Processed tx ${processedTx.hash} with ${tx.enqueuedPublicFunctionCalls.length} public calls`,
-          {
-            txHash: processedTx.hash,
-            txFee: processedTx.txEffect.transactionFee.toBigInt(),
-            revertCode: processedTx.txEffect.revertCode.getCode(),
-            revertReason: processedTx.revertReason,
-            gasUsed: processedTx.gasUsed,
-            publicDataWriteCount: processedTx.txEffect.publicDataWrites.length,
-            nullifierCount: processedTx.txEffect.nullifiers.length,
-            noteHashCount: processedTx.txEffect.noteHashes.length,
-            contractClassLogCount: processedTx.txEffect.contractClassLogs.getTotalLogCount(),
-            unencryptedLogCount: processedTx.txEffect.unencryptedLogs.getTotalLogCount(),
-            privateLogCount: processedTx.txEffect.privateLogs.length,
-            l2ToL1MessageCount: processedTx.txEffect.l2ToL1Msgs.length,
-          },
-        );
-
-        // Commit the state updates from this transaction
-        await this.worldStateDB.commit();
-
-        // Re-validate the transaction
-        if (txValidator) {
-          // Only accept processed transactions that are not double-spends,
-          // public functions emitting nullifiers would pass earlier check but fail here.
-          // Note that we're checking all nullifiers generated in the private execution twice,
-          // we could store the ones already checked and skip them here as an optimization.
-          const [_, invalid] = await txValidator.validateTxs([processedTx]);
-          if (invalid.length) {
-            throw new Error(`Transaction ${invalid[0].hash} invalid after processing public functions`);
-          }
-        }
-        // if we were given a handler then send the transaction to it for block building or proving
-        if (processedTxHandler) {
-          await processedTxHandler.addNewTx(processedTx);
-        }
-        // Update the state so that the next tx in the loop has the correct .startState
-        // NB: before this change, all .startStates were actually incorrect, but the issue was never caught because we either:
-        // a) had only 1 tx with public calls per block, so this loop had len 1
-        // b) always had a txHandler with the same db passed to it as this.db, which updated the db in buildBaseRollupHints in this loop
-        // To see how this ^ happens, move back to one shared db in test_context and run orchestrator_multi_public_functions.test.ts
-        // The below is taken from buildBaseRollupHints:
-        await this.db.appendLeaves(
-          MerkleTreeId.NOTE_HASH_TREE,
-          padArrayEnd(processedTx.txEffect.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
-        );
-        try {
-          await this.db.batchInsert(
-            MerkleTreeId.NULLIFIER_TREE,
-            padArrayEnd(processedTx.txEffect.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX).map(n => n.toBuffer()),
-            NULLIFIER_SUBTREE_HEIGHT,
-          );
-        } catch (error) {
-          if (txValidator) {
-            // Ideally the validator has already caught this above, but just in case:
-            throw new Error(`Transaction ${processedTx.hash} invalid after processing public functions`);
-          } else {
-            // We have no validator and assume this call should blindly process txs with duplicates being caught later
-            this.log.warn(`Detected duplicate nullifier after public processing for: ${processedTx.hash}.`);
-          }
-        }
-
-        await this.db.sequentialInsert(
-          MerkleTreeId.PUBLIC_DATA_TREE,
-          processedTx.txEffect.publicDataWrites.map(x => x.toBuffer()),
-        );
+        const [processedTx, returnValues] = await this.processTx(tx, txValidator, deadline);
         result.push(processedTx);
-        returns = returns.concat(returnValues ?? []);
+        returns = returns.concat(returnValues);
+        totalGas = totalGas.add(processedTx.gasUsed.publicGas);
       } catch (err: any) {
+        if (err?.name === 'PublicProcessorTimeoutError') {
+          this.log.warn(`Stopping tx processing due to timeout.`);
+          break;
+        }
         const errorMessage = err instanceof Error ? err.message : 'Unknown error';
         this.log.warn(`Failed to process tx ${tx.getTxHash()}: ${errorMessage} ${err?.stack}`);
 
-        failed.push({
-          tx,
-          error: err instanceof Error ? err : new Error(errorMessage),
-        });
+        failed.push({ tx, error: err instanceof Error ? err : new Error(errorMessage) });
         returns.push(new NestedProcessReturnValues([]));
       }
     }
 
+    const duration = timer.s();
+    const rate = duration > 0 ? totalGas.l2Gas / duration : 0;
+    this.metrics.recordAllTxs(totalGas, rate);
+
     return [result, failed, returns];
   }
 
+  @trackSpan('PublicProcessor.processTx', tx => ({ [Attributes.TX_HASH]: tx.tryGetTxHash()?.toString() }))
+  private async processTx(
+    tx: Tx,
+    txValidator?: TxValidator<ProcessedTx>,
+    deadline?: Date,
+  ): Promise<[ProcessedTx, NestedProcessReturnValues[]]> {
+    const [time, [processedTx, returnValues]] = await elapsed(() => this.processTxWithinDeadline(tx, deadline));
+
+    this.log.verbose(
+      !tx.hasPublicCalls()
+        ? `Processed tx ${processedTx.hash} with no public calls in ${time}ms`
+        : `Processed tx ${processedTx.hash} with ${tx.enqueuedPublicFunctionCalls.length} public calls in ${time}ms`,
+      {
+        txHash: processedTx.hash,
+        txFee: processedTx.txEffect.transactionFee.toBigInt(),
+        revertCode: processedTx.txEffect.revertCode.getCode(),
+        revertReason: processedTx.revertReason,
+        gasUsed: processedTx.gasUsed,
+        publicDataWriteCount: processedTx.txEffect.publicDataWrites.length,
+        nullifierCount: processedTx.txEffect.nullifiers.length,
+        noteHashCount: processedTx.txEffect.noteHashes.length,
+        contractClassLogCount: processedTx.txEffect.contractClassLogs.getTotalLogCount(),
+        unencryptedLogCount: processedTx.txEffect.unencryptedLogs.getTotalLogCount(),
+        privateLogCount: processedTx.txEffect.privateLogs.length,
+        l2ToL1MessageCount: processedTx.txEffect.l2ToL1Msgs.length,
+        durationMs: time,
+      },
+    );
+
+    // Commit the state updates from this transaction
+    await this.worldStateDB.commit();
+
+    // Re-validate the transaction
+    if (txValidator) {
+      // Only accept processed transactions that are not double-spends,
+      // public functions emitting nullifiers would pass earlier check but fail here.
+      // Note that we're checking all nullifiers generated in the private execution twice,
+      // we could store the ones already checked and skip them here as an optimization.
+      const [_, invalid] = await txValidator.validateTxs([processedTx]);
+      if (invalid.length) {
+        throw new Error(`Transaction ${invalid[0].hash} invalid after processing public functions`);
+      }
+    }
+    // Update the state so that the next tx in the loop has the correct .startState
+    // NB: before this change, all .startStates were actually incorrect, but the issue was never caught because we either:
+    // a) had only 1 tx with public calls per block, so this loop had len 1
+    // b) always had a txHandler with the same db passed to it as this.db, which updated the db in buildBaseRollupHints in this loop
+    // To see how this ^ happens, move back to one shared db in test_context and run orchestrator_multi_public_functions.test.ts
+    // The below is taken from buildBaseRollupHints:
+    const treeInsertionStart = process.hrtime.bigint();
+    await this.db.appendLeaves(
+      MerkleTreeId.NOTE_HASH_TREE,
+      padArrayEnd(processedTx.txEffect.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
+    );
+    try {
+      await this.db.batchInsert(
+        MerkleTreeId.NULLIFIER_TREE,
+        padArrayEnd(processedTx.txEffect.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX).map(n => n.toBuffer()),
+        NULLIFIER_SUBTREE_HEIGHT,
+      );
+    } catch (error) {
+      if (txValidator) {
+        // Ideally the validator has already caught this above, but just in case:
+        throw new Error(`Transaction ${processedTx.hash} invalid after processing public functions`);
+      } else {
+        // We have no validator and assume this call should blindly process txs with duplicates being caught later
+        this.log.warn(`Detected duplicate nullifier after public processing for: ${processedTx.hash}.`);
+      }
+    }
+
+    await this.db.sequentialInsert(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      processedTx.txEffect.publicDataWrites.map(x => x.toBuffer()),
+    );
+    const treeInsertionEnd = process.hrtime.bigint();
+    this.metrics.recordTreeInsertions(Number(treeInsertionEnd - treeInsertionStart) / 1_000);
+
+    return [processedTx, returnValues ?? []];
+  }
+
+  /** Processes the given tx within deadline. Returns timeout if deadline is hit. */
+  private async processTxWithinDeadline(
+    tx: Tx,
+    deadline?: Date,
+  ): Promise<[ProcessedTx, NestedProcessReturnValues[] | undefined]> {
+    const processFn: () => Promise<[ProcessedTx, NestedProcessReturnValues[] | undefined]> = tx.hasPublicCalls()
+      ? () => this.processTxWithPublicCalls(tx)
+      : () => this.processPrivateOnlyTx(tx);
+
+    if (!deadline) {
+      return await processFn();
+    }
+
+    const timeout = +deadline - this.dateProvider.now();
+    this.log.debug(`Processing tx ${tx.getTxHash().toString()} within ${timeout}ms`, {
+      deadline: deadline.toISOString(),
+      now: new Date(this.dateProvider.now()).toISOString(),
+      txHash: tx.getTxHash().toString(),
+    });
+
+    if (timeout < 0) {
+      throw new PublicProcessorTimeoutError();
+    }
+
+    return await executeTimeout(
+      () => processFn(),
+      timeout,
+      () => new PublicProcessorTimeoutError(),
+    );
+  }
+
   /**
    * Creates the public data write for paying the tx fee.
    * This is used in private only txs, since for txs with public calls
@@ -241,7 +324,7 @@ export class PublicProcessor {
   @trackSpan('PublicProcessor.processPrivateOnlyTx', (tx: Tx) => ({
     [Attributes.TX_HASH]: tx.getTxHash().toString(),
   }))
-  private async processPrivateOnlyTx(tx: Tx): Promise<[ProcessedTx]> {
+  private async processPrivateOnlyTx(tx: Tx): Promise<[ProcessedTx, undefined]> {
     const gasFees = this.globalVariables.gasFees;
     const transactionFee = tx.data.gasUsed.computeFee(gasFees);
 
@@ -260,7 +343,7 @@ export class PublicProcessor {
         .filter(log => ContractClassRegisteredEvent.isContractClassRegisteredEvent(log.data))
         .map(log => ContractClassRegisteredEvent.fromLog(log.data)),
     );
-    return [processedTx];
+    return [processedTx, undefined];
   }
 
   @trackSpan('PublicProcessor.processTxWithPublicCalls', tx => ({
@@ -294,7 +377,7 @@ export class PublicProcessor {
 
     const phaseCount = processedPhases.length;
     const durationMs = timer.ms();
-    this.metrics.recordTx(phaseCount, durationMs);
+    this.metrics.recordTx(phaseCount, durationMs, gasUsed.publicGas);
 
     const processedTx = makeProcessedTxFromTxWithPublicCalls(tx, avmProvingRequest, gasUsed, revertCode, revertReason);
 

package/src/public/public_processor_metrics.ts

@@ -1,7 +1,9 @@
 import { type TxExecutionPhase } from '@aztec/circuit-types';
-import { type ContractClassRegisteredEvent } from '@aztec/protocol-contracts';
+import { type Gas } from '@aztec/circuits.js';
+import { type ContractClassRegisteredEvent } from '@aztec/protocol-contracts/class-registerer';
 import {
   Attributes,
+  type Gauge,
   type Histogram,
   Metrics,
   type TelemetryClient,
@@ -21,6 +23,12 @@ export class PublicProcessorMetrics {
   private phaseCount: UpDownCounter;
 
   private bytecodeDeployed: Histogram;
+  private totalGas: Gauge;
+  private totalGasHistogram: Histogram;
+  private gasRate: Histogram;
+  private txGas: Histogram;
+
+  private treeInsertionDuration: Histogram;
 
   constructor(client: TelemetryClient, name = 'PublicProcessor') {
     this.tracer = client.getTracer(name);
@@ -54,6 +62,32 @@ export class PublicProcessorMetrics {
       description: 'Size of deployed bytecode',
       unit: 'By',
     });
+
+    this.totalGas = meter.createGauge(Metrics.PUBLIC_PROCESSOR_TOTAL_GAS, {
+      description: 'Total gas used in block',
+      unit: 'gas',
+    });
+
+    this.totalGasHistogram = meter.createHistogram(Metrics.PUBLIC_PROCESSOR_TOTAL_GAS_HISTOGRAM, {
+      description: 'Total gas used in block as histogram',
+      unit: 'gas/block',
+    });
+
+    this.txGas = meter.createHistogram(Metrics.PUBLIC_PROCESSOR_TX_GAS, {
+      description: 'Gas used in transaction',
+      unit: 'gas/tx',
+    });
+
+    this.gasRate = meter.createHistogram(Metrics.PUBLIC_PROCESSOR_GAS_RATE, {
+      description: 'L2 gas per second for complete block',
+      unit: 'gas/s',
+    });
+
+    this.treeInsertionDuration = meter.createHistogram(Metrics.PUBLIC_PROCESSOR_TREE_INSERTION, {
+      description: 'How long it takes for tree insertion',
+      unit: 'us',
+      valueType: ValueType.INT,
+    });
   }
 
   recordPhaseDuration(phaseName: TxExecutionPhase, durationMs: number) {
@@ -61,12 +95,36 @@ export class PublicProcessorMetrics {
     this.phaseDuration.record(Math.ceil(durationMs), { [Attributes.TX_PHASE_NAME]: phaseName });
   }
 
-  recordTx(phaseCount: number, durationMs: number) {
+  recordTx(phaseCount: number, durationMs: number, gasUsed: Gas) {
     this.txPhaseCount.add(phaseCount);
     this.txDuration.record(Math.ceil(durationMs));
     this.txCount.add(1, {
       [Attributes.OK]: true,
     });
+    this.txGas.record(gasUsed.daGas, {
+      [Attributes.GAS_DIMENSION]: 'DA',
+    });
+    this.txGas.record(gasUsed.l2Gas, {
+      [Attributes.GAS_DIMENSION]: 'L2',
+    });
+  }
+
+  recordAllTxs(totalGas: Gas, gasRate: number) {
+    this.totalGas.record(totalGas.daGas, {
+      [Attributes.GAS_DIMENSION]: 'DA',
+    });
+    this.totalGas.record(totalGas.l2Gas, {
+      [Attributes.GAS_DIMENSION]: 'L2',
+    });
+    this.gasRate.record(gasRate, {
+      [Attributes.GAS_DIMENSION]: 'L2',
+    });
+    this.totalGasHistogram.record(totalGas.daGas, {
+      [Attributes.GAS_DIMENSION]: 'DA',
+    });
+    this.totalGasHistogram.record(totalGas.l2Gas, {
+      [Attributes.GAS_DIMENSION]: 'L2',
+    });
   }
 
   recordFailedTx() {
@@ -89,4 +147,8 @@ export class PublicProcessorMetrics {
       this.bytecodeDeployed.record(totalBytecode);
     }
   }
+
+  recordTreeInsertions(durationUs: number) {
+    this.treeInsertionDuration.record(Math.ceil(durationUs));
+  }
 }

package/src/public/public_tx_context.ts

@@ -10,7 +10,6 @@ import {
   TxHash,
 } from '@aztec/circuit-types';
 import {
-  AppendOnlyTreeSnapshot,
   AvmCircuitInputs,
   type AvmCircuitPublicInputs,
   type AztecAddress,
@@ -19,12 +18,15 @@ import {
   type GasSettings,
   type GlobalVariables,
   MAX_L2_GAS_PER_TX_PUBLIC_PORTION,
+  MAX_NOTE_HASHES_PER_TX,
+  MAX_NULLIFIERS_PER_TX,
   type PrivateToPublicAccumulatedData,
   type PublicCallRequest,
   PublicCircuitPublicInputs,
   RevertCode,
   type StateReference,
   TreeSnapshots,
+  computeTransactionFee,
   countAccumulatedItems,
 } from '@aztec/circuits.js';
 import { type Logger, createLogger } from '@aztec/foundation/log';
@@ -91,7 +93,7 @@ export class PublicTxContext {
     const previousAccumulatedDataArrayLengths = new SideEffectArrayLengths(
       /*publicDataWrites*/ 0,
       /*protocolPublicDataWrites*/ 0,
-      countAccumulatedItems(nonRevertibleAccumulatedDataFromPrivate.noteHashes),
+      /*noteHashes*/ 0,
       /*nullifiers=*/ 0,
       countAccumulatedItems(nonRevertibleAccumulatedDataFromPrivate.l2ToL1Msgs),
       /*unencryptedLogsHashes*/ 0,
@@ -102,7 +104,12 @@ export class PublicTxContext {
     );
 
     // Transaction level state manager that will be forked for revertible phases.
-    const txStateManager = await AvmPersistableStateManager.create(worldStateDB, enqueuedCallTrace, doMerkleOperations);
+    const txStateManager = await AvmPersistableStateManager.create(
+      worldStateDB,
+      enqueuedCallTrace,
+      doMerkleOperations,
+      fetchTxHash(nonRevertibleAccumulatedDataFromPrivate),
+    );
 
     const gasSettings = tx.data.constants.txContext.gasSettings;
     const gasUsedByPrivate = tx.data.gasUsed;
@@ -186,12 +193,7 @@ export class PublicTxContext {
    * @returns The transaction's hash.
    */
   getTxHash(): TxHash {
-    // Private kernel functions are executed client side and for this reason tx hash is already set as first nullifier
-    const firstNullifier = this.nonRevertibleAccumulatedDataFromPrivate.nullifiers[0];
-    if (!firstNullifier || firstNullifier.isZero()) {
-      throw new Error(`Cannot get tx hash since first nullifier is missing`);
-    }
-    return new TxHash(firstNullifier.toBuffer());
+    return fetchTxHash(this.nonRevertibleAccumulatedDataFromPrivate);
   }
 
   /**
@@ -280,6 +282,15 @@ export class PublicTxContext {
     return this.getTotalGasUsed().sub(teardownGasLimits).add(this.teardownGasUsed);
   }
 
+  /**
+   * Compute the public gas used using the actual gas used during teardown instead
+   * of the teardown gas limit.
+   */
+  getActualPublicGasUsed(): Gas {
+    assert(this.halted, 'Can only compute actual gas used after tx execution ends');
+    return this.gasUsedByPublic.add(this.teardownGasUsed);
+  }
+
   /**
    * Get the transaction fee as is available to the specified phase.
    * Only teardown should have access to the actual transaction fee.
@@ -297,12 +308,15 @@ export class PublicTxContext {
    * Should only be called during or after teardown.
    */
   private getTransactionFeeUnsafe(): Fr {
-    const txFee = this.getTotalGasUsed().computeFee(this.globalVariables.gasFees);
+    const gasUsed = this.getTotalGasUsed();
+    const txFee = computeTransactionFee(this.globalVariables.gasFees, this.gasSettings, gasUsed);
+
     this.log.debug(`Computed tx fee`, {
       txFee,
-      gasUsed: inspect(this.getTotalGasUsed()),
+      gasUsed: inspect(gasUsed),
       gasFees: inspect(this.globalVariables.gasFees),
     });
+
     return txFee;
   }
 
@@ -311,16 +325,29 @@ export class PublicTxContext {
    */
   private generateAvmCircuitPublicInputs(endStateReference: StateReference): AvmCircuitPublicInputs {
     assert(this.halted, 'Can only get AvmCircuitPublicInputs after tx execution ends');
-    const ephemeralTrees = this.state.getActiveStateManager().merkleTrees.treeMap;
-
-    const getAppendSnaphot = (id: MerkleTreeId) => {
-      const tree = ephemeralTrees.get(id)!;
-      return new AppendOnlyTreeSnapshot(tree.getRoot(), Number(tree.leafCount));
-    };
-
-    const noteHashTree = getAppendSnaphot(MerkleTreeId.NOTE_HASH_TREE);
-    const nullifierTree = getAppendSnaphot(MerkleTreeId.NULLIFIER_TREE);
-    const publicDataTree = getAppendSnaphot(MerkleTreeId.PUBLIC_DATA_TREE);
+    const ephemeralTrees = this.state.getActiveStateManager().merkleTrees;
+
+    const noteHashTree = ephemeralTrees.getTreeSnapshot(MerkleTreeId.NOTE_HASH_TREE);
+    const nullifierTree = ephemeralTrees.getTreeSnapshot(MerkleTreeId.NULLIFIER_TREE);
+    const publicDataTree = ephemeralTrees.getTreeSnapshot(MerkleTreeId.PUBLIC_DATA_TREE);
+    // Pad the note hash and nullifier trees
+    const paddedNoteHashTreeSize =
+      this.startStateReference.partial.noteHashTree.nextAvailableLeafIndex + MAX_NOTE_HASHES_PER_TX;
+    if (noteHashTree.nextAvailableLeafIndex > paddedNoteHashTreeSize) {
+      throw new Error(
+        `Inserted too many leaves in note hash tree: ${noteHashTree.nextAvailableLeafIndex} > ${paddedNoteHashTreeSize}`,
+      );
+    }
+    noteHashTree.nextAvailableLeafIndex = paddedNoteHashTreeSize;
+
+    const paddedNullifierTreeSize =
+      this.startStateReference.partial.nullifierTree.nextAvailableLeafIndex + MAX_NULLIFIERS_PER_TX;
+    if (nullifierTree.nextAvailableLeafIndex > paddedNullifierTreeSize) {
+      throw new Error(
+        `Inserted too many leaves in nullifier tree: ${nullifierTree.nextAvailableLeafIndex} > ${paddedNullifierTreeSize}`,
+      );
+    }
+    nullifierTree.nextAvailableLeafIndex = paddedNullifierTreeSize;
 
     const endTreeSnapshots = new TreeSnapshots(
       endStateReference.l1ToL2MessageTree,
@@ -425,3 +452,12 @@ function applyMaxToAvailableGas(availableGas: Gas) {
     /*l2Gas=*/ Math.min(availableGas.l2Gas, MAX_L2_GAS_PER_TX_PUBLIC_PORTION),
   );
 }
+
+function fetchTxHash(nonRevertibleAccumulatedData: PrivateToPublicAccumulatedData): TxHash {
+  // Private kernel functions are executed client side and for this reason tx hash is already set as first nullifier
+  const firstNullifier = nonRevertibleAccumulatedData.nullifiers[0];
+  if (!firstNullifier || firstNullifier.isZero()) {
+    throw new Error(`Cannot get tx hash since first nullifier is missing`);
+  }
+  return new TxHash(firstNullifier.toBuffer());
+}