feat: public data tree overrides for simulation (core)

yarn-project/aztec-node/src/aztec-node/public_data_overrides.test.ts

@@ -0,0 +1,85 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { PublicDataWrite } from '@aztec/stdlib/avm';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
+import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
+import { MerkleTreeId } from '@aztec/stdlib/trees';
+
+import { jest } from '@jest/globals';
+import { type MockProxy, mock } from 'jest-mock-extended';
+
+import { applyPublicDataOverrides } from './public_data_overrides.js';
+
+describe('applyPublicDataOverrides', () => {
+  let fork: MockProxy<MerkleTreeWriteOperations>;
+
+  beforeEach(() => {
+    fork = mock<MerkleTreeWriteOperations>();
+    fork.sequentialInsert.mockResolvedValue({ lowLeavesWitnesses: [], insertionWitnesses: [] } as any);
+  });
+
+  it('does nothing when overrides is undefined', async () => {
+    await applyPublicDataOverrides(fork, undefined);
+    expect(fork.sequentialInsert).not.toHaveBeenCalled();
+  });
+
+  it('does nothing when overrides is empty', async () => {
+    await applyPublicDataOverrides(fork, []);
+    expect(fork.sequentialInsert).not.toHaveBeenCalled();
+  });
+
+  it('inserts a single override at the correct leaf slot', async () => {
+    const contract = await AztecAddress.random();
+    const slot = Fr.random();
+    const value = Fr.random();
+
+    await applyPublicDataOverrides(fork, [{ contract, slot, value }]);
+
+    const expectedLeafSlot = await computePublicDataTreeLeafSlot(contract, slot);
+    const expectedWrite = new PublicDataWrite(expectedLeafSlot, value);
+
+    expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
+    expect(fork.sequentialInsert).toHaveBeenCalledWith(MerkleTreeId.PUBLIC_DATA_TREE, [expectedWrite.toBuffer()]);
+  });
+
+  it('inserts multiple overrides in a single batch call', async () => {
+    const contract = await AztecAddress.random();
+    const overrides = [
+      { contract, slot: Fr.random(), value: Fr.random() },
+      { contract, slot: Fr.random(), value: Fr.random() },
+    ];
+
+    await applyPublicDataOverrides(fork, overrides);
+
+    const expectedWrites = await Promise.all(
+      overrides.map(async o => {
+        const leafSlot = await computePublicDataTreeLeafSlot(o.contract, o.slot);
+        return new PublicDataWrite(leafSlot, o.value);
+      }),
+    );
+
+    expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
+    expect(fork.sequentialInsert).toHaveBeenCalledWith(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      expectedWrites.map(w => w.toBuffer()),
+    );
+  });
+
+  it('passes duplicate (contract, slot) writes through — last write wins via tree semantics', async () => {
+    const contract = await AztecAddress.random();
+    const slot = Fr.random();
+    const firstValue = Fr.random();
+    const secondValue = Fr.random();
+
+    await applyPublicDataOverrides(fork, [
+      { contract, slot, value: firstValue },
+      { contract, slot, value: secondValue },
+    ]);
+
+    // Both writes are passed to sequentialInsert; the tree handles last-wins ordering.
+    expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
+    const [treeId, leaves] = fork.sequentialInsert.mock.calls[0] as [MerkleTreeId, Buffer[]];
+    expect(treeId).toBe(MerkleTreeId.PUBLIC_DATA_TREE);
+    expect(leaves).toHaveLength(2);
+  });
+});

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

@@ -0,0 +1,35 @@
+import { PublicDataWrite } from '@aztec/stdlib/avm';
+import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
+import type { PublicDataTreeOverride } from '@aztec/stdlib/interfaces/client';
+import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
+import { MerkleTreeId } from '@aztec/stdlib/trees';
+
+/**
+ * Injects public-state overrides into an ephemeral world-state fork before simulation.
+ *
+ * Each override is written via the same `sequentialInsert` path the public processor
+ * uses during real transaction execution, so low-leaf updates and root coherence are
+ * handled identically for both simulation and proof generation.
+ *
+ * The fork is ephemeral — these writes never reach the committed world state.
+ */
+export async function applyPublicDataOverrides(
+  fork: MerkleTreeWriteOperations,
+  publicDataOverrides: PublicDataTreeOverride[] | undefined,
+): Promise<void> {
+  if (!publicDataOverrides?.length) {
+    return;
+  }
+
+  const writes = await Promise.all(
+    publicDataOverrides.map(async o => {
+      const leafSlot = await computePublicDataTreeLeafSlot(o.contract, o.slot);
+      return new PublicDataWrite(leafSlot, o.value);
+    }),
+  );
+
+  await fork.sequentialInsert(
+    MerkleTreeId.PUBLIC_DATA_TREE,
+    writes.map(w => w.toBuffer()),
+  );
+}

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

@@ -75,6 +75,7 @@ import {
   type AztecNodeDebug,
   type GetContractClassLogsResponse,
   type GetPublicLogsResponse,
+  type PublicDataTreeOverride,
 } from '@aztec/stdlib/interfaces/client';
 import {
   type AllowedElement,
@@ -131,6 +132,7 @@ import { createSentinel } from '../sentinel/factory.js';
 import { Sentinel } from '../sentinel/sentinel.js';
 import { type AztecNodeConfig, createKeyStoreForValidator } from './config.js';
 import { NodeMetrics } from './node_metrics.js';
+import { applyPublicDataOverrides } from './public_data_overrides.js';
 
 /**
  * The aztec node.
@@ -1312,11 +1314,17 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
   /**
    * Simulates the public part of a transaction with the current state.
    * @param tx - The transaction to simulate.
+   * @param skipFeeEnforcement - If true, fee enforcement is skipped.
+   * @param publicDataOverrides - Optional public state overrides injected into the ephemeral fork before simulation.
    **/
   @trackSpan('AztecNodeService.simulatePublicCalls', (tx: Tx) => ({
     [Attributes.TX_HASH]: tx.getTxHash().toString(),
   }))
-  public async simulatePublicCalls(tx: Tx, skipFeeEnforcement = false): Promise<PublicSimulationOutput> {
+  public async simulatePublicCalls(
+    tx: Tx,
+    skipFeeEnforcement = false,
+    publicDataOverrides?: PublicDataTreeOverride[],
+  ): Promise<PublicSimulationOutput> {
     // Check total gas limit for simulation
     const gasSettings = tx.data.constants.txContext.gasSettings;
     const txGasLimit = gasSettings.gasLimits.l2Gas;
@@ -1361,6 +1369,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
     await this.worldStateSynchronizer.syncImmediate(latestBlockNumber);
     const merkleTreeFork = await this.worldStateSynchronizer.fork();
     try {
+      await applyPublicDataOverrides(merkleTreeFork, publicDataOverrides);
       const config = PublicSimulatorConfig.from({
         skipFeeEnforcement,
         collectDebugLogs: true,

yarn-project/stdlib/src/interfaces/aztec-node.ts

@@ -67,6 +67,7 @@ import {
   type GetPublicLogsResponse,
   GetPublicLogsResponseSchema,
 } from './get_logs_response.js';
+import { type PublicDataTreeOverride, PublicDataTreeOverrideSchema } from './public_data_tree_override.js';
 import { type WorldStateSyncStatus, WorldStateSyncStatusSchema } from './world_state.js';
 
 /**
@@ -462,8 +463,14 @@ export interface AztecNode
    * Simulates the public part of a transaction with the current state.
    * This currently just checks that the transaction execution succeeds.
    * @param tx - The transaction to simulate.
+   * @param skipFeeEnforcement - If true, fee enforcement is skipped.
+   * @param publicDataOverrides - Optional public state overrides injected into the ephemeral world-state fork before simulation.
    **/
-  simulatePublicCalls(tx: Tx, skipFeeEnforcement?: boolean): Promise<PublicSimulationOutput>;
+  simulatePublicCalls(
+    tx: Tx,
+    skipFeeEnforcement?: boolean,
+    publicDataOverrides?: PublicDataTreeOverride[],
+  ): Promise<PublicSimulationOutput>;
 
   /**
    * Returns true if the transaction is valid for inclusion at the current state. Valid transactions can be
@@ -659,7 +666,10 @@ export const AztecNodeApiSchema: ApiSchemaFor<AztecNode> = {
     .args(schemas.EthAddress, optional(schemas.SlotNumber), optional(schemas.SlotNumber))
     .returns(SingleValidatorStatsSchema.optional()),
 
-  simulatePublicCalls: z.function().args(Tx.schema, optional(z.boolean())).returns(PublicSimulationOutput.schema),
+  simulatePublicCalls: z
+    .function()
+    .args(Tx.schema, optional(z.boolean()), optional(z.array(PublicDataTreeOverrideSchema)))
+    .returns(PublicSimulationOutput.schema),
 
   isValidTx: z
     .function()

yarn-project/stdlib/src/interfaces/client.ts

@@ -4,3 +4,4 @@ export * from './aztec-node-debug.js';
 export * from './private_kernel_prover.js';
 export * from './get_logs_response.js';
 export * from './api_limit.js';
+export * from './public_data_tree_override.js';

yarn-project/stdlib/src/interfaces/public_data_tree_override.ts

@@ -0,0 +1,22 @@
+import type { Fr } from '@aztec/foundation/curves/bn254';
+
+import { z } from 'zod';
+
+import type { AztecAddress } from '../aztec-address/index.js';
+import { schemas } from '../schemas/schemas.js';
+
+/** A single public-state override to inject into a world-state fork before simulation. */
+export type PublicDataTreeOverride = {
+  /** Contract that owns the storage slot. */
+  contract: AztecAddress;
+  /** Raw storage slot within the contract (not yet hashed into a tree key). */
+  slot: Fr;
+  /** Value to place at that slot for the duration of the simulation. */
+  value: Fr;
+};
+
+export const PublicDataTreeOverrideSchema = z.object({
+  contract: schemas.AztecAddress,
+  slot: schemas.Fr,
+  value: schemas.Fr,
+});