Skip to content

Commit 5daac75

Browse files
committed
feat: public storage overrides for simulation via stateOverrides
Adds a `stateOverrides?: StateOverrides` parameter to `simulatePublicCalls` and threads it through aztec.js → wallet-sdk → PXE → node. The new shape is modeled after Foundry / Geth `eth_call` state overrides — a single bundled argument that can be extended with new override flavors (e.g. contract class/instance shadows, nullifier overrides) without changing the RPC. This first cut populates only `publicStorage`: ```ts type StateOverrides = { publicStorage?: PublicStorageOverride[]; }; type PublicStorageOverride = { contract: AztecAddress; slot: Fr; // raw storage slot, not leafSlot — derived internally value: Fr; }; ``` The node opens an ephemeral world-state fork per simulation and writes each override via `applyPublicDataOverrides(fork, publicStorage)` (computes `computePublicDataTreeLeafSlot(contract, slot)` then `sequentialInsert`). Writes never reach committed state. Includes the e2e wiring fix in TestWallet.simulateViaEntrypoint to forward the option through to PXE so e2e tests can exercise the full chain.
1 parent 0a2a38c commit 5daac75

15 files changed

Lines changed: 276 additions & 10 deletions

File tree

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import { Fr } from '@aztec/foundation/curves/bn254';
2+
import { PublicDataWrite } from '@aztec/stdlib/avm';
3+
import { AztecAddress } from '@aztec/stdlib/aztec-address';
4+
import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
5+
import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
6+
import { MerkleTreeId } from '@aztec/stdlib/trees';
7+
8+
import { jest } from '@jest/globals';
9+
import { type MockProxy, mock } from 'jest-mock-extended';
10+
11+
import { applyPublicDataOverrides } from './public_data_overrides.js';
12+
13+
describe('applyPublicDataOverrides', () => {
14+
let fork: MockProxy<MerkleTreeWriteOperations>;
15+
16+
beforeEach(() => {
17+
fork = mock<MerkleTreeWriteOperations>();
18+
fork.sequentialInsert.mockResolvedValue({ lowLeavesWitnesses: [], insertionWitnesses: [] } as any);
19+
});
20+
21+
it('does nothing when overrides is undefined', async () => {
22+
await applyPublicDataOverrides(fork, undefined);
23+
expect(fork.sequentialInsert).not.toHaveBeenCalled();
24+
});
25+
26+
it('does nothing when overrides is empty', async () => {
27+
await applyPublicDataOverrides(fork, []);
28+
expect(fork.sequentialInsert).not.toHaveBeenCalled();
29+
});
30+
31+
it('inserts a single override at the correct leaf slot', async () => {
32+
const contract = await AztecAddress.random();
33+
const slot = Fr.random();
34+
const value = Fr.random();
35+
36+
await applyPublicDataOverrides(fork, [{ contract, slot, value }]);
37+
38+
const expectedLeafSlot = await computePublicDataTreeLeafSlot(contract, slot);
39+
const expectedWrite = new PublicDataWrite(expectedLeafSlot, value);
40+
41+
expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
42+
expect(fork.sequentialInsert).toHaveBeenCalledWith(MerkleTreeId.PUBLIC_DATA_TREE, [expectedWrite.toBuffer()]);
43+
});
44+
45+
it('inserts multiple overrides in a single batch call', async () => {
46+
const contract = await AztecAddress.random();
47+
const overrides = [
48+
{ contract, slot: Fr.random(), value: Fr.random() },
49+
{ contract, slot: Fr.random(), value: Fr.random() },
50+
];
51+
52+
await applyPublicDataOverrides(fork, overrides);
53+
54+
const expectedWrites = await Promise.all(
55+
overrides.map(async o => {
56+
const leafSlot = await computePublicDataTreeLeafSlot(o.contract, o.slot);
57+
return new PublicDataWrite(leafSlot, o.value);
58+
}),
59+
);
60+
61+
expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
62+
expect(fork.sequentialInsert).toHaveBeenCalledWith(
63+
MerkleTreeId.PUBLIC_DATA_TREE,
64+
expectedWrites.map(w => w.toBuffer()),
65+
);
66+
});
67+
68+
it('passes duplicate (contract, slot) writes through — last write wins via tree semantics', async () => {
69+
const contract = await AztecAddress.random();
70+
const slot = Fr.random();
71+
const firstValue = Fr.random();
72+
const secondValue = Fr.random();
73+
74+
await applyPublicDataOverrides(fork, [
75+
{ contract, slot, value: firstValue },
76+
{ contract, slot, value: secondValue },
77+
]);
78+
79+
// Both writes are passed to sequentialInsert; the tree handles last-wins ordering.
80+
expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
81+
const [treeId, leaves] = fork.sequentialInsert.mock.calls[0] as [MerkleTreeId, Buffer[]];
82+
expect(treeId).toBe(MerkleTreeId.PUBLIC_DATA_TREE);
83+
expect(leaves).toHaveLength(2);
84+
});
85+
});
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import { PublicDataWrite } from '@aztec/stdlib/avm';
2+
import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
3+
import type { PublicStorageOverride } from '@aztec/stdlib/interfaces/client';
4+
import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
5+
import { MerkleTreeId } from '@aztec/stdlib/trees';
6+
7+
/**
8+
* Injects public-state overrides into an ephemeral world-state fork before simulation.
9+
*
10+
* Each override is written via the same `sequentialInsert` path the public processor
11+
* uses during real transaction execution, so low-leaf updates and root coherence are
12+
* handled identically for both simulation and proof generation.
13+
*
14+
* The fork is ephemeral — these writes never reach the committed world state.
15+
*/
16+
export async function applyPublicDataOverrides(
17+
fork: MerkleTreeWriteOperations,
18+
publicStorage: PublicStorageOverride[] | undefined,
19+
): Promise<void> {
20+
if (!publicStorage?.length) {
21+
return;
22+
}
23+
24+
const writes = await Promise.all(
25+
publicStorage.map(async o => {
26+
const leafSlot = await computePublicDataTreeLeafSlot(o.contract, o.slot);
27+
return new PublicDataWrite(leafSlot, o.value);
28+
}),
29+
);
30+
31+
await fork.sequentialInsert(
32+
MerkleTreeId.PUBLIC_DATA_TREE,
33+
writes.map(w => w.toBuffer()),
34+
);
35+
}

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ import type {
8080
CheckpointResponse,
8181
GetContractClassLogsResponse,
8282
GetPublicLogsResponse,
83+
StateOverrides,
8384
} from '@aztec/stdlib/interfaces/client';
8485
import { AztecNodeAdminConfigSchema } from '@aztec/stdlib/interfaces/client';
8586
import {
@@ -143,6 +144,7 @@ import {
143144
} from './block_response_helpers.js';
144145
import { type AztecNodeConfig, createKeyStoreForValidator } from './config.js';
145146
import { NodeMetrics } from './node_metrics.js';
147+
import { applyPublicDataOverrides } from './public_data_overrides.js';
146148

147149
/**
148150
* The aztec node.
@@ -1481,11 +1483,17 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
14811483
/**
14821484
* Simulates the public part of a transaction with the current state.
14831485
* @param tx - The transaction to simulate.
1486+
* @param skipFeeEnforcement - If true, fee enforcement is skipped.
1487+
* @param stateOverrides - Optional state overrides applied to the ephemeral fork and contract DB before simulation.
14841488
**/
14851489
@trackSpan('AztecNodeService.simulatePublicCalls', (tx: Tx) => ({
14861490
[Attributes.TX_HASH]: tx.getTxHash().toString(),
14871491
}))
1488-
public async simulatePublicCalls(tx: Tx, skipFeeEnforcement = false): Promise<PublicSimulationOutput> {
1492+
public async simulatePublicCalls(
1493+
tx: Tx,
1494+
skipFeeEnforcement = false,
1495+
stateOverrides?: StateOverrides,
1496+
): Promise<PublicSimulationOutput> {
14891497
// Check total gas limit for simulation
14901498
const gasSettings = tx.data.constants.txContext.gasSettings;
14911499
const txGasLimit = gasSettings.gasLimits.l2Gas;
@@ -1530,6 +1538,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
15301538
await this.worldStateSynchronizer.syncImmediate(latestBlockNumber);
15311539
const merkleTreeFork = await this.worldStateSynchronizer.fork();
15321540
try {
1541+
await applyPublicDataOverrides(merkleTreeFork, stateOverrides?.publicStorage);
15331542
const config = PublicSimulatorConfig.from({
15341543
skipFeeEnforcement,
15351544
collectDebugLogs: true,

yarn-project/aztec.js/src/api/wallet.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,4 +78,11 @@ export { AccountManager } from '../wallet/account_manager.js';
7878

7979
export { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
8080

81+
export {
82+
type PublicStorageOverride,
83+
PublicStorageOverrideSchema,
84+
type StateOverrides,
85+
StateOverridesSchema,
86+
} from '@aztec/stdlib/interfaces/client';
87+
8188
export { type DeployAccountOptions, DeployAccountMethod } from '../wallet/deploy_account_method.js';

yarn-project/aztec.js/src/contract/interaction_options.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import type { FieldsOf } from '@aztec/foundation/types';
33
import type { AuthWitness } from '@aztec/stdlib/auth-witness';
44
import { AztecAddress } from '@aztec/stdlib/aztec-address';
55
import type { GasSettings, ManaUsageEstimate } from '@aztec/stdlib/gas';
6+
import type { StateOverrides } from '@aztec/stdlib/interfaces/client';
67
import {
78
type Capsule,
89
OFFCHAIN_MESSAGE_IDENTIFIER,
@@ -157,6 +158,8 @@ export type SimulateInteractionOptions = Omit<SendInteractionOptions, 'fee'> & {
157158
/** Whether to include metadata such as performance statistics (e.g. timing information of the different circuits and oracles) and gas estimation
158159
* in the simulation result, in addition to the return value and offchain effects */
159160
includeMetadata?: boolean;
161+
/** Pre-simulation state overrides applied to the ephemeral fork and contract DB. */
162+
stateOverrides?: StateOverrides;
160163
};
161164

162165
/**

yarn-project/aztec.js/src/wallet/wallet.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import { AuthWitness } from '@aztec/stdlib/auth-witness';
1313
import type { AztecAddress } from '@aztec/stdlib/aztec-address';
1414
import { type ContractInstanceWithAddress, ContractInstanceWithAddressSchema } from '@aztec/stdlib/contract';
1515
import { Gas, ManaUsageEstimate } from '@aztec/stdlib/gas';
16+
import { StateOverridesSchema } from '@aztec/stdlib/interfaces/client';
1617
import { LogId } from '@aztec/stdlib/logs';
1718
import { AbiDecodedSchema, type ApiSchemaFor, optional, schemas, zodFor } from '@aztec/stdlib/schemas';
1819
import type { ExecutionPayload, InTx } from '@aztec/stdlib/tx';
@@ -335,6 +336,7 @@ export const SimulateOptionsSchema = z.object({
335336
skipFeeEnforcement: optional(z.boolean()),
336337
includeMetadata: optional(z.boolean()),
337338
additionalScopes: optional(z.array(schemas.AztecAddress)),
339+
stateOverrides: optional(StateOverridesSchema),
338340
});
339341

340342
export const ProfileOptionsSchema = SimulateOptionsSchema.extend({

yarn-project/end-to-end/src/e2e_avm_simulator.test.ts

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
import { AztecAddress } from '@aztec/aztec.js/addresses';
22
import { BatchCall, type ContractInstanceWithAddress } from '@aztec/aztec.js/contracts';
33
import { Fr } from '@aztec/aztec.js/fields';
4+
import type { AztecNode } from '@aztec/aztec.js/node';
45
import { TxExecutionResult } from '@aztec/aztec.js/tx';
6+
import type { PublicStorageOverride } from '@aztec/aztec.js/wallet';
57
import type { Wallet } from '@aztec/aztec.js/wallet';
68
import { AvmInitializerTestContract } from '@aztec/noir-test-contracts.js/AvmInitializerTest';
79
import { AvmTestContract } from '@aztec/noir-test-contracts.js/AvmTest';
@@ -16,13 +18,15 @@ describe('e2e_avm_simulator', () => {
1618
jest.setTimeout(TIMEOUT);
1719

1820
let wallet: Wallet;
21+
let aztecNode: AztecNode;
1922
let defaultAccountAddress: AztecAddress;
2023
let teardown: () => Promise<void>;
2124

2225
beforeAll(async () => {
2326
({
2427
teardown,
2528
wallet,
29+
aztecNode,
2630
accounts: [defaultAccountAddress],
2731
} = await setup(1));
2832
await ensureAccountContractsPublished(wallet, [defaultAccountAddress]);
@@ -249,6 +253,51 @@ describe('e2e_avm_simulator', () => {
249253
});
250254
});
251255

256+
describe('publicDataOverrides', () => {
257+
// AvmTestContract: `single` is the first storage variable and lives at raw slot 1.
258+
const SINGLE_SLOT = new Fr(1n);
259+
let avmContract: AvmTestContract;
260+
261+
beforeEach(async () => {
262+
({ contract: avmContract } = await AvmTestContract.deploy(wallet).send({ from: defaultAccountAddress }));
263+
});
264+
265+
it('simulated read of an unwritten slot returns the override; real storage is untouched', async () => {
266+
const overrideValue = new Fr(0xdeadbeefn);
267+
const publicStorage: PublicStorageOverride[] = [
268+
{ contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue },
269+
];
270+
271+
const simResult = await avmContract.methods
272+
.read_storage_single()
273+
.simulate({ from: defaultAccountAddress, stateOverrides: { publicStorage } });
274+
expect(simResult.result).toEqual(overrideValue.toBigInt());
275+
276+
// Real state is untouched — the slot was never written.
277+
const realValue = await aztecNode.getPublicStorageAt('latest', avmContract.address, SINGLE_SLOT);
278+
expect(realValue.toBigInt()).toEqual(0n);
279+
});
280+
281+
it('simulated read returns the override when a slot was previously written by a real tx', async () => {
282+
const realValue = new Fr(100n);
283+
await avmContract.methods.set_storage_single(realValue).send({ from: defaultAccountAddress });
284+
285+
const overrideValue = new Fr(999n);
286+
const publicStorage: PublicStorageOverride[] = [
287+
{ contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue },
288+
];
289+
290+
const simResult = await avmContract.methods
291+
.read_storage_single()
292+
.simulate({ from: defaultAccountAddress, stateOverrides: { publicStorage } });
293+
expect(simResult.result).toEqual(overrideValue.toBigInt());
294+
295+
// Real storage still holds the original written value.
296+
const storedValue = await aztecNode.getPublicStorageAt('latest', avmContract.address, SINGLE_SLOT);
297+
expect(storedValue.toBigInt()).toEqual(realValue.toBigInt());
298+
});
299+
});
300+
252301
describe('AvmInitializerTestContract', () => {
253302
let avmContract: AvmInitializerTestContract;
254303

yarn-project/end-to-end/src/test-wallet/test_wallet.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,8 @@ export class TestWallet extends BaseWallet {
260260
executionPayload: ExecutionPayload,
261261
opts: SimulateViaEntrypointOptions,
262262
): Promise<TxSimulationResultWithAppOffset> {
263-
const { from, feeOptions, additionalScopes, skipTxValidation, skipFeeEnforcement, sendMessagesAs } = opts;
263+
const { from, feeOptions, additionalScopes, skipTxValidation, skipFeeEnforcement, sendMessagesAs, stateOverrides } =
264+
opts;
264265
const scopes = this.scopesFrom(from, additionalScopes);
265266
const skipKernels = this.simulationMode !== 'full';
266267
const useOverride = this.simulationMode === 'kernelless-override';
@@ -311,6 +312,7 @@ export class TestWallet extends BaseWallet {
311312
overrides,
312313
scopes,
313314
senderForTags: this.senderForTagsFrom(from, sendMessagesAs),
315+
stateOverrides,
314316
});
315317
const appCallOffset = await this.computeAppCallOffset(from, feeOptions);
316318
return TxSimulationResultWithAppOffset.fromResultAndOffset(result, appCallOffset);

yarn-project/pxe/src/pxe.ts

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import {
2727
getContractClassFromArtifact,
2828
} from '@aztec/stdlib/contract';
2929
import { SimulationError } from '@aztec/stdlib/errors';
30-
import type { AztecNode, PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
30+
import type { AztecNode, PrivateKernelProver, StateOverrides } from '@aztec/stdlib/interfaces/client';
3131
import type {
3232
PrivateExecutionStep,
3333
PrivateKernelExecutionProofOutput,
@@ -125,6 +125,8 @@ export type SimulateTxOpts = {
125125
scopes: AztecAddress[];
126126
/** Sender address used to derive discovery tags for private messages (notes, events, logs) this tx emits. */
127127
senderForTags?: AztecAddress;
128+
/** Pre-simulation state overrides applied to the ephemeral fork and contract DB. */
129+
stateOverrides?: StateOverrides;
128130
};
129131

130132
/** Options for PXE.executeUtility. */
@@ -467,11 +469,11 @@ export class PXE {
467469
* It can also be used for estimating gas in the future.
468470
* @param tx - The transaction to be simulated.
469471
*/
470-
async #simulatePublicCalls(tx: Tx, skipFeeEnforcement: boolean) {
472+
async #simulatePublicCalls(tx: Tx, skipFeeEnforcement: boolean, stateOverrides?: StateOverrides) {
471473
// Simulating public calls can throw if the TX fails in a phase that doesn't allow reverts (setup)
472474
// Or return as reverted if it fails in a phase that allows reverts (app logic, teardown)
473475
try {
474-
const result = await this.node.simulatePublicCalls(tx, skipFeeEnforcement);
476+
const result = await this.node.simulatePublicCalls(tx, skipFeeEnforcement, stateOverrides);
475477
if (result.revertReason) {
476478
throw result.revertReason;
477479
}
@@ -956,6 +958,7 @@ export class PXE {
956958
overrides,
957959
scopes,
958960
senderForTags,
961+
stateOverrides,
959962
}: SimulateTxOpts,
960963
): Promise<TxSimulationResult> {
961964
// We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
@@ -1037,7 +1040,7 @@ export class PXE {
10371040
let publicOutput: PublicSimulationOutput | undefined;
10381041
if (simulatePublic && publicInputs.forPublic) {
10391042
const publicSimulationTimer = new Timer();
1040-
publicOutput = await this.#simulatePublicCalls(simulatedTx, skipFeeEnforcement);
1043+
publicOutput = await this.#simulatePublicCalls(simulatedTx, skipFeeEnforcement, stateOverrides);
10411044
publicSimulationTime = publicSimulationTimer.ms();
10421045
if (publicOutput?.debugLogs?.length) {
10431046
await displayDebugLogs(publicOutput.debugLogs, addr => this.contractStore.getDebugContractName(addr));

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ import {
7979
type GetPublicLogsResponse,
8080
GetPublicLogsResponseSchema,
8181
} from './get_logs_response.js';
82+
import { type StateOverrides, StateOverridesSchema } from './state_overrides.js';
8283
import { type WorldStateSyncStatus, WorldStateSyncStatusSchema } from './world_state.js';
8384

8485
/**
@@ -465,8 +466,15 @@ export interface AztecNode {
465466
* Simulates the public part of a transaction with the current state.
466467
* This currently just checks that the transaction execution succeeds.
467468
* @param tx - The transaction to simulate.
469+
* @param skipFeeEnforcement - If true, fee enforcement is skipped.
470+
* @param stateOverrides - Optional state overrides applied to the ephemeral world-state fork and
471+
* contract DB before simulation. Bundles public-storage writes and contract class/instance shadows.
468472
**/
469-
simulatePublicCalls(tx: Tx, skipFeeEnforcement?: boolean): Promise<PublicSimulationOutput>;
473+
simulatePublicCalls(
474+
tx: Tx,
475+
skipFeeEnforcement?: boolean,
476+
stateOverrides?: StateOverrides,
477+
): Promise<PublicSimulationOutput>;
470478

471479
/**
472480
* Returns true if the transaction is valid for inclusion at the current state. Valid transactions can be
@@ -669,7 +677,10 @@ export const AztecNodeApiSchema: ApiSchemaFor<AztecNode> = {
669677
.args(schemas.EthAddress, optional(schemas.SlotNumber), optional(schemas.SlotNumber))
670678
.returns(SingleValidatorStatsSchema.optional()),
671679

672-
simulatePublicCalls: z.function().args(Tx.schema, optional(z.boolean())).returns(PublicSimulationOutput.schema),
680+
simulatePublicCalls: z
681+
.function()
682+
.args(Tx.schema, optional(z.boolean()), optional(StateOverridesSchema))
683+
.returns(PublicSimulationOutput.schema),
673684

674685
isValidTx: z
675686
.function()

0 commit comments

Comments
 (0)