@aztec/simulator 0.76.4 → 0.77.0-testnet-ignition.17

package/src/public/public_tx_context.ts

@@ -1,48 +1,52 @@
 import {
-  type AvmProvingRequest,
-  MerkleTreeId,
-  type MerkleTreeReadOperations,
-  ProvingRequestType,
-  type PublicExecutionRequest,
-  type SimulationError,
-  type Tx,
-  TxExecutionPhase,
-  type TxHash,
-} from '@aztec/circuit-types';
-import {
-  AvmCircuitInputs,
-  type AvmCircuitPublicInputs,
-  type AztecAddress,
-  Fr,
-  Gas,
-  type GasSettings,
-  type GlobalVariables,
   MAX_L2_GAS_PER_TX_PUBLIC_PORTION,
   MAX_L2_TO_L1_MSGS_PER_TX,
   MAX_NOTE_HASHES_PER_TX,
   MAX_NULLIFIERS_PER_TX,
   MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
+  NULLIFIER_SUBTREE_HEIGHT,
+} from '@aztec/constants';
+import { padArrayEnd } from '@aztec/foundation/collection';
+import { Fr } from '@aztec/foundation/fields';
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import { assertLength } from '@aztec/foundation/serialize';
+import {
+  AvmCircuitInputs,
+  type AvmCircuitPublicInputs,
+  type AvmProvingRequest,
+  PublicDataWrite,
+  RevertCode,
+} from '@aztec/stdlib/avm';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { SimulationError } from '@aztec/stdlib/errors';
+import { computeTransactionFee } from '@aztec/stdlib/fees';
+import { Gas, GasSettings } from '@aztec/stdlib/gas';
+import type { MerkleTreeReadOperations } from '@aztec/stdlib/interfaces/server';
+import {
   PrivateToAvmAccumulatedData,
   PrivateToAvmAccumulatedDataArrayLengths,
   type PrivateToPublicAccumulatedData,
   PublicCallRequest,
-  PublicDataWrite,
-  RevertCode,
-  type StateReference,
-  TreeSnapshots,
-  computeTransactionFee,
   countAccumulatedItems,
   mergeAccumulatedData,
-} from '@aztec/circuits.js';
-import { padArrayEnd } from '@aztec/foundation/collection';
-import { type Logger, createLogger } from '@aztec/foundation/log';
-import { assertLength } from '@aztec/foundation/serialize';
+} from '@aztec/stdlib/kernel';
+import { ProvingRequestType } from '@aztec/stdlib/proofs';
+import { MerkleTreeId } from '@aztec/stdlib/trees';
+import {
+  type GlobalVariables,
+  type PublicExecutionRequest,
+  type StateReference,
+  TreeSnapshots,
+  type Tx,
+  TxExecutionPhase,
+  type TxHash,
+} from '@aztec/stdlib/tx';
 
 import { strict as assert } from 'assert';
 import { inspect } from 'util';
 
-import { AvmPersistableStateManager } from '../avm/index.js';
-import { type WorldStateDB } from './public_db_sources.js';
+import { AvmPersistableStateManager } from './avm/index.js';
+import type { WorldStateDB } from './public_db_sources.js';
 import { SideEffectArrayLengths, SideEffectTrace } from './side_effect_trace.js';
 import { getCallRequestsByPhase, getExecutionRequestsByPhase } from './utils.js';
 
@@ -111,12 +115,7 @@ export class PublicTxContext {
     const firstNullifier = nonRevertibleAccumulatedDataFromPrivate.nullifiers[0];
 
     // Transaction level state manager that will be forked for revertible phases.
-    const txStateManager = await AvmPersistableStateManager.create(
-      worldStateDB,
-      trace,
-      doMerkleOperations,
-      firstNullifier,
-    );
+    const txStateManager = AvmPersistableStateManager.create(worldStateDB, trace, doMerkleOperations, firstNullifier);
 
     const gasSettings = tx.data.constants.txContext.gasSettings;
     const gasUsedByPrivate = tx.data.gasUsed;
@@ -149,9 +148,9 @@ export class PublicTxContext {
    * All phases have been processed.
    * Actual transaction fee and actual total consumed gas can now be queried.
    */
-  halt() {
+  async halt() {
     if (this.state.isForked()) {
-      this.state.mergeForkedState();
+      await this.state.mergeForkedState();
     }
     this.halted = true;
   }
@@ -325,37 +324,7 @@ export class PublicTxContext {
    */
   private async generateAvmCircuitPublicInputs(endStateReference: StateReference): Promise<AvmCircuitPublicInputs> {
     assert(this.halted, 'Can only get AvmCircuitPublicInputs after tx execution ends');
-    const ephemeralTrees = this.state.getActiveStateManager().merkleTrees;
-
-    const noteHashTree = await ephemeralTrees.getTreeSnapshot(MerkleTreeId.NOTE_HASH_TREE);
-    const nullifierTree = await ephemeralTrees.getTreeSnapshot(MerkleTreeId.NULLIFIER_TREE);
-    const publicDataTree = await 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,
-      noteHashTree,
-      nullifierTree,
-      publicDataTree,
-    );
+    const stateManager = this.state.getActiveStateManager();
 
     const startTreeSnapshots = new TreeSnapshots(
       this.startStateReference.l1ToL2MessageTree,
@@ -364,6 +333,14 @@ export class PublicTxContext {
       this.startStateReference.partial.publicDataTree,
     );
 
+    // Will be patched/padded at the end of this fn
+    const endTreeSnapshots = new TreeSnapshots(
+      endStateReference.l1ToL2MessageTree,
+      endStateReference.partial.noteHashTree,
+      endStateReference.partial.nullifierTree,
+      endStateReference.partial.publicDataTree,
+    );
+
     const avmCircuitPublicInputs = this.trace.toAvmCircuitPublicInputs(
       this.globalVariables,
       startTreeSnapshots,
@@ -427,6 +404,24 @@ export class PublicTxContext {
       PublicDataWrite.empty(),
       MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
     );
+    const numNoteHashesToPad =
+      MAX_NOTE_HASHES_PER_TX - countAccumulatedItems(avmCircuitPublicInputs.accumulatedData.noteHashes);
+    await stateManager.db.appendLeaves(MerkleTreeId.NOTE_HASH_TREE, padArrayEnd([], Fr.ZERO, numNoteHashesToPad));
+    const numNullifiersToPad =
+      MAX_NULLIFIERS_PER_TX - countAccumulatedItems(avmCircuitPublicInputs.accumulatedData.nullifiers);
+    await stateManager.db.batchInsert(
+      MerkleTreeId.NULLIFIER_TREE,
+      padArrayEnd([], Fr.ZERO, numNullifiersToPad).map(nullifier => nullifier.toBuffer()),
+      NULLIFIER_SUBTREE_HEIGHT,
+    );
+
+    const paddedState = await stateManager.db.getStateReference();
+    avmCircuitPublicInputs.endTreeSnapshots = new TreeSnapshots(
+      paddedState.l1ToL2MessageTree,
+      paddedState.partial.noteHashTree,
+      paddedState.partial.nullifierTree,
+      paddedState.partial.publicDataTree,
+    );
 
     return avmCircuitPublicInputs;
   }
@@ -467,10 +462,10 @@ class PhaseStateManager {
     this.log = createLogger(`simulator:public_phase_state_manager`);
   }
 
-  fork() {
+  async fork() {
     assert(!this.currentlyActiveStateManager, 'Cannot fork when already forked');
     this.log.debug(`Forking phase state manager`);
-    this.currentlyActiveStateManager = this.txStateManager.fork();
+    this.currentlyActiveStateManager = await this.txStateManager.fork();
   }
 
   getActiveStateManager() {
@@ -481,18 +476,18 @@ class PhaseStateManager {
     return !!this.currentlyActiveStateManager;
   }
 
-  mergeForkedState() {
+  async mergeForkedState() {
     assert(this.currentlyActiveStateManager, 'No forked state to merge');
     this.log.debug(`Merging in forked state`);
-    this.txStateManager.merge(this.currentlyActiveStateManager!);
+    await this.txStateManager.merge(this.currentlyActiveStateManager!);
     // Drop the forked state manager now that it is merged
     this.currentlyActiveStateManager = undefined;
   }
 
-  discardForkedState() {
+  async discardForkedState() {
     this.log.debug(`Discarding forked state`);
     assert(this.currentlyActiveStateManager, 'No forked state to discard');
-    this.txStateManager.reject(this.currentlyActiveStateManager!);
+    await this.txStateManager.reject(this.currentlyActiveStateManager!);
     // Drop the forked state manager. We don't want it!
     this.currentlyActiveStateManager = undefined;
   }

package/src/public/public_tx_simulator.ts

@@ -1,29 +1,31 @@
-import {
-  type AvmProvingRequest,
-  type GasUsed,
-  type MerkleTreeReadOperations,
-  NestedProcessReturnValues,
-  type PublicExecutionRequest,
-  type SimulationError,
-  type Tx,
-  TxExecutionPhase,
-} from '@aztec/circuit-types';
-import { type AvmSimulationStats } from '@aztec/circuit-types/stats';
-import { type Fr, type Gas, type GlobalVariables, type PublicCallRequest, type RevertCode } from '@aztec/circuits.js';
+import type { Fr } from '@aztec/foundation/fields';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
 import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import { computeFeePayerBalanceStorageSlot } from '@aztec/protocol-contracts/fee-juice';
+import type { AvmProvingRequest, RevertCode } from '@aztec/stdlib/avm';
+import type { SimulationError } from '@aztec/stdlib/errors';
+import type { Gas, GasUsed } from '@aztec/stdlib/gas';
+import type { MerkleTreeReadOperations } from '@aztec/stdlib/interfaces/server';
+import type { PublicCallRequest } from '@aztec/stdlib/kernel';
+import type { AvmSimulationStats } from '@aztec/stdlib/stats';
+import {
+  type GlobalVariables,
+  NestedProcessReturnValues,
+  PublicExecutionRequest,
+  Tx,
+  TxExecutionPhase,
+} from '@aztec/stdlib/tx';
 import { Attributes, type TelemetryClient, type Tracer, getTelemetryClient, trackSpan } from '@aztec/telemetry-client';
 
 import { strict as assert } from 'assert';
 
-import { type AvmFinalizedCallResult } from '../avm/avm_contract_call_result.js';
-import { type AvmPersistableStateManager, AvmSimulator } from '../avm/index.js';
-import { NullifierCollisionError } from '../avm/journal/nullifiers.js';
 import { getPublicFunctionDebugName } from '../common/debug_fn_name.js';
+import type { AvmFinalizedCallResult } from './avm/avm_contract_call_result.js';
+import { type AvmPersistableStateManager, AvmSimulator } from './avm/index.js';
+import { NullifierCollisionError } from './avm/journal/nullifiers.js';
 import { ExecutorMetrics } from './executor_metrics.js';
-import { computeFeePayerBalanceStorageSlot } from './fee_payment.js';
-import { type WorldStateDB } from './public_db_sources.js';
+import type { WorldStateDB } from './public_db_sources.js';
 import { PublicTxContext } from './public_tx_context.js';
 
 export type ProcessedPhase = {
@@ -54,7 +56,7 @@ export class PublicTxSimulator {
     private worldStateDB: WorldStateDB,
     private globalVariables: GlobalVariables,
     private doMerkleOperations: boolean = false,
-    private enforceFeePayment: boolean = true,
+    private skipFeeEnforcement: boolean = false,
     telemetryClient: TelemetryClient = getTelemetryClient(),
   ) {
     this.log = createLogger(`simulator:public_tx_simulator`);
@@ -70,6 +72,8 @@ export class PublicTxSimulator {
    * @returns The result of the transaction's public execution.
    */
   public async simulate(tx: Tx): Promise<PublicTxResult> {
+    const startTime = process.hrtime.bigint();
+
     const txHash = await tx.getTxHash();
     this.log.debug(`Simulating ${tx.enqueuedPublicFunctionCalls.length} public calls for tx ${txHash}`, { txHash });
 
@@ -114,7 +118,7 @@ export class PublicTxSimulator {
       processedPhases.push(teardownResult);
     }
 
-    context.halt();
+    await context.halt();
     await this.payFee(context);
 
     const endStateReference = await this.db.getStateReference();
@@ -129,15 +133,19 @@ export class PublicTxSimulator {
       // FIXME: we shouldn't need to directly modify worldStateDb here!
       await this.worldStateDB.removeNewContracts(tx, true);
       // FIXME(dbanks12): should not be changing immutable tx
-      tx.filterRevertedLogs(tx.data.forPublic!.nonRevertibleAccumulatedData);
+      await tx.filterRevertedLogs();
     }
 
+    const endTime = process.hrtime.bigint();
+    this.log.debug(`Public TX simulator took ${Number(endTime - startTime) / 1_000_000} ms\n`);
+
     return {
       avmProvingRequest,
       gasUsed: {
         totalGas: context.getActualGasUsed(),
         teardownGas: context.teardownGasUsed,
         publicGas: context.getActualPublicGasUsed(),
+        billedGas: context.getTotalGasUsed(),
       },
       revertCode,
       revertReason: context.revertReason,
@@ -166,11 +174,11 @@ export class PublicTxSimulator {
 
     if (result.reverted) {
       // Drop the currently active forked state manager and rollback to end of setup.
-      context.state.discardForkedState();
+      await context.state.discardForkedState();
     } else {
       if (!context.hasPhase(TxExecutionPhase.TEARDOWN)) {
         // Nothing to do after this (no teardown), so merge state updates now instead of letting teardown handle it.
-        context.state.mergeForkedState();
+        await context.state.mergeForkedState();
       }
     }
 
@@ -186,17 +194,17 @@ export class PublicTxSimulator {
     if (!context.state.isForked()) {
       // If state isn't forked (app logic reverted), fork now
       // so we can rollback to the end of setup if teardown reverts.
-      context.state.fork();
+      await context.state.fork();
     }
 
     const result = await this.simulatePhase(TxExecutionPhase.TEARDOWN, context);
 
     if (result.reverted) {
       // Drop the currently active forked state manager and rollback to end of setup.
-      context.state.discardForkedState();
+      await context.state.discardForkedState();
     } else {
       // Merge state updates from teardown,
-      context.state.mergeForkedState();
+      await context.state.mergeForkedState();
     }
 
     return result;
@@ -395,7 +403,7 @@ export class PublicTxSimulator {
    */
   public async insertRevertiblesFromPrivate(context: PublicTxContext) {
     // Fork the state manager so we can rollback to end of setup if app logic reverts.
-    context.state.fork();
+    await context.state.fork();
     const stateManager = context.state.getActiveStateManager();
     try {
       await stateManager.writeSiloedNullifiersFromPrivate(context.revertibleAccumulatedDataFromPrivate.nullifiers);
@@ -433,7 +441,7 @@ export class PublicTxSimulator {
     // When mocking the balance of the fee payer, the circuit should not be able to prove the simulation
 
     if (currentBalance.lt(txFee)) {
-      if (this.enforceFeePayment) {
+      if (!this.skipFeeEnforcement) {
         throw new Error(
           `Not enough balance for fee payer to pay for transaction (got ${currentBalance.toBigInt()} needs ${txFee.toBigInt()})`,
         );

package/src/public/side_effect_trace.ts

@@ -1,23 +1,5 @@
 import {
-  AvmAccumulatedData,
-  AvmAppendTreeHint,
-  AvmCircuitPublicInputs,
-  AvmContractBytecodeHints,
-  AvmContractInstanceHint,
-  AvmEnqueuedCallHint,
-  AvmExecutionHints,
-  AvmNullifierReadTreeHint,
-  AvmNullifierWriteTreeHint,
-  AvmPublicDataReadTreeHint,
-  AvmPublicDataWriteTreeHint,
-  type AztecAddress,
-  type ContractClassIdPreimage,
-  EthAddress,
-  type Gas,
-  type GasSettings,
-  type GlobalVariables,
   L1_TO_L2_MSG_TREE_HEIGHT,
-  L2ToL1Message,
   MAX_ENQUEUED_CALLS_PER_TX,
   MAX_L2_TO_L1_MSGS_PER_TX,
   MAX_NOTE_HASHES_PER_TX,
@@ -28,33 +10,49 @@ import {
   MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
   NOTE_HASH_TREE_HEIGHT,
   NULLIFIER_TREE_HEIGHT,
-  NoteHash,
-  Nullifier,
-  NullifierLeafPreimage,
   PROTOCOL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
   PUBLIC_DATA_TREE_HEIGHT,
   PUBLIC_LOG_DATA_SIZE_IN_FIELDS,
-  PrivateToAvmAccumulatedData,
-  PrivateToAvmAccumulatedDataArrayLengths,
-  PublicCallRequest,
-  PublicDataTreeLeafPreimage,
-  PublicDataUpdateRequest,
-  PublicDataWrite,
-  PublicLog,
-  ScopedL2ToL1Message,
-  SerializableContractInstance,
-  type TreeSnapshots,
-} from '@aztec/circuits.js';
-import { computePublicDataTreeLeafSlot } from '@aztec/circuits.js/hash';
+} from '@aztec/constants';
 import { padArrayEnd } from '@aztec/foundation/collection';
+import { EthAddress } from '@aztec/foundation/eth-address';
 import { Fr } from '@aztec/foundation/fields';
-import { jsonStringify } from '@aztec/foundation/json-rpc';
 import { createLogger } from '@aztec/foundation/log';
+import {
+  AvmAccumulatedData,
+  AvmAppendTreeHint,
+  AvmCircuitPublicInputs,
+  AvmContractClassHint,
+  AvmContractInstanceHint,
+  AvmEnqueuedCallHint,
+  AvmExecutionHints,
+  AvmNullifierReadTreeHint,
+  AvmNullifierWriteTreeHint,
+  AvmPublicDataReadTreeHint,
+  AvmPublicDataWriteTreeHint,
+  PublicDataUpdateRequest,
+  PublicDataWrite,
+} from '@aztec/stdlib/avm';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { type ContractClassWithCommitment, SerializableContractInstance } from '@aztec/stdlib/contract';
+import type { Gas, GasSettings } from '@aztec/stdlib/gas';
+import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
+import {
+  NoteHash,
+  Nullifier,
+  PrivateToAvmAccumulatedData,
+  PrivateToAvmAccumulatedDataArrayLengths,
+  PublicCallRequest,
+} from '@aztec/stdlib/kernel';
+import { PublicLog } from '@aztec/stdlib/logs';
+import { L2ToL1Message, ScopedL2ToL1Message } from '@aztec/stdlib/messaging';
+import { NullifierLeafPreimage, PublicDataTreeLeafPreimage } from '@aztec/stdlib/trees';
+import type { GlobalVariables, TreeSnapshots } from '@aztec/stdlib/tx';
 
 import { strict as assert } from 'assert';
 
 import { SideEffectLimitReachedError } from './side_effect_errors.js';
-import { type PublicSideEffectTraceInterface } from './side_effect_trace_interface.js';
+import type { PublicSideEffectTraceInterface } from './side_effect_trace_interface.js';
 import { UniqueClassIds } from './unique_class_ids.js';
 
 const emptyPublicDataPath = () => new Array(PUBLIC_DATA_TREE_HEIGHT).fill(Fr.zero());
@@ -124,8 +122,8 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
      *  otherwise the public kernel can fail to prove because TX limits are breached.
      */
     private readonly previousSideEffectArrayLengths: SideEffectArrayLengths = SideEffectArrayLengths.empty(),
-    /** We need to track the set of class IDs used for bytecode retrieval to deduplicate and enforce limits. */
-    private gotBytecodeFromClassIds: UniqueClassIds = new UniqueClassIds(),
+    /** We need to track the set of class IDs used, to enforce limits. */
+    private uniqueClassIds: UniqueClassIds = new UniqueClassIds(),
   ) {
     this.sideEffectCounter = startSideEffectCounter;
     this.avmCircuitHints = AvmExecutionHints.empty();
@@ -142,7 +140,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
         this.previousSideEffectArrayLengths.l2ToL1Msgs + this.l2ToL1Messages.length,
         this.previousSideEffectArrayLengths.publicLogs + this.publicLogs.length,
       ),
-      this.gotBytecodeFromClassIds.fork(),
+      this.uniqueClassIds.fork(),
     );
   }
 
@@ -168,22 +166,10 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
   }
 
   private mergeHints(forkedTrace: this) {
-    this.gotBytecodeFromClassIds.acceptAndMerge(forkedTrace.gotBytecodeFromClassIds);
-
+    this.uniqueClassIds.acceptAndMerge(forkedTrace.uniqueClassIds);
     this.avmCircuitHints.enqueuedCalls.items.push(...forkedTrace.avmCircuitHints.enqueuedCalls.items);
-
     this.avmCircuitHints.contractInstances.items.push(...forkedTrace.avmCircuitHints.contractInstances.items);
-
-    // merge in contract bytecode hints
-    // UniqueClassIds should prevent duplication
-    for (const [contractClassId, bytecodeHint] of forkedTrace.avmCircuitHints.contractBytecodeHints) {
-      assert(
-        !this.avmCircuitHints.contractBytecodeHints.has(contractClassId),
-        'Bug preventing duplication of contract bytecode hints',
-      );
-      this.avmCircuitHints.contractBytecodeHints.set(contractClassId, bytecodeHint);
-    }
-
+    this.avmCircuitHints.contractClasses.items.push(...forkedTrace.avmCircuitHints.contractClasses.items);
     this.avmCircuitHints.publicDataReads.items.push(...forkedTrace.avmCircuitHints.publicDataReads.items);
     this.avmCircuitHints.publicDataWrites.items.push(...forkedTrace.avmCircuitHints.publicDataWrites.items);
     this.avmCircuitHints.nullifierReads.items.push(...forkedTrace.avmCircuitHints.nullifierReads.items);
@@ -371,114 +357,55 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
     contractAddress: AztecAddress,
     exists: boolean,
     instance: SerializableContractInstance = SerializableContractInstance.default(),
-    lowLeafPreimage: NullifierLeafPreimage = NullifierLeafPreimage.empty(),
-    lowLeafIndex: Fr = Fr.zero(),
-    lowLeafPath: Fr[] = emptyNullifierPath(),
+    updateMembershipHint: AvmPublicDataReadTreeHint = AvmPublicDataReadTreeHint.empty(),
+    updatePreimage: Fr[] = [],
   ) {
-    const membershipHint = new AvmNullifierReadTreeHint(lowLeafPreimage, lowLeafIndex, lowLeafPath);
     this.avmCircuitHints.contractInstances.items.push(
       new AvmContractInstanceHint(
         contractAddress,
         exists,
         instance.salt,
         instance.deployer,
-        instance.contractClassId,
+        instance.currentContractClassId,
+        instance.originalContractClassId,
         instance.initializationHash,
         instance.publicKeys,
-        membershipHint,
+        updateMembershipHint,
+        updatePreimage,
       ),
     );
     this.log.trace(`Tracing contract instance retrieval (counter=${this.sideEffectCounter})`);
     this.incrementSideEffectCounter();
   }
 
-  // This tracing function gets called everytime we start simulation/execution.
-  // This happens both when starting a new top-level trace and the start of every forked trace
-  // We use this to collect the AvmContractBytecodeHints
-  public traceGetBytecode(
-    contractAddress: AztecAddress,
-    exists: boolean,
-    bytecode: Buffer = Buffer.alloc(0),
-    contractInstance: SerializableContractInstance = SerializableContractInstance.default(),
-    contractClass: ContractClassIdPreimage = {
-      artifactHash: Fr.zero(),
-      privateFunctionsRoot: Fr.zero(),
-      publicBytecodeCommitment: Fr.zero(),
-    },
-    lowLeafPreimage: NullifierLeafPreimage = NullifierLeafPreimage.empty(),
-    lowLeafIndex: Fr = Fr.zero(),
-    lowLeafPath: Fr[] = emptyNullifierPath(),
-  ) {
-    // FIXME: The way we are hinting contract bytecodes is fundamentally broken.
-    // We are mapping contract class ID to a bytecode hint
-    // But a bytecode hint is tied to a contract INSTANCE.
-    // What if you encounter another contract instance with the same class ID?
-    // We can't hint that instance too since there is already an entry in the hints set that class ID.
-    // But without that instance hinted, the circuit can't prove that the called contract address
-    // actually corresponds to any class ID.
-
-    const membershipHint = new AvmNullifierReadTreeHint(lowLeafPreimage, lowLeafIndex, lowLeafPath);
-    const instance = new AvmContractInstanceHint(
-      contractAddress,
-      exists,
-      contractInstance.salt,
-      contractInstance.deployer,
-      contractInstance.contractClassId,
-      contractInstance.initializationHash,
-      contractInstance.publicKeys,
-      membershipHint,
-    );
-
-    // Always hint the contract instance separately from the bytecode hint.
-    // Since the bytecode hints are keyed by class ID, we need to hint the instance separately
-    // since there might be multiple instances hinted for the same class ID.
-    this.avmCircuitHints.contractInstances.items.push(instance);
-    this.log.trace(
-      `Tracing contract instance for bytecode retrieval: exists=${exists}, instance=${jsonStringify(contractInstance)}`,
-    );
-
+  public traceGetContractClass(contractClassId: Fr, exists: boolean, contractClass?: ContractClassWithCommitment) {
     if (!exists) {
-      // this ensures there are no duplicates
-      this.log.debug(`Contract address ${contractAddress} does not exist. Not tracing bytecode & class ID.`);
-      return;
-    }
-    // We already hinted this bytecode. No need to
-    // Don't we still need to hint if the class ID already exists?
-    // Because the circuit needs to prove that the called contract address corresponds to the class ID.
-    // To do so, the circuit needs to know the class ID in the
-    if (this.gotBytecodeFromClassIds.has(contractInstance.contractClassId.toString())) {
-      // this ensures there are no duplicates
-      this.log.trace(
-        `Contract class id ${contractInstance.contractClassId.toString()} already exists in previous hints`,
+      this.avmCircuitHints.contractClasses.items.push(
+        new AvmContractClassHint(contractClassId, exists, Fr.zero(), Fr.zero(), Fr.zero(), Buffer.alloc(0)),
       );
-      return;
-    }
+    } else if (!this.uniqueClassIds.has(contractClassId.toString())) {
+      if (this.uniqueClassIds.size() >= MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS) {
+        this.log.debug(`Bytecode retrieval failure for contract class ID ${contractClassId} (limit reached)`);
+        throw new SideEffectLimitReachedError(
+          'contract calls to unique class IDs',
+          MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS,
+        );
+      }
 
-    // If we could actually allow contract calls after the limit was reached, we would hint even if we have
-    // surpassed the limit of unique class IDs (still trace the failed bytecode retrieval)
-    // because the circuit needs to know the class ID to know when the limit is hit.
-    // BUT, the issue with this approach is that the sequencer could lie and say "this call was to a new class ID",
-    // and the circuit cannot prove that it's not true without deriving the class ID from bytecode,
-    // proving that it corresponds to the called contract address, and proving that the class ID wasn't already
-    // present/used. That would require more bytecode hashing which is exactly what this limit exists to avoid.
-    if (this.gotBytecodeFromClassIds.size() >= MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS) {
-      this.log.debug(
-        `Bytecode retrieval failure for contract class ID ${contractInstance.contractClassId.toString()} (limit reached)`,
+      this.uniqueClassIds.add(contractClassId.toString());
+      this.avmCircuitHints.contractClasses.items.push(
+        new AvmContractClassHint(
+          contractClassId,
+          exists,
+          contractClass!.artifactHash,
+          contractClass!.privateFunctionsRoot,
+          contractClass!.publicBytecodeCommitment,
+          contractClass!.packedBytecode,
+        ),
       );
-      throw new SideEffectLimitReachedError(
-        'contract calls to unique class IDs',
-        MAX_PUBLIC_CALLS_TO_UNIQUE_CONTRACT_CLASS_IDS,
-      );
-    }
 
-    this.log.trace(`Tracing bytecode & contract class for bytecode retrieval: class=${jsonStringify(contractClass)}`);
-    this.avmCircuitHints.contractBytecodeHints.set(
-      contractInstance.contractClassId.toString(),
-      new AvmContractBytecodeHints(bytecode, instance, contractClass),
-    );
-    // After adding the bytecode hint, mark the classId as retrieved to avoid duplication.
-    // The above map alone isn't sufficient because we need to check the parent trace's (and its parent) as well.
-    this.gotBytecodeFromClassIds.add(contractInstance.contractClassId.toString());
+      this.incrementSideEffectCounter();
+    }
   }
 
   /**