@aztec/simulator 0.66.0 → 0.67.1-devnet

package/src/public/public_processor.ts

@@ -4,7 +4,6 @@ import {
   type MerkleTreeWriteOperations,
   NestedProcessReturnValues,
   type ProcessedTx,
-  type ProcessedTxHandler,
   Tx,
   TxExecutionPhase,
   type TxValidator,
@@ -13,20 +12,21 @@ import {
 } from '@aztec/circuit-types';
 import {
   type AztecAddress,
+  type BlockHeader,
   type ContractDataSource,
   Fr,
   type GlobalVariables,
-  type Header,
   MAX_NOTE_HASHES_PER_TX,
   MAX_NULLIFIERS_PER_TX,
   NULLIFIER_SUBTREE_HEIGHT,
   PublicDataWrite,
 } from '@aztec/circuits.js';
 import { padArrayEnd } from '@aztec/foundation/collection';
-import { createDebugLogger } from '@aztec/foundation/log';
+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 { 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';
@@ -47,13 +47,19 @@ export class PublicProcessorFactory {
    */
   public create(
     merkleTree: MerkleTreeWriteOperations,
-    maybeHistoricalHeader: Header | undefined,
+    maybeHistoricalHeader: BlockHeader | undefined,
     globalVariables: GlobalVariables,
   ): PublicProcessor {
     const historicalHeader = maybeHistoricalHeader ?? merkleTree.getInitialHeader();
 
     const worldStateDB = new WorldStateDB(merkleTree, this.contractDataSource);
-    const publicTxSimulator = new PublicTxSimulator(merkleTree, worldStateDB, this.telemetryClient, globalVariables);
+    const publicTxSimulator = new PublicTxSimulator(
+      merkleTree,
+      worldStateDB,
+      this.telemetryClient,
+      globalVariables,
+      /*doMerkleOperations=*/ true,
+    );
 
     return new PublicProcessor(
       merkleTree,
@@ -70,16 +76,16 @@ export class PublicProcessorFactory {
  * 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,
     protected globalVariables: GlobalVariables,
-    protected historicalHeader: Header,
+    protected historicalHeader: BlockHeader,
     protected worldStateDB: WorldStateDB,
     protected publicTxSimulator: PublicTxSimulator,
     telemetryClient: TelemetryClient,
-    private log = createDebugLogger('aztec:simulator:public-processor'),
+    private log = createLogger('simulator:public-processor'),
   ) {
     this.metrics = new PublicProcessorMetrics(telemetryClient, 'PublicProcessor');
   }
@@ -97,7 +103,6 @@ export class PublicProcessor {
   public async process(
     txs: Tx[],
     maxTransactions = txs.length,
-    processedTxHandler?: ProcessedTxHandler,
     txValidator?: TxValidator<ProcessedTx>,
   ): Promise<[ProcessedTx[], FailedTx[], NestedProcessReturnValues[]]> {
     // The processor modifies the tx objects in place, so we need to clone them.
@@ -112,66 +117,9 @@ export class PublicProcessor {
         break;
       }
       try {
-        const [processedTx, returnValues] = !tx.hasPublicCalls()
-          ? await this.processPrivateOnlyTx(tx)
-          : await this.processTxWithPublicCalls(tx);
-        this.log.debug(`Processed tx`, {
-          txHash: processedTx.hash,
-          historicalHeaderHash: processedTx.constants.historicalHeader.hash(),
-          blockNumber: processedTx.constants.globalVariables.blockNumber,
-          lastArchiveRoot: processedTx.constants.historicalHeader.lastArchive.root,
-        });
-
-        // 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);
         result.push(processedTx);
-        returns = returns.concat(returnValues ?? []);
+        returns = returns.concat(returnValues);
       } catch (err: any) {
         const errorMessage = err instanceof Error ? err.message : 'Unknown error';
         this.log.warn(`Failed to process tx ${tx.getTxHash()}: ${errorMessage} ${err?.stack}`);
@@ -187,17 +135,89 @@ export class PublicProcessor {
     return [result, failed, returns];
   }
 
+  @trackSpan('PublicProcessor.processTx', tx => ({ [Attributes.TX_HASH]: tx.tryGetTxHash()?.toString() }))
+  private async processTx(
+    tx: Tx,
+    txValidator?: TxValidator<ProcessedTx>,
+  ): Promise<[ProcessedTx, NestedProcessReturnValues[]]> {
+    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`);
+      }
+    }
+    // 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()),
+    );
+
+    return [processedTx, returnValues ?? []];
+  }
+
   /**
-   * Creates the final set of data update requests for the transaction. This includes the
-   * set of public data update requests as returned by the public kernel, plus a data update
-   * request for updating fee balance. It also updates the local public state db.
-   * See build_or_patch_payment_update_request in base_rollup_inputs.nr for more details.
+   * Creates the public data write for paying the tx fee.
+   * This is used in private only txs, since for txs with public calls
+   * the avm handles the fee payment itself.
    */
-  private async getFeePaymentPublicDataWrite(
-    publicDataWrites: PublicDataWrite[],
-    txFee: Fr,
-    feePayer: AztecAddress,
-  ): Promise<PublicDataWrite | undefined> {
+  private async getFeePaymentPublicDataWrite(txFee: Fr, feePayer: AztecAddress): Promise<PublicDataWrite | undefined> {
     if (feePayer.isZero()) {
       this.log.debug(`No one is paying the fee of ${txFee.toBigInt()}`);
       return;
@@ -209,11 +229,7 @@ export class PublicProcessor {
 
     this.log.debug(`Deducting ${txFee.toBigInt()} balance in Fee Juice for ${feePayer}`);
 
-    const existingBalanceWrite = publicDataWrites.find(write => write.leafSlot.equals(leafSlot));
-
-    const balance = existingBalanceWrite
-      ? existingBalanceWrite.value
-      : await this.worldStateDB.storageRead(feeJuiceAddress, balanceSlot);
+    const balance = await this.worldStateDB.storageRead(feeJuiceAddress, balanceSlot);
 
     if (balance.lt(txFee)) {
       throw new Error(
@@ -234,12 +250,7 @@ export class PublicProcessor {
     const gasFees = this.globalVariables.gasFees;
     const transactionFee = tx.data.gasUsed.computeFee(gasFees);
 
-    const accumulatedData = tx.data.forRollup!.end;
-    const feePaymentPublicDataWrite = await this.getFeePaymentPublicDataWrite(
-      accumulatedData.publicDataWrites,
-      transactionFee,
-      tx.data.feePayer,
-    );
+    const feePaymentPublicDataWrite = await this.getFeePaymentPublicDataWrite(transactionFee, tx.data.feePayer);
 
     const processedTx = makeProcessedTxFromPrivateOnlyTx(
       tx,
@@ -247,6 +258,13 @@ export class PublicProcessor {
       feePaymentPublicDataWrite,
       this.globalVariables,
     );
+
+    this.metrics.recordClassRegistration(
+      ...tx.contractClassLogs
+        .unrollLogs()
+        .filter(log => ContractClassRegisteredEvent.isContractClassRegisteredEvent(log.data))
+        .map(log => ContractClassRegisteredEvent.fromLog(log.data)),
+    );
     return [processedTx];
   }
 
@@ -283,21 +301,7 @@ export class PublicProcessor {
     const durationMs = timer.ms();
     this.metrics.recordTx(phaseCount, durationMs);
 
-    const data = avmProvingRequest.inputs.output;
-    const feePaymentPublicDataWrite = await this.getFeePaymentPublicDataWrite(
-      data.accumulatedData.publicDataWrites,
-      data.transactionFee,
-      tx.data.feePayer,
-    );
-
-    const processedTx = makeProcessedTxFromTxWithPublicCalls(
-      tx,
-      avmProvingRequest,
-      feePaymentPublicDataWrite,
-      gasUsed,
-      revertCode,
-      revertReason,
-    );
+    const processedTx = makeProcessedTxFromTxWithPublicCalls(tx, avmProvingRequest, gasUsed, revertCode, revertReason);
 
     const returnValues = processedPhases.find(({ phase }) => phase === TxExecutionPhase.APP_LOGIC)?.returnValues ?? [];
 

package/src/public/public_processor_metrics.ts

@@ -1,5 +1,5 @@
 import { type TxExecutionPhase } from '@aztec/circuit-types';
-import { type ContractClassRegisteredEvent } from '@aztec/protocol-contracts';
+import { type ContractClassRegisteredEvent } from '@aztec/protocol-contracts/class-registerer';
 import {
   Attributes,
   type Histogram,

package/src/public/public_tx_context.ts

@@ -10,31 +10,33 @@ import {
   TxHash,
 } from '@aztec/circuit-types';
 import {
-  AppendOnlyTreeSnapshot,
   AvmCircuitInputs,
   type AvmCircuitPublicInputs,
+  type AztecAddress,
   Fr,
   Gas,
   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 DebugLogger, createDebugLogger } from '@aztec/foundation/log';
+import { type Logger, createLogger } from '@aztec/foundation/log';
 
 import { strict as assert } from 'assert';
 import { inspect } from 'util';
 
 import { AvmPersistableStateManager } from '../avm/index.js';
-import { DualSideEffectTrace } from './dual_side_effect_trace.js';
 import { PublicEnqueuedCallSideEffectTrace, SideEffectArrayLengths } from './enqueued_call_side_effect_trace.js';
 import { type WorldStateDB } from './public_db_sources.js';
-import { PublicSideEffectTrace } from './side_effect_trace.js';
 import { generateAvmCircuitPublicInputs } from './transitional_adapters.js';
 import { getCallRequestsByPhase, getExecutionRequestsByPhase } from './utils.js';
 
@@ -42,10 +44,10 @@ import { getCallRequestsByPhase, getExecutionRequestsByPhase } from './utils.js'
  * The transaction-level context for public execution.
  */
 export class PublicTxContext {
-  private log: DebugLogger;
+  private log: Logger;
 
   /* Gas used including private, teardown gas _limit_, setup and app logic */
-  private gasUsed: Gas;
+  private gasUsedByPublic: Gas = Gas.empty();
   /* Gas actually used during teardown (different from limit) */
   public teardownGasUsed: Gas = Gas.empty();
 
@@ -62,8 +64,9 @@ export class PublicTxContext {
     public readonly state: PhaseStateManager,
     private readonly globalVariables: GlobalVariables,
     private readonly startStateReference: StateReference,
-    private readonly startGasUsed: Gas,
     private readonly gasSettings: GasSettings,
+    private readonly gasUsedByPrivate: Gas,
+    private readonly gasAllocatedToPublic: Gas,
     private readonly setupCallRequests: PublicCallRequest[],
     private readonly appLogicCallRequests: PublicCallRequest[],
     private readonly teardownCallRequests: PublicCallRequest[],
@@ -72,10 +75,10 @@ export class PublicTxContext {
     private readonly teardownExecutionRequests: PublicExecutionRequest[],
     public readonly nonRevertibleAccumulatedDataFromPrivate: PrivateToPublicAccumulatedData,
     public readonly revertibleAccumulatedDataFromPrivate: PrivateToPublicAccumulatedData,
+    public readonly feePayer: AztecAddress,
     public trace: PublicEnqueuedCallSideEffectTrace, // FIXME(dbanks12): should be private
   ) {
-    this.log = createDebugLogger(`aztec:public_tx_context`);
-    this.gasUsed = startGasUsed;
+    this.log = createLogger(`simulator:public_tx_context`);
   }
 
   public static async create(
@@ -87,10 +90,10 @@ export class PublicTxContext {
   ) {
     const nonRevertibleAccumulatedDataFromPrivate = tx.data.forPublic!.nonRevertibleAccumulatedData;
 
-    const innerCallTrace = new PublicSideEffectTrace();
     const previousAccumulatedDataArrayLengths = new SideEffectArrayLengths(
       /*publicDataWrites*/ 0,
-      countAccumulatedItems(nonRevertibleAccumulatedDataFromPrivate.noteHashes),
+      /*protocolPublicDataWrites*/ 0,
+      /*noteHashes*/ 0,
       /*nullifiers=*/ 0,
       countAccumulatedItems(nonRevertibleAccumulatedDataFromPrivate.l2ToL1Msgs),
       /*unencryptedLogsHashes*/ 0,
@@ -99,17 +102,27 @@ export class PublicTxContext {
       /*startSideEffectCounter=*/ 0,
       previousAccumulatedDataArrayLengths,
     );
-    const trace = new DualSideEffectTrace(innerCallTrace, enqueuedCallTrace);
 
     // Transaction level state manager that will be forked for revertible phases.
-    const txStateManager = await AvmPersistableStateManager.create(worldStateDB, trace, doMerkleOperations);
+    const txStateManager = await AvmPersistableStateManager.create(
+      worldStateDB,
+      enqueuedCallTrace,
+      doMerkleOperations,
+      fetchTxHash(nonRevertibleAccumulatedDataFromPrivate),
+    );
+
+    const gasSettings = tx.data.constants.txContext.gasSettings;
+    const gasUsedByPrivate = tx.data.gasUsed;
+    // Gas allocated to public is "whatever's left" after private, but with some max applied.
+    const gasAllocatedToPublic = applyMaxToAvailableGas(gasSettings.gasLimits.sub(gasUsedByPrivate));
 
     return new PublicTxContext(
       new PhaseStateManager(txStateManager),
       globalVariables,
       await db.getStateReference(),
-      tx.data.gasUsed,
-      tx.data.constants.txContext.gasSettings,
+      gasSettings,
+      gasUsedByPrivate,
+      gasAllocatedToPublic,
       getCallRequestsByPhase(tx, TxExecutionPhase.SETUP),
       getCallRequestsByPhase(tx, TxExecutionPhase.APP_LOGIC),
       getCallRequestsByPhase(tx, TxExecutionPhase.TEARDOWN),
@@ -118,6 +131,7 @@ export class PublicTxContext {
       getExecutionRequestsByPhase(tx, TxExecutionPhase.TEARDOWN),
       tx.data.forPublic!.nonRevertibleAccumulatedData,
       tx.data.forPublic!.revertibleAccumulatedData,
+      tx.data.feePayer,
       enqueuedCallTrace,
     );
   }
@@ -179,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);
   }
 
   /**
@@ -230,13 +239,14 @@ export class PublicTxContext {
   }
 
   /**
-   * How much gas is left for the specified phase?
+   * How much gas is left as of the specified phase?
    */
-  getGasLeftForPhase(phase: TxExecutionPhase): Gas {
+  getGasLeftAtPhase(phase: TxExecutionPhase): Gas {
     if (phase === TxExecutionPhase.TEARDOWN) {
-      return this.gasSettings.teardownGasLimits;
+      return applyMaxToAvailableGas(this.gasSettings.teardownGasLimits);
     } else {
-      return this.gasSettings.gasLimits.sub(this.gasUsed);
+      const gasLeftForPublic = this.gasAllocatedToPublic.sub(this.gasUsedByPublic);
+      return gasLeftForPublic;
     }
   }
 
@@ -247,10 +257,18 @@ export class PublicTxContext {
     if (phase === TxExecutionPhase.TEARDOWN) {
       this.teardownGasUsed = this.teardownGasUsed.add(gas);
     } else {
-      this.gasUsed = this.gasUsed.add(gas);
+      this.gasUsedByPublic = this.gasUsedByPublic.add(gas);
     }
   }
 
+  /**
+   * The gasUsed by public and private,
+   * as if the entire teardown gas limit was consumed.
+   */
+  getTotalGasUsed(): Gas {
+    return this.gasUsedByPrivate.add(this.gasUsedByPublic);
+  }
+
   /**
    * Compute the gas used using the actual gas used during teardown instead
    * of the teardown gas limit.
@@ -261,14 +279,7 @@ export class PublicTxContext {
     assert(this.halted, 'Can only compute actual gas used after tx execution ends');
     const requireTeardown = this.teardownCallRequests.length > 0;
     const teardownGasLimits = requireTeardown ? this.gasSettings.teardownGasLimits : Gas.empty();
-    return this.gasUsed.sub(teardownGasLimits).add(this.teardownGasUsed);
-  }
-
-  /**
-   * The gasUsed as if the entire teardown gas limit was consumed.
-   */
-  getGasUsedForFee(): Gas {
-    return this.gasUsed;
+    return this.getTotalGasUsed().sub(teardownGasLimits).add(this.teardownGasUsed);
   }
 
   /**
@@ -288,12 +299,15 @@ export class PublicTxContext {
    * Should only be called during or after teardown.
    */
   private getTransactionFeeUnsafe(): Fr {
-    const txFee = this.gasUsed.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.gasUsed),
+      gasUsed: inspect(gasUsed),
       gasFees: inspect(this.globalVariables.gasFees),
     });
+
     return txFee;
   }
 
@@ -302,16 +316,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,
@@ -324,15 +351,16 @@ export class PublicTxContext {
       this.trace,
       this.globalVariables,
       this.startStateReference,
-      this.startGasUsed,
+      /*startGasUsed=*/ this.gasUsedByPrivate,
       this.gasSettings,
+      this.feePayer,
       this.setupCallRequests,
       this.appLogicCallRequests,
       this.teardownCallRequests,
       this.nonRevertibleAccumulatedDataFromPrivate,
       this.revertibleAccumulatedDataFromPrivate,
       endTreeSnapshots,
-      /*endGasUsed=*/ this.gasUsed,
+      /*endGasUsed=*/ this.getTotalGasUsed(),
       this.getTransactionFeeUnsafe(),
       this.revertCode,
     );
@@ -367,12 +395,12 @@ export class PublicTxContext {
  * transaction level one.
  */
 class PhaseStateManager {
-  private log: DebugLogger;
+  private log: Logger;
 
   private currentlyActiveStateManager: AvmPersistableStateManager | undefined;
 
   constructor(private readonly txStateManager: AvmPersistableStateManager) {
-    this.log = createDebugLogger(`aztec:public_phase_state_manager`);
+    this.log = createLogger(`simulator:public_phase_state_manager`);
   }
 
   fork() {
@@ -405,3 +433,22 @@ class PhaseStateManager {
     this.currentlyActiveStateManager = undefined;
   }
 }
+
+/**
+ * Apply L2 gas maximum.
+ */
+function applyMaxToAvailableGas(availableGas: Gas) {
+  return new Gas(
+    /*daGas=*/ availableGas.daGas,
+    /*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());
+}