@aztec/simulator 0.87.6 → 1.0.0-nightly.20250604

package/src/public/public_processor/guarded_merkle_tree.ts

@@ -0,0 +1,148 @@
+import { SerialQueue } from '@aztec/foundation/queue';
+import type { IndexedTreeLeafPreimage, SiblingPath } from '@aztec/foundation/trees';
+import type {
+  BatchInsertionResult,
+  IndexedTreeId,
+  MerkleTreeId,
+  MerkleTreeLeafType,
+  MerkleTreeWriteOperations,
+  SequentialInsertionResult,
+  TreeInfo,
+} from '@aztec/stdlib/trees';
+import type { BlockHeader, StateReference } from '@aztec/stdlib/tx';
+
+/**
+ * Wraps an instance of `MerkleTreeWriteOperations` to allow the sequencer to gate access.
+ * If transactions execution goes past the deadline, the simulator will continue to execute and update the world state
+ * The public processor however requires that the world state remain constant after the deadline in order to finalise the block
+ * The public processor provides this implementation of MerkleTreeWriteOperations to the simulator
+ */
+
+export class GuardedMerkleTreeOperations implements MerkleTreeWriteOperations {
+  private isStopped = false;
+  private serialQueue = new SerialQueue();
+
+  constructor(private target: MerkleTreeWriteOperations) {
+    this.serialQueue.start();
+  }
+
+  private guard() {
+    if (this.isStopped) {
+      throw new Error('Merkle tree access has been stopped');
+    }
+  }
+
+  // Executes the provided function only if the guard is not stopped.
+  private guardAndPush<T>(fn: () => Promise<T>): Promise<T> {
+    this.guard();
+    return this.serialQueue.put(() => {
+      this.guard();
+      return fn();
+    });
+  }
+
+  public getUnderlyingFork(): MerkleTreeWriteOperations {
+    return this.target;
+  }
+
+  // Stops all further access to the merkle trees via this object
+  async stop(): Promise<void> {
+    await this.serialQueue.put(() => {
+      this.isStopped = true;
+      return Promise.resolve();
+    });
+    return this.serialQueue.end();
+  }
+
+  // Proxy all methods to the target
+  appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<ID>[]): Promise<void> {
+    return this.guardAndPush(() => this.target.appendLeaves(treeId, leaves));
+  }
+
+  updateArchive(header: BlockHeader): Promise<void> {
+    return this.guardAndPush(() => this.target.updateArchive(header));
+  }
+  batchInsert<TreeHeight extends number, SubtreeSiblingPathHeight extends number, ID extends IndexedTreeId>(
+    treeId: ID,
+    leaves: Buffer[],
+    subtreeHeight: number,
+  ): Promise<BatchInsertionResult<TreeHeight, SubtreeSiblingPathHeight>> {
+    return this.guardAndPush(() => this.target.batchInsert(treeId, leaves, subtreeHeight));
+  }
+  sequentialInsert<TreeHeight extends number, ID extends IndexedTreeId>(
+    treeId: ID,
+    leaves: Buffer[],
+  ): Promise<SequentialInsertionResult<TreeHeight>> {
+    return this.guardAndPush(() => this.target.sequentialInsert(treeId, leaves));
+  }
+  close(): Promise<void> {
+    return this.guardAndPush(() => this.target.close());
+  }
+  getTreeInfo(treeId: MerkleTreeId): Promise<TreeInfo> {
+    return this.guardAndPush(() => this.target.getTreeInfo(treeId));
+  }
+  getStateReference(): Promise<StateReference> {
+    return this.guardAndPush(() => this.target.getStateReference());
+  }
+  getInitialHeader(): BlockHeader {
+    return this.target.getInitialHeader();
+  }
+  getSiblingPath<N extends number>(treeId: MerkleTreeId, index: bigint): Promise<SiblingPath<N>> {
+    return this.guardAndPush(() => this.target.getSiblingPath(treeId, index));
+  }
+  getPreviousValueIndex<ID extends IndexedTreeId>(
+    treeId: ID,
+    value: bigint,
+  ): Promise<{ index: bigint; alreadyPresent: boolean } | undefined> {
+    return this.guardAndPush(() => this.target.getPreviousValueIndex(treeId, value));
+  }
+  getLeafPreimage<ID extends IndexedTreeId>(treeId: ID, index: bigint): Promise<IndexedTreeLeafPreimage | undefined> {
+    return this.guardAndPush(() => this.target.getLeafPreimage(treeId, index));
+  }
+  findLeafIndices<ID extends MerkleTreeId>(
+    treeId: ID,
+    values: MerkleTreeLeafType<ID>[],
+  ): Promise<(bigint | undefined)[]> {
+    return this.guardAndPush(() => this.target.findLeafIndices(treeId, values));
+  }
+  findLeafIndicesAfter<ID extends MerkleTreeId>(
+    treeId: ID,
+    values: MerkleTreeLeafType<ID>[],
+    startIndex: bigint,
+  ): Promise<(bigint | undefined)[]> {
+    return this.guardAndPush(() => this.target.findLeafIndicesAfter(treeId, values, startIndex));
+  }
+  getLeafValue<ID extends MerkleTreeId>(
+    treeId: ID,
+    index: bigint,
+  ): Promise<MerkleTreeLeafType<typeof treeId> | undefined> {
+    return this.guardAndPush(() => this.target.getLeafValue(treeId, index));
+  }
+  getBlockNumbersForLeafIndices<ID extends MerkleTreeId>(
+    treeId: ID,
+    leafIndices: bigint[],
+  ): Promise<(bigint | undefined)[]> {
+    return this.guardAndPush(() => this.target.getBlockNumbersForLeafIndices(treeId, leafIndices));
+  }
+  createCheckpoint(): Promise<void> {
+    return this.guardAndPush(() => this.target.createCheckpoint());
+  }
+  commitCheckpoint(): Promise<void> {
+    return this.guardAndPush(() => this.target.commitCheckpoint());
+  }
+  revertCheckpoint(): Promise<void> {
+    return this.guardAndPush(() => this.target.revertCheckpoint());
+  }
+  commitAllCheckpoints(): Promise<void> {
+    return this.guardAndPush(() => this.target.commitAllCheckpoints());
+  }
+  revertAllCheckpoints(): Promise<void> {
+    return this.guardAndPush(() => this.target.revertAllCheckpoints());
+  }
+  findSiblingPaths<ID extends MerkleTreeId, N extends number>(
+    treeId: ID,
+    values: MerkleTreeLeafType<ID>[],
+  ): Promise<(SiblingPath<N> | undefined)[]> {
+    return this.guardAndPush(() => this.target.findSiblingPaths(treeId, values));
+  }
+}

package/src/public/public_processor/public_processor.ts

@@ -10,7 +10,11 @@ import { PublicDataWrite } from '@aztec/stdlib/avm';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import type { ContractDataSource } from '@aztec/stdlib/contract';
 import { Gas } from '@aztec/stdlib/gas';
-import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
+import type {
+  MerkleTreeWriteOperations,
+  PublicProcessorLimits,
+  PublicProcessorValidator,
+} from '@aztec/stdlib/interfaces/server';
 import { MerkleTreeId } from '@aztec/stdlib/trees';
 import {
   type FailedTx,
@@ -35,6 +39,7 @@ import { ForkCheckpoint } from '@aztec/world-state/native';
 
 import { PublicContractsDB, PublicTreesDB } from '../public_db_sources.js';
 import { type PublicTxSimulator, TelemetryPublicTxSimulator } from '../public_tx_simulator/index.js';
+import { GuardedMerkleTreeOperations } from './guarded_merkle_tree.js';
 import { PublicProcessorMetrics } from './public_processor_metrics.js';
 
 /**
@@ -61,8 +66,10 @@ export class PublicProcessorFactory {
     clientInitiatedSimulation: boolean = false,
   ): PublicProcessor {
     const contractsDB = new PublicContractsDB(this.contractDataSource);
+
+    const guardedFork = new GuardedMerkleTreeOperations(merkleTree);
     const publicTxSimulator = this.createPublicTxSimulator(
-      merkleTree,
+      guardedFork,
       contractsDB,
       globalVariables,
       /*doMerkleOperations=*/ true,
@@ -72,7 +79,7 @@ export class PublicProcessorFactory {
 
     return new PublicProcessor(
       globalVariables,
-      merkleTree,
+      guardedFork,
       contractsDB,
       publicTxSimulator,
       this.dateProvider,
@@ -116,7 +123,7 @@ export class PublicProcessor implements Traceable {
 
   constructor(
     protected globalVariables: GlobalVariables,
-    private merkleTree: MerkleTreeWriteOperations,
+    private guardedMerkleTree: GuardedMerkleTreeOperations,
     protected contractsDB: PublicContractsDB,
     protected publicTxSimulator: PublicTxSimulator,
     private dateProvider: DateProvider,
@@ -139,16 +146,8 @@ export class PublicProcessor implements Traceable {
    */
   public async process(
     txs: Iterable<Tx> | AsyncIterable<Tx>,
-    limits: {
-      maxTransactions?: number;
-      maxBlockSize?: number;
-      maxBlockGas?: Gas;
-      deadline?: Date;
-    } = {},
-    validator: {
-      preprocessValidator?: TxValidator<Tx>;
-      nullifierCache?: { addNullifiers: (nullifiers: Buffer[]) => void };
-    } = {},
+    limits: PublicProcessorLimits = {},
+    validator: PublicProcessorValidator = {},
   ): Promise<[ProcessedTx[], FailedTx[], Tx[], NestedProcessReturnValues[]]> {
     const { maxTransactions, maxBlockSize, deadline, maxBlockGas } = limits;
     const { preprocessValidator, nullifierCache } = validator;
@@ -226,7 +225,9 @@ export class PublicProcessor implements Traceable {
       // We checkpoint the transaction here, then within the try/catch we
       // 1. Revert the checkpoint if the tx fails or needs to be discarded for any reason
       // 2. Commit the transaction in the finally block. Note that by using the ForkCheckpoint lifecycle only the first commit/revert takes effect
-      const checkpoint = await ForkCheckpoint.new(this.merkleTree);
+      // By doing this, every transaction starts on a fresh checkpoint and it's state updates only make it to the fork if this checkpoint is committed.
+      // Note: We use the underlying fork here not the guarded one, this ensures that it's not impacted by stopping the guarded version
+      const checkpoint = await ForkCheckpoint.new(this.guardedMerkleTree.getUnderlyingFork());
 
       try {
         const [processedTx, returnValues] = await this.processTx(tx, deadline);
@@ -256,19 +257,37 @@ export class PublicProcessor implements Traceable {
         totalBlockGas = totalBlockGas.add(processedTx.gasUsed.totalGas);
         totalSizeInBytes += txSize;
       } catch (err: any) {
-        // Roll back state to start of TX before proceeding to next TX
-        await checkpoint.revert();
         if (err?.name === 'PublicProcessorTimeoutError') {
           this.log.warn(`Stopping tx processing due to timeout.`);
+          // We hit the transaction execution deadline.
+          // There may still be a transaction executing. We stop the guarded fork to prevent any further access to the world state.
+          await this.guardedMerkleTree.stop();
+
+          // We now know there can't be any further access to world state. The fork is in a state where there is:
+          // 1. At least one outstanding checkpoint that has not been committed (the one created before we processed the tx).
+          // 2. Possible state updates on that checkpoint or any others created during execution.
+
+          // First we revert a checkpoint as managed by the ForkCheckpoint. This will revert whatever is the current checkpoint
+          // which may not be the one originally created by this object. But that is ok, we do this to fulfil the ForkCheckpoint
+          // lifecycle expectations and ensure it doesn't attempt to commit later on.
+          await checkpoint.revert();
+
+          // Now we want to revert any/all remaining checkpoints, destroying any outstanding state updates.
+          // This needs to be done directly on the underlying fork as the guarded fork has been stopped.
+          await this.guardedMerkleTree.getUnderlyingFork().revertAllCheckpoints();
+
+          // We should now be in a position where the fork is in a clean state and no further updates can be made to it.
           break;
         }
+
+        // Roll back state to start of TX before proceeding to next TX
+        await checkpoint.revert();
         const errorMessage = err instanceof Error ? err.message : 'Unknown error';
         this.log.warn(`Failed to process tx ${txHash.toString()}: ${errorMessage} ${err?.stack}`);
-
         failed.push({ tx, error: err instanceof Error ? err : new Error(errorMessage) });
         returns.push(new NestedProcessReturnValues([]));
       } finally {
-        // Base case is we always commit the checkpoint. Using the ForkCheckpoint means this has no effect if the tx was reverted
+        // Base case is we always commit the checkpoint. Using the ForkCheckpoint means this has no effect if the tx was previously reverted
         await checkpoint.commit();
         // The tx-level contracts cache should not live on to the next tx
         this.contractsDB.clearContractsForTx();
@@ -330,12 +349,12 @@ export class PublicProcessor implements Traceable {
     // 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.merkleTree.appendLeaves(
+    await this.guardedMerkleTree.appendLeaves(
       MerkleTreeId.NOTE_HASH_TREE,
       padArrayEnd(processedTx.txEffect.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
     );
     try {
-      await this.merkleTree.batchInsert(
+      await this.guardedMerkleTree.batchInsert(
         MerkleTreeId.NULLIFIER_TREE,
         padArrayEnd(processedTx.txEffect.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX).map(n => n.toBuffer()),
         NULLIFIER_SUBTREE_HEIGHT,
@@ -396,7 +415,7 @@ export class PublicProcessor implements Traceable {
     const balanceSlot = await computeFeePayerBalanceStorageSlot(feePayer);
     const leafSlot = await computeFeePayerBalanceLeafSlot(feePayer);
     // This high-level db is used as a convenient helper. It could be done with the merkleTree directly.
-    const treesDB = new PublicTreesDB(this.merkleTree);
+    const treesDB = new PublicTreesDB(this.guardedMerkleTree);
 
     this.log.debug(`Deducting ${txFee.toBigInt()} balance in Fee Juice for ${feePayer}`);
 
@@ -430,13 +449,9 @@ export class PublicProcessor implements Traceable {
       this.globalVariables,
     );
 
-    const siloedContractClassLogs = await tx.filterContractClassLogs(
-      tx.data.getNonEmptyContractClassLogsHashes(),
-      true,
-    );
-
     this.metrics.recordClassRegistration(
-      ...siloedContractClassLogs
+      ...tx
+        .getContractClassLogs()
         .filter(log => ContractClassRegisteredEvent.isContractClassRegisteredEvent(log))
         .map(log => ContractClassRegisteredEvent.fromLog(log)),
     );
@@ -474,13 +489,11 @@ export class PublicProcessor implements Traceable {
       }
     });
 
-    const siloedContractClassLogs = await tx.filterContractClassLogs(
-      tx.data.getNonEmptyContractClassLogsHashes(),
-      true,
-    );
-
+    const contractClassLogs = revertCode.isOK()
+      ? tx.getContractClassLogs()
+      : tx.getSplitContractClassLogs(false /* revertible */);
     this.metrics.recordClassRegistration(
-      ...siloedContractClassLogs
+      ...contractClassLogs
         .filter(log => ContractClassRegisteredEvent.isContractClassRegisteredEvent(log))
         .map(log => ContractClassRegisteredEvent.fromLog(log)),
     );

package/src/public/public_tx_simulator/public_tx_context.ts

@@ -1,6 +1,5 @@
 import {
   MAX_ENQUEUED_CALLS_PER_TX,
-  MAX_L2_GAS_PER_TX_PUBLIC_PORTION,
   MAX_L2_TO_L1_MSGS_PER_TX,
   MAX_NOTE_HASHES_PER_TX,
   MAX_NULLIFIERS_PER_TX,
@@ -10,7 +9,14 @@ import {
 import { padArrayEnd } from '@aztec/foundation/collection';
 import { Fr } from '@aztec/foundation/fields';
 import { type Logger, createLogger } from '@aztec/foundation/log';
-import { AvmAccumulatedData, AvmCircuitPublicInputs, PublicDataWrite, RevertCode } from '@aztec/stdlib/avm';
+import {
+  AvmAccumulatedData,
+  AvmAccumulatedDataArrayLengths,
+  AvmCircuitPublicInputs,
+  PublicDataWrite,
+  RevertCode,
+  clampGasSettingsForAVM,
+} 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';
@@ -18,8 +24,9 @@ import { Gas, GasSettings } from '@aztec/stdlib/gas';
 import {
   PrivateToAvmAccumulatedData,
   PrivateToAvmAccumulatedDataArrayLengths,
-  type PrivateToPublicAccumulatedData,
+  PrivateToPublicAccumulatedData,
   PublicCallRequest,
+  PublicCallRequestArrayLengths,
   countAccumulatedItems,
 } from '@aztec/stdlib/kernel';
 import { PublicLog } from '@aztec/stdlib/logs';
@@ -66,8 +73,10 @@ export class PublicTxContext {
     private readonly startTreeSnapshots: TreeSnapshots,
     private readonly globalVariables: GlobalVariables,
     private readonly gasSettings: GasSettings,
+    private readonly clampedGasSettings: GasSettings,
     private readonly gasUsedByPrivate: Gas,
     private readonly gasAllocatedToPublic: Gas,
+    private readonly gasAllocatedToPublicTeardown: Gas,
     private readonly setupCallRequests: PublicCallRequestWithCalldata[],
     private readonly appLogicCallRequests: PublicCallRequestWithCalldata[],
     private readonly teardownCallRequests: PublicCallRequestWithCalldata[],
@@ -105,7 +114,9 @@ export class PublicTxContext {
     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));
+    const clampedGasSettings = clampGasSettingsForAVM(gasSettings, gasUsedByPrivate);
+    const gasAllocatedToPublic = clampedGasSettings.gasLimits.sub(gasUsedByPrivate);
+    const gasAllocatedToPublicTeardown = clampedGasSettings.teardownGasLimits;
 
     return new PublicTxContext(
       await tx.getTxHash(),
@@ -113,8 +124,10 @@ export class PublicTxContext {
       await txStateManager.getTreeSnapshots(),
       globalVariables,
       gasSettings,
+      clampedGasSettings,
       gasUsedByPrivate,
       gasAllocatedToPublic,
+      gasAllocatedToPublicTeardown,
       getCallRequestsWithCalldataByPhase(tx, TxExecutionPhase.SETUP),
       getCallRequestsWithCalldataByPhase(tx, TxExecutionPhase.APP_LOGIC),
       getCallRequestsWithCalldataByPhase(tx, TxExecutionPhase.TEARDOWN),
@@ -210,7 +223,8 @@ export class PublicTxContext {
    */
   getGasLeftAtPhase(phase: TxExecutionPhase): Gas {
     if (phase === TxExecutionPhase.TEARDOWN) {
-      return applyMaxToAvailableGas(this.gasSettings.teardownGasLimits);
+      const gasLeftForPublicTeardown = this.gasAllocatedToPublicTeardown.sub(this.teardownGasUsed);
+      return gasLeftForPublicTeardown;
     } else {
       const gasLeftForPublic = this.gasAllocatedToPublic.sub(this.gasUsedByPublic);
       return gasLeftForPublic;
@@ -356,19 +370,32 @@ export class PublicTxContext {
     // This converts the private accumulated data to the avm accumulated data format.
     const convertAccumulatedData = (from: PrivateToPublicAccumulatedData) =>
       new PrivateToAvmAccumulatedData(from.noteHashes, from.nullifiers, from.l2ToL1Msgs);
-    const getArrayLengths = (from: PrivateToPublicAccumulatedData) =>
+    const getPreviousAccumulatedDataArrayLengths = (from: PrivateToPublicAccumulatedData) =>
       new PrivateToAvmAccumulatedDataArrayLengths(
         countAccumulatedItems(from.noteHashes),
         countAccumulatedItems(from.nullifiers),
         countAccumulatedItems(from.l2ToL1Msgs),
       );
+    const getAvmAccumulatedDataArrayLengths = (from: AvmAccumulatedData) =>
+      new AvmAccumulatedDataArrayLengths(
+        from.noteHashes.length,
+        from.nullifiers.length,
+        from.l2ToL1Msgs.length,
+        from.publicLogs.length,
+        from.publicDataWrites.length,
+      );
 
     return new AvmCircuitPublicInputs(
       this.globalVariables,
       this.startTreeSnapshots,
       /*startGasUsed=*/ this.gasUsedByPrivate,
-      this.gasSettings,
+      this.clampedGasSettings,
       this.feePayer,
+      /*publicCallRequestArrayLengths=*/ new PublicCallRequestArrayLengths(
+        this.setupCallRequests.length,
+        this.appLogicCallRequests.length,
+        this.teardownCallRequests.length > 0,
+      ),
       /*publicSetupCallRequests=*/ padArrayEnd(
         this.setupCallRequests.map(r => r.request),
         PublicCallRequest.empty(),
@@ -382,12 +409,13 @@ export class PublicTxContext {
       /*publicTeardownCallRequests=*/ this.teardownCallRequests.length > 0
         ? this.teardownCallRequests[0].request
         : PublicCallRequest.empty(),
-      getArrayLengths(this.nonRevertibleAccumulatedDataFromPrivate),
-      getArrayLengths(this.revertibleAccumulatedDataFromPrivate),
+      getPreviousAccumulatedDataArrayLengths(this.nonRevertibleAccumulatedDataFromPrivate),
+      getPreviousAccumulatedDataArrayLengths(this.revertibleAccumulatedDataFromPrivate),
       convertAccumulatedData(this.nonRevertibleAccumulatedDataFromPrivate),
       convertAccumulatedData(this.revertibleAccumulatedDataFromPrivate),
       endTreeSnapshots,
       this.getTotalGasUsed(),
+      getAvmAccumulatedDataArrayLengths(accumulatedData),
       accumulatedData,
       /*transactionFee=*/ this.getTransactionFeeUnsafe(),
       /*isReverted=*/ !this.revertCode.isOK(),
@@ -444,13 +472,3 @@ 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),
-  );
-}

package/src/public/public_tx_simulator/public_tx_simulator.ts

@@ -121,9 +121,6 @@ export class PublicTxSimulator {
 
       const revertCode = context.getFinalRevertCode();
 
-      if (!revertCode.isOK()) {
-        await tx.filterRevertedLogs();
-      }
       // Commit contracts from this TX to the block-level cache and clear tx cache
       // If the tx reverted, only commit non-revertible contracts
       // NOTE: You can't create contracts in public, so this is only relevant for private-created contracts

package/src/public/side_effect_trace.ts

@@ -82,6 +82,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
     private readonly previousSideEffectArrayLengths: SideEffectArrayLengths = SideEffectArrayLengths.empty(),
     /** We need to track the set of class IDs used, to enforce limits. */
     private uniqueClassIds: UniqueClassIds = new UniqueClassIds(),
+    private writtenPublicDataSlots: Set<string> = new Set(),
   ) {
     this.sideEffectCounter = startSideEffectCounter;
   }
@@ -98,6 +99,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
         this.previousSideEffectArrayLengths.publicLogs + this.publicLogs.length,
       ),
       this.uniqueClassIds.fork(),
+      new Set(this.writtenPublicDataSlots),
     );
   }
 
@@ -111,6 +113,8 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
 
     this.sideEffectCounter = forkedTrace.sideEffectCounter;
     this.uniqueClassIds.acceptAndMerge(forkedTrace.uniqueClassIds);
+    // Accept even if reverted, since the user already paid for the writes
+    this.writtenPublicDataSlots = new Set(forkedTrace.writtenPublicDataSlots);
 
     if (!reverted) {
       this.publicDataWrites.push(...forkedTrace.publicDataWrites);
@@ -170,6 +174,15 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
       `Traced public data write (address=${contractAddress}, slot=${slot}): value=${value} (counter=${this.sideEffectCounter}, isProtocol:${protocolWrite})`,
     );
     this.incrementSideEffectCounter();
+    this.writtenPublicDataSlots.add(this.computePublicDataSlotKey(contractAddress, slot));
+  }
+
+  private computePublicDataSlotKey(contractAddress: AztecAddress, slot: Fr): string {
+    return `${contractAddress.toString()}:${slot.toString()}`;
+  }
+
+  public isStorageCold(contractAddress: AztecAddress, slot: Fr): boolean {
+    return !this.writtenPublicDataSlots.has(this.computePublicDataSlotKey(contractAddress, slot));
   }
 
   public traceNewNoteHash(noteHash: Fr) {

package/src/public/side_effect_trace_interface.ts

@@ -12,6 +12,7 @@ export interface PublicSideEffectTraceInterface {
     value: Fr,
     protocolWrite: boolean,
   ): Promise<void>;
+  isStorageCold(contractAddress: AztecAddress, slot: Fr): boolean;
   traceNewNoteHash(uniqueNoteHash: Fr): void;
   getNoteHashCount(): number;
   traceNewNullifier(siloedNullifier: Fr): void;

package/src/public/state_manager/state_manager.ts

@@ -148,6 +148,10 @@ export class PublicPersistableStateManager {
     await this.trace.tracePublicStorageWrite(contractAddress, slot, value, protocolWrite);
   }
 
+  public isStorageCold(contractAddress: AztecAddress, slot: Fr): boolean {
+    return this.trace.isStorageCold(contractAddress, slot);
+  }
+
   /**
    * Read from public storage.
    *