From c199c88098f37f1fef3b097aec15fdf88fd9d863 Mon Sep 17 00:00:00 2001 From: dbanks12 Date: Fri, 1 May 2026 15:34:41 +0000 Subject: [PATCH 1/2] feat: public storage overrides for simulation via SimulationOverrides MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends `SimulationOverrides` with a `publicStorage` field, plumbed through `.simulate({ overrides })` on aztec.js → wallet → PXE → `AztecNode.simulatePublicCalls`. Each entry writes a `(contract, slot, value)` into the public-data tree of the ephemeral world-state fork before the tx runs; real chain state is untouched. --- .../docs/aztec-js/how_to_test.md | 16 ++++ .../docs/resources/migration_notes.md | 14 ++++ .../aztec-node/public_data_overrides.test.ts | 84 +++++++++++++++++++ .../src/aztec-node/public_data_overrides.ts | 35 ++++++++ .../aztec-node/src/aztec-node/server.ts | 11 ++- yarn-project/aztec.js/src/api/wallet.ts | 3 + .../src/contract/interaction_options.ts | 3 + yarn-project/aztec.js/src/wallet/wallet.ts | 2 + .../end-to-end/src/e2e_avm_simulator.test.ts | 49 +++++++++++ .../end-to-end/src/test-wallet/test_wallet.ts | 7 +- yarn-project/pxe/src/pxe.ts | 12 ++- .../stdlib/src/interfaces/aztec-node.ts | 15 +++- yarn-project/stdlib/src/interfaces/client.ts | 1 + .../src/interfaces/public_storage_override.ts | 26 ++++++ yarn-project/stdlib/src/tx/simulated_tx.ts | 18 ++-- .../wallet-sdk/src/base-wallet/base_wallet.ts | 5 +- .../wallet-sdk/src/base-wallet/utils.ts | 10 ++- 17 files changed, 293 insertions(+), 18 deletions(-) create mode 100644 yarn-project/aztec-node/src/aztec-node/public_data_overrides.test.ts create mode 100644 yarn-project/aztec-node/src/aztec-node/public_data_overrides.ts create mode 100644 yarn-project/stdlib/src/interfaces/public_storage_override.ts diff --git a/docs/docs-developers/docs/aztec-js/how_to_test.md b/docs/docs-developers/docs/aztec-js/how_to_test.md index d22443b724e2..1ae388802997 100644 --- a/docs/docs-developers/docs/aztec-js/how_to_test.md +++ b/docs/docs-developers/docs/aztec-js/how_to_test.md @@ -67,6 +67,22 @@ Test that invalid operations revert as expected: Use `.simulate()` to test reverts without spending gas. The simulation will throw if the transaction would fail onchain. +## Simulating with overrides + +`.simulate()` accepts an `overrides` option that injects values into the simulator's (ephemeral) world-state fork and contract DB before the call runs. The override is scoped to that single simulation and thrown away afterwards. + +Override a public-storage slot: + +```typescript +const result = await contract.methods.read_balance(account).simulate({ + overrides: { + publicStorage: [{ contract: contract.address, slot: BALANCE_SLOT, value: new Fr(1_000_000n) }], + }, +}); +``` + +Use this to set up state preconditions, reproduce production bugs against pinned storage, or exercise rare value branches without orchestrating the contract calls that produce them. + ## Further reading - [How to read contract data](./how_to_read_data.md) diff --git a/docs/docs-developers/docs/resources/migration_notes.md b/docs/docs-developers/docs/resources/migration_notes.md index d7e55d1360e7..23271b2be237 100644 --- a/docs/docs-developers/docs/resources/migration_notes.md +++ b/docs/docs-developers/docs/resources/migration_notes.md @@ -164,6 +164,20 @@ If you set `Noir: Nargo Path` in the VS Code Noir extension to `$HOME/.aztec/cur + overrides = { contracts: { [addr.toString()]: { instance: { ...instance, currentContractClassId: stubClassId } } } }; ``` +### [Aztec.js] `simulate` accepts `overrides` for testing "what if storage value was X?" + +`Contract.methods.foo(...).simulate(...)` now accepts an `overrides` option that injects values into the simulator's (ephemeral) world-state fork and contract DB before the call runs. The supported field is `publicStorage`, which writes a `(contract, slot, value)` into the public-data tree as if a previous tx had set it. Overrides are thrown away after simulation completes. + +```typescript +const result = await contract.methods.read_balance(account).simulate({ + overrides: { + publicStorage: [{ contract: contract.address, slot: BALANCE_SLOT, value: new Fr(1_000_000n) }], + }, +}); +``` + +The same option flows through `wallet.simulateTx` and eventually to `simulatePublicCalls` RPC on `AztecNode`. + ### [PXE] `proveTx` takes an options bag `PXE.proveTx` used to accept `scopes` as a positional argument; it now takes an options bag consistent with `simulateTx` and `profileTx`, and adds an optional `senderForTags` field. Update direct callers: diff --git a/yarn-project/aztec-node/src/aztec-node/public_data_overrides.test.ts b/yarn-project/aztec-node/src/aztec-node/public_data_overrides.test.ts new file mode 100644 index 000000000000..9944b2ccc033 --- /dev/null +++ b/yarn-project/aztec-node/src/aztec-node/public_data_overrides.test.ts @@ -0,0 +1,84 @@ +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 { type MockProxy, mock } from 'jest-mock-extended'; + +import { applyPublicDataOverrides } from './public_data_overrides.js'; + +describe('applyPublicDataOverrides', () => { + let fork: MockProxy; + + beforeEach(() => { + fork = mock(); + 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); + }); +}); diff --git a/yarn-project/aztec-node/src/aztec-node/public_data_overrides.ts b/yarn-project/aztec-node/src/aztec-node/public_data_overrides.ts new file mode 100644 index 000000000000..b72bc4df0695 --- /dev/null +++ b/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 { PublicStorageOverride } 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. + * + * Writes never reach committed world state — the fork is thrown away after simulation. + */ +export async function applyPublicDataOverrides( + fork: MerkleTreeWriteOperations, + publicStorage: PublicStorageOverride[] | undefined, +): Promise { + if (!publicStorage?.length) { + return; + } + + const writes = await Promise.all( + publicStorage.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()), + ); +} diff --git a/yarn-project/aztec-node/src/aztec-node/server.ts b/yarn-project/aztec-node/src/aztec-node/server.ts index 804569df9d55..a396db829f3d 100644 --- a/yarn-project/aztec-node/src/aztec-node/server.ts +++ b/yarn-project/aztec-node/src/aztec-node/server.ts @@ -106,6 +106,7 @@ import { type GlobalVariableBuilder as GlobalVariableBuilderInterface, type IndexedTxEffect, PublicSimulationOutput, + type SimulationOverrides, Tx, type TxHash, TxReceipt, @@ -146,6 +147,7 @@ import { } from './block_response_helpers.js'; import { type AztecNodeConfig, createKeyStoreForValidator } from './config.js'; import { NodeMetrics } from './node_metrics.js'; +import { applyPublicDataOverrides } from './public_data_overrides.js'; /** * The aztec node. @@ -1440,11 +1442,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 overrides - Optional pre-simulation overrides applied to the ephemeral fork and contract DB. **/ @trackSpan('AztecNodeService.simulatePublicCalls', (tx: Tx) => ({ [Attributes.TX_HASH]: tx.getTxHash().toString(), })) - public async simulatePublicCalls(tx: Tx, skipFeeEnforcement = false): Promise { + public async simulatePublicCalls( + tx: Tx, + skipFeeEnforcement = false, + overrides?: SimulationOverrides, + ): Promise { // Check total gas limit for simulation const gasSettings = tx.data.constants.txContext.gasSettings; const txGasLimit = gasSettings.gasLimits.l2Gas; @@ -1489,6 +1497,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb await this.worldStateSynchronizer.syncImmediate(latestBlockNumber); const merkleTreeFork = await this.worldStateSynchronizer.fork(); try { + await applyPublicDataOverrides(merkleTreeFork, overrides?.publicStorage); const config = PublicSimulatorConfig.from({ skipFeeEnforcement, collectDebugLogs: true, diff --git a/yarn-project/aztec.js/src/api/wallet.ts b/yarn-project/aztec.js/src/api/wallet.ts index db844b6f32ff..62bbb6f01696 100644 --- a/yarn-project/aztec.js/src/api/wallet.ts +++ b/yarn-project/aztec.js/src/api/wallet.ts @@ -78,4 +78,7 @@ export { AccountManager } from '../wallet/account_manager.js'; export { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js'; +export { type PublicStorageOverride, PublicStorageOverrideSchema } from '@aztec/stdlib/interfaces/client'; +export { SimulationOverrides } from '@aztec/stdlib/tx'; + export { type DeployAccountOptions, DeployAccountMethod } from '../wallet/deploy_account_method.js'; diff --git a/yarn-project/aztec.js/src/contract/interaction_options.ts b/yarn-project/aztec.js/src/contract/interaction_options.ts index 7eeee0ce8842..98583a42f52b 100644 --- a/yarn-project/aztec.js/src/contract/interaction_options.ts +++ b/yarn-project/aztec.js/src/contract/interaction_options.ts @@ -7,6 +7,7 @@ import { type Capsule, OFFCHAIN_MESSAGE_IDENTIFIER, type OffchainEffect, + type SimulationOverrides, type SimulationStats, type TxHash, type TxReceipt, @@ -157,6 +158,8 @@ export type SimulateInteractionOptions = Omit & { /** Whether to include metadata such as performance statistics (e.g. timing information of the different circuits and oracles) and gas estimation * in the simulation result, in addition to the return value and offchain effects */ includeMetadata?: boolean; + /** Pre-simulation overrides applied to the ephemeral fork and contract DB (publicStorage writes, contract instance overrides). */ + overrides?: SimulationOverrides; }; /** diff --git a/yarn-project/aztec.js/src/wallet/wallet.ts b/yarn-project/aztec.js/src/wallet/wallet.ts index 9fd080aaec98..a9f35e5815ab 100644 --- a/yarn-project/aztec.js/src/wallet/wallet.ts +++ b/yarn-project/aztec.js/src/wallet/wallet.ts @@ -19,6 +19,7 @@ import type { ExecutionPayload, InTx } from '@aztec/stdlib/tx'; import { Capsule, HashedValues, + SimulationOverrides, TxHash, TxProfileResult, TxReceipt, @@ -335,6 +336,7 @@ export const SimulateOptionsSchema = z.object({ skipFeeEnforcement: optional(z.boolean()), includeMetadata: optional(z.boolean()), additionalScopes: optional(z.array(schemas.AztecAddress)), + overrides: optional(SimulationOverrides.schema), }); export const ProfileOptionsSchema = SimulateOptionsSchema.extend({ diff --git a/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts b/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts index 05eef0fdf403..a8b507dd21d7 100644 --- a/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts +++ b/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts @@ -1,7 +1,9 @@ import { AztecAddress } from '@aztec/aztec.js/addresses'; import { BatchCall, type ContractInstanceWithAddress } from '@aztec/aztec.js/contracts'; import { Fr } from '@aztec/aztec.js/fields'; +import type { AztecNode } from '@aztec/aztec.js/node'; import { TxExecutionResult } from '@aztec/aztec.js/tx'; +import type { PublicStorageOverride } from '@aztec/aztec.js/wallet'; import type { Wallet } from '@aztec/aztec.js/wallet'; import { AvmInitializerTestContract } from '@aztec/noir-test-contracts.js/AvmInitializerTest'; import { AvmTestContract } from '@aztec/noir-test-contracts.js/AvmTest'; @@ -16,6 +18,7 @@ describe('e2e_avm_simulator', () => { jest.setTimeout(TIMEOUT); let wallet: Wallet; + let aztecNode: AztecNode; let defaultAccountAddress: AztecAddress; let teardown: () => Promise; @@ -23,6 +26,7 @@ describe('e2e_avm_simulator', () => { ({ teardown, wallet, + aztecNode, accounts: [defaultAccountAddress], } = await setup(1)); await ensureAccountContractsPublished(wallet, [defaultAccountAddress]); @@ -251,6 +255,51 @@ describe('e2e_avm_simulator', () => { }); }); + describe('publicDataOverrides', () => { + // AvmTestContract: `single` is the first storage variable and lives at raw slot 1. + const SINGLE_SLOT = new Fr(1n); + let avmContract: AvmTestContract; + + beforeEach(async () => { + ({ contract: avmContract } = await AvmTestContract.deploy(wallet).send({ from: defaultAccountAddress })); + }); + + it('simulated read of an unwritten slot returns the override; real storage is untouched', async () => { + const overrideValue = new Fr(0xdeadbeefn); + const publicStorage: PublicStorageOverride[] = [ + { contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue }, + ]; + + const simResult = await avmContract.methods + .read_storage_single() + .simulate({ from: defaultAccountAddress, overrides: { publicStorage } }); + expect(simResult.result).toEqual(overrideValue.toBigInt()); + + // Real state is untouched — the slot was never written. + const realValue = await aztecNode.getPublicStorageAt('latest', avmContract.address, SINGLE_SLOT); + expect(realValue.toBigInt()).toEqual(0n); + }); + + it('simulated read returns the override when a slot was previously written by a real tx', async () => { + const realValue = new Fr(100n); + await avmContract.methods.set_storage_single(realValue).send({ from: defaultAccountAddress }); + + const overrideValue = new Fr(999n); + const publicStorage: PublicStorageOverride[] = [ + { contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue }, + ]; + + const simResult = await avmContract.methods + .read_storage_single() + .simulate({ from: defaultAccountAddress, overrides: { publicStorage } }); + expect(simResult.result).toEqual(overrideValue.toBigInt()); + + // Real storage still holds the original written value. + const storedValue = await aztecNode.getPublicStorageAt('latest', avmContract.address, SINGLE_SLOT); + expect(storedValue.toBigInt()).toEqual(realValue.toBigInt()); + }); + }); + describe('AvmInitializerTestContract', () => { let avmContract: AvmInitializerTestContract; diff --git a/yarn-project/end-to-end/src/test-wallet/test_wallet.ts b/yarn-project/end-to-end/src/test-wallet/test_wallet.ts index 1363f6df4160..f559bdda136f 100644 --- a/yarn-project/end-to-end/src/test-wallet/test_wallet.ts +++ b/yarn-project/end-to-end/src/test-wallet/test_wallet.ts @@ -287,11 +287,14 @@ export class TestWallet extends BaseWallet { : executionPayload; const chainInfo = await this.getChainInfo(); - let overrides: SimulationOverrides | undefined; + let overrides = opts.overrides; let txRequest: TxExecutionRequest; if (useOverride) { const accountOverrides = await this.buildAccountOverrides(scopes); - overrides = new SimulationOverrides(accountOverrides); + overrides = new SimulationOverrides({ + publicStorage: overrides?.publicStorage, + contracts: { ...overrides?.contracts, ...accountOverrides }, + }); } if (from === NO_FROM) { diff --git a/yarn-project/pxe/src/pxe.ts b/yarn-project/pxe/src/pxe.ts index a030004dee51..5c678f8cacfe 100644 --- a/yarn-project/pxe/src/pxe.ts +++ b/yarn-project/pxe/src/pxe.ts @@ -120,7 +120,11 @@ export type SimulateTxOpts = { skipFeeEnforcement?: boolean; /** If true, kernel logic is emulated in TS for simulation */ skipKernels?: boolean; - /** State overrides for the simulation, such as contract instances and artifacts. Requires skipKernels: true */ + /** + * Pre-simulation overrides applied to the ephemeral fork and contract DB. Bundles publicStorage + * writes (no skipKernels required) and per-address (instance, artifact?) overrides used by both + * AVM-side public dispatch and PXE-side ACIR private dispatch (requires skipKernels: true). + */ overrides?: SimulationOverrides; /** Addresses whose private state and keys are accessible during private execution */ scopes: AztecAddress[]; @@ -484,11 +488,11 @@ export class PXE { * It can also be used for estimating gas in the future. * @param tx - The transaction to be simulated. */ - async #simulatePublicCalls(tx: Tx, skipFeeEnforcement: boolean) { + async #simulatePublicCalls(tx: Tx, skipFeeEnforcement: boolean, overrides?: SimulationOverrides) { // Simulating public calls can throw if the TX fails in a phase that doesn't allow reverts (setup) // Or return as reverted if it fails in a phase that allows reverts (app logic, teardown) try { - const result = await this.node.simulatePublicCalls(tx, skipFeeEnforcement); + const result = await this.node.simulatePublicCalls(tx, skipFeeEnforcement, overrides); if (result.revertReason) { throw result.revertReason; } @@ -1054,7 +1058,7 @@ export class PXE { let publicOutput: PublicSimulationOutput | undefined; if (simulatePublic && publicInputs.forPublic) { const publicSimulationTimer = new Timer(); - publicOutput = await this.#simulatePublicCalls(simulatedTx, skipFeeEnforcement); + publicOutput = await this.#simulatePublicCalls(simulatedTx, skipFeeEnforcement, overrides); publicSimulationTime = publicSimulationTimer.ms(); if (publicOutput?.debugLogs?.length) { await displayDebugLogs(publicOutput.debugLogs, addr => this.contractStore.getDebugContractName(addr)); diff --git a/yarn-project/stdlib/src/interfaces/aztec-node.ts b/yarn-project/stdlib/src/interfaces/aztec-node.ts index 4e71f392d04a..0488bb98c30f 100644 --- a/yarn-project/stdlib/src/interfaces/aztec-node.ts +++ b/yarn-project/stdlib/src/interfaces/aztec-node.ts @@ -46,6 +46,7 @@ import { BlockHeader, type IndexedTxEffect, PublicSimulationOutput, + SimulationOverrides, Tx, TxHash, TxReceipt, @@ -464,8 +465,15 @@ 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 overrides - Optional pre-simulation overrides applied to the ephemeral fork and contract DB + * (publicStorage writes, contract instance overrides). **/ - simulatePublicCalls(tx: Tx, skipFeeEnforcement?: boolean): Promise; + simulatePublicCalls( + tx: Tx, + skipFeeEnforcement?: boolean, + overrides?: SimulationOverrides, + ): Promise; /** * Returns true if the transaction is valid for inclusion at the current state. Valid transactions can be @@ -668,7 +676,10 @@ export const AztecNodeApiSchema: ApiSchemaFor = { .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(SimulationOverrides.schema)) + .returns(PublicSimulationOutput.schema), isValidTx: z .function() diff --git a/yarn-project/stdlib/src/interfaces/client.ts b/yarn-project/stdlib/src/interfaces/client.ts index 2fafa5bcd151..88fa0f93b402 100644 --- a/yarn-project/stdlib/src/interfaces/client.ts +++ b/yarn-project/stdlib/src/interfaces/client.ts @@ -9,3 +9,4 @@ export * from './l1_publish_info.js'; export * from './private_kernel_prover.js'; export * from './get_logs_response.js'; export * from './api_limit.js'; +export * from './public_storage_override.js'; diff --git a/yarn-project/stdlib/src/interfaces/public_storage_override.ts b/yarn-project/stdlib/src/interfaces/public_storage_override.ts new file mode 100644 index 000000000000..355b6f67871b --- /dev/null +++ b/yarn-project/stdlib/src/interfaces/public_storage_override.ts @@ -0,0 +1,26 @@ +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-storage override to inject before simulation. Identifies a slot in + * (contract, raw slot) space, not by leafSlot — the simulator computes the leafSlot internally + * via `computePublicDataTreeLeafSlot`. + */ +export type PublicStorageOverride = { + /** 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 PublicStorageOverrideSchema = z.object({ + contract: schemas.AztecAddress, + slot: schemas.Fr, + value: schemas.Fr, +}); diff --git a/yarn-project/stdlib/src/tx/simulated_tx.ts b/yarn-project/stdlib/src/tx/simulated_tx.ts index 7983036ade3a..0ffd03c5674a 100644 --- a/yarn-project/stdlib/src/tx/simulated_tx.ts +++ b/yarn-project/stdlib/src/tx/simulated_tx.ts @@ -9,6 +9,7 @@ import { } from '../contract/interfaces/contract_instance.js'; import { Gas } from '../gas/gas.js'; import type { GasUsed } from '../gas/gas_used.js'; +import { type PublicStorageOverride, PublicStorageOverrideSchema } from '../interfaces/public_storage_override.js'; import { PrivateKernelTailCircuitPublicInputs } from '../kernel/private_kernel_tail_circuit_public_inputs.js'; import { ChonkProof } from '../proofs/chonk_proof.js'; import type { OffchainEffect } from './offchain_effect.js'; @@ -32,24 +33,27 @@ import { Tx } from './tx.js'; export type ContractOverrides = Record; /* - * Optional values that can be overridden during simulation. In order to simulate a transaction with these - * set, it *must* be run without the kernel circuits, or validations will fail. + * Optional values that can be overridden during simulation. `publicStorage` writes to the public-data + * tree fork. `contracts` overrides contract instances in the contract DB. + * In order to simulate a transaction with these set, it *must* be run without the kernel circuits, + * or validations will fail. */ export class SimulationOverrides { + public publicStorage?: PublicStorageOverride[]; public contracts?: ContractOverrides; - constructor(contracts: ContractOverrides = {}) { - this.contracts = Object.keys(contracts).length > 0 ? contracts : undefined; + constructor(args: { publicStorage?: PublicStorageOverride[]; contracts?: ContractOverrides } = {}) { + this.publicStorage = args.publicStorage?.length ? args.publicStorage : undefined; + this.contracts = args.contracts && Object.keys(args.contracts).length > 0 ? args.contracts : undefined; } static get schema() { return z .object({ + publicStorage: optional(z.array(PublicStorageOverrideSchema)), contracts: optional(z.record(z.string(), z.object({ instance: ContractInstanceWithAddressSchema }))), }) - .transform(({ contracts }) => { - return new SimulationOverrides(contracts); - }); + .transform(args => new SimulationOverrides(args)); } } diff --git a/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts b/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts index 1e4c329fbc65..5d1fc6827c05 100644 --- a/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts +++ b/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts @@ -85,7 +85,7 @@ export type FeeOptions = { /** Options for `simulateViaEntrypoint`. */ export type SimulateViaEntrypointOptions = Pick< SimulateOptions, - 'from' | 'additionalScopes' | 'skipTxValidation' | 'skipFeeEnforcement' | 'sendMessagesAs' + 'from' | 'additionalScopes' | 'skipTxValidation' | 'skipFeeEnforcement' | 'sendMessagesAs' | 'overrides' > & { /** Fee options for the entrypoint */ feeOptions: FeeOptions; @@ -365,6 +365,7 @@ export abstract class BaseWallet implements Wallet { skipFeeEnforcement: opts.skipFeeEnforcement, scopes: this.scopesFrom(opts.from, opts.additionalScopes), senderForTags: this.senderForTagsFrom(opts.from, opts.sendMessagesAs), + overrides: opts.overrides, }); const appCallOffset = await this.computeAppCallOffset(opts.from, opts.feeOptions); return TxSimulationResultWithAppOffset.fromResultAndOffset(result, appCallOffset); @@ -428,6 +429,7 @@ export abstract class BaseWallet implements Wallet { blockHeader, opts.skipFeeEnforcement ?? true, this.getContractName.bind(this), + opts.overrides, ) : Promise.resolve([]), remainingCalls.length > 0 @@ -438,6 +440,7 @@ export abstract class BaseWallet implements Wallet { skipTxValidation: opts.skipTxValidation, skipFeeEnforcement: opts.skipFeeEnforcement ?? true, sendMessagesAs: opts.sendMessagesAs, + overrides: opts.overrides, }) : Promise.resolve(null), ]); diff --git a/yarn-project/wallet-sdk/src/base-wallet/utils.ts b/yarn-project/wallet-sdk/src/base-wallet/utils.ts index 6bc1423f8aba..0a22864c4b61 100644 --- a/yarn-project/wallet-sdk/src/base-wallet/utils.ts +++ b/yarn-project/wallet-sdk/src/base-wallet/utils.ts @@ -25,6 +25,7 @@ import { PrivateCallExecutionResult, PrivateExecutionResult, PublicSimulationOutput, + type SimulationOverrides, Tx, TxContext, TxSimulationResult, @@ -65,6 +66,8 @@ export function extractOptimizablePublicStaticCalls(payload: ExecutionPayload): * @param gasSettings - Gas settings for the transaction. * @param blockHeader - Block header to use as anchor. * @param skipFeeEnforcement - Whether to skip fee enforcement during simulation. + * @param getContractName - Resolver for contract names (used for debug log display). + * @param overrides - Optional pre-simulation overrides applied to the ephemeral fork and contract DB. * @returns TxSimulationResult with public return values. */ async function simulateBatchViaNode( @@ -76,6 +79,7 @@ async function simulateBatchViaNode( blockHeader: BlockHeader, skipFeeEnforcement: boolean, getContractName: ContractNameResolver, + overrides?: SimulationOverrides, ): Promise { const txContext = new TxContext(chainInfo.chainId, chainInfo.version, gasSettings); @@ -143,7 +147,7 @@ async function simulateBatchViaNode( publicFunctionCalldata: publicFunctionCalldata, }); - const publicOutput = await node.simulatePublicCalls(tx, skipFeeEnforcement); + const publicOutput = await node.simulatePublicCalls(tx, skipFeeEnforcement, overrides); if (publicOutput.revertReason) { throw publicOutput.revertReason; @@ -166,6 +170,8 @@ async function simulateBatchViaNode( * @param gasSettings - Gas settings for the transaction. * @param blockHeader - Block header to use as anchor. * @param skipFeeEnforcement - Whether to skip fee enforcement during simulation. + * @param getContractName - Resolver for contract names (used for debug log display). + * @param overrides - Optional pre-simulation overrides applied to the ephemeral fork and contract DB. * @returns Array of TxSimulationResult, one per batch. */ export async function simulateViaNode( @@ -177,6 +183,7 @@ export async function simulateViaNode( blockHeader: BlockHeader, skipFeeEnforcement: boolean = true, getContractName: ContractNameResolver, + overrides?: SimulationOverrides, ): Promise { const batches: FunctionCall[][] = []; @@ -196,6 +203,7 @@ export async function simulateViaNode( blockHeader, skipFeeEnforcement, getContractName, + overrides, ); results.push(result); } From 1b062ca2f5b4d849fe345a11bf03c7fe497166ae Mon Sep 17 00:00:00 2001 From: dbanks12 Date: Wed, 6 May 2026 20:28:39 +0000 Subject: [PATCH 2/2] fix(wallets): pass contracts via options bag to SimulationOverrides EmbeddedWallet was still calling the old positional-argument form of the SimulationOverrides constructor, so contract-instance overrides for the stub account were silently dropped, and simulation re-entered the real account entrypoint with no auth witness. --- docs/docs-developers/docs/resources/migration_notes.md | 7 +++++++ yarn-project/wallets/src/embedded/embedded_wallet.ts | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/docs-developers/docs/resources/migration_notes.md b/docs/docs-developers/docs/resources/migration_notes.md index 23271b2be237..f1333f13d158 100644 --- a/docs/docs-developers/docs/resources/migration_notes.md +++ b/docs/docs-developers/docs/resources/migration_notes.md @@ -178,6 +178,13 @@ const result = await contract.methods.read_balance(account).simulate({ The same option flows through `wallet.simulateTx` and eventually to `simulatePublicCalls` RPC on `AztecNode`. +Direct callers of the `SimulationOverrides` constructor must switch from a positional `contracts` argument to an options bag: + +```diff +- new SimulationOverrides(contracts); ++ new SimulationOverrides({ contracts }); +``` + ### [PXE] `proveTx` takes an options bag `PXE.proveTx` used to accept `scopes` as a positional argument; it now takes an options bag consistent with `simulateTx` and `profileTx`, and adds an optional `senderForTags` field. Update direct callers: diff --git a/yarn-project/wallets/src/embedded/embedded_wallet.ts b/yarn-project/wallets/src/embedded/embedded_wallet.ts index fd63f0613a18..bfddd9907030 100644 --- a/yarn-project/wallets/src/embedded/embedded_wallet.ts +++ b/yarn-project/wallets/src/embedded/embedded_wallet.ts @@ -271,7 +271,7 @@ export class EmbeddedWallet extends BaseWallet { const chainInfo = await this.getChainInfo(); const accountOverrides = await this.buildAccountOverrides(scopes); - const overrides = new SimulationOverrides(accountOverrides); + const overrides = new SimulationOverrides({ contracts: accountOverrides }); let txRequest: TxExecutionRequest; if (from === NO_FROM) {