Skip to content

Commit 7c70c29

Browse files
dbanks12dbanks12
committed
feat: public data tree overrides for simulation
1 parent 6e05d8b commit 7c70c29

13 files changed

Lines changed: 241 additions & 9 deletions

File tree

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

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+
});

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

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 { PublicDataTreeOverride } 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+
publicDataOverrides: PublicDataTreeOverride[] | undefined,
19+
): Promise<void> {
20+
if (!publicDataOverrides?.length) {
21+
return;
22+
}
23+
24+
const writes = await Promise.all(
25+
publicDataOverrides.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
@@ -75,6 +75,7 @@ import {
7575
type AztecNodeDebug,
7676
type GetContractClassLogsResponse,
7777
type GetPublicLogsResponse,
78+
type PublicDataTreeOverride,
7879
} from '@aztec/stdlib/interfaces/client';
7980
import {
8081
type AllowedElement,
@@ -131,6 +132,7 @@ import { createSentinel } from '../sentinel/factory.js';
131132
import { Sentinel } from '../sentinel/sentinel.js';
132133
import { type AztecNodeConfig, createKeyStoreForValidator } from './config.js';
133134
import { NodeMetrics } from './node_metrics.js';
135+
import { applyPublicDataOverrides } from './public_data_overrides.js';
134136

135137
/**
136138
* The aztec node.
@@ -1312,11 +1314,17 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
13121314
/**
13131315
* Simulates the public part of a transaction with the current state.
13141316
* @param tx - The transaction to simulate.
1317+
* @param skipFeeEnforcement - If true, fee enforcement is skipped.
1318+
* @param publicDataOverrides - Optional public state overrides injected into the ephemeral fork before simulation.
13151319
**/
13161320
@trackSpan('AztecNodeService.simulatePublicCalls', (tx: Tx) => ({
13171321
[Attributes.TX_HASH]: tx.getTxHash().toString(),
13181322
}))
1319-
public async simulatePublicCalls(tx: Tx, skipFeeEnforcement = false): Promise<PublicSimulationOutput> {
1323+
public async simulatePublicCalls(
1324+
tx: Tx,
1325+
skipFeeEnforcement = false,
1326+
publicDataOverrides?: PublicDataTreeOverride[],
1327+
): Promise<PublicSimulationOutput> {
13201328
// Check total gas limit for simulation
13211329
const gasSettings = tx.data.constants.txContext.gasSettings;
13221330
const txGasLimit = gasSettings.gasLimits.l2Gas;
@@ -1361,6 +1369,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
13611369
await this.worldStateSynchronizer.syncImmediate(latestBlockNumber);
13621370
const merkleTreeFork = await this.worldStateSynchronizer.fork();
13631371
try {
1372+
await applyPublicDataOverrides(merkleTreeFork, publicDataOverrides);
13641373
const config = PublicSimulatorConfig.from({
13651374
skipFeeEnforcement,
13661375
collectDebugLogs: true,

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

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

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

81+
export { type PublicDataTreeOverride, PublicDataTreeOverrideSchema } from '@aztec/stdlib/interfaces/client';
82+
8183
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 { PublicDataTreeOverride } 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+
/** Public state overrides injected into the ephemeral world-state fork before simulation. */
162+
publicDataOverrides?: PublicDataTreeOverride[];
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 { type PublicDataTreeOverride, PublicDataTreeOverrideSchema } 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+
publicDataOverrides: optional(z.array(PublicDataTreeOverrideSchema)),
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 { PublicDataTreeOverride } 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 publicDataOverrides: PublicDataTreeOverride[] = [
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, publicDataOverrides });
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 publicDataOverrides: PublicDataTreeOverride[] = [
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, publicDataOverrides });
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/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, PublicDataTreeOverride } 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+
/** Public state overrides injected into the ephemeral world-state fork before simulation. */
129+
publicDataOverrides?: PublicDataTreeOverride[];
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, publicDataOverrides?: PublicDataTreeOverride[]) {
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, publicDataOverrides);
475477
if (result.revertReason) {
476478
throw result.revertReason;
477479
}
@@ -956,6 +958,7 @@ export class PXE {
956958
overrides,
957959
scopes,
958960
senderForTags,
961+
publicDataOverrides,
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, publicDataOverrides);
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: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ import {
6767
type GetPublicLogsResponse,
6868
GetPublicLogsResponseSchema,
6969
} from './get_logs_response.js';
70+
import { type PublicDataTreeOverride, PublicDataTreeOverrideSchema } from './public_data_tree_override.js';
7071
import { type WorldStateSyncStatus, WorldStateSyncStatusSchema } from './world_state.js';
7172

7273
/**
@@ -462,8 +463,14 @@ export interface AztecNode
462463
* Simulates the public part of a transaction with the current state.
463464
* This currently just checks that the transaction execution succeeds.
464465
* @param tx - The transaction to simulate.
466+
* @param skipFeeEnforcement - If true, fee enforcement is skipped.
467+
* @param publicDataOverrides - Optional public state overrides injected into the ephemeral world-state fork before simulation.
465468
**/
466-
simulatePublicCalls(tx: Tx, skipFeeEnforcement?: boolean): Promise<PublicSimulationOutput>;
469+
simulatePublicCalls(
470+
tx: Tx,
471+
skipFeeEnforcement?: boolean,
472+
publicDataOverrides?: PublicDataTreeOverride[],
473+
): Promise<PublicSimulationOutput>;
467474

468475
/**
469476
* 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> = {
659666
.args(schemas.EthAddress, optional(schemas.SlotNumber), optional(schemas.SlotNumber))
660667
.returns(SingleValidatorStatsSchema.optional()),
661668

662-
simulatePublicCalls: z.function().args(Tx.schema, optional(z.boolean())).returns(PublicSimulationOutput.schema),
669+
simulatePublicCalls: z
670+
.function()
671+
.args(Tx.schema, optional(z.boolean()), optional(z.array(PublicDataTreeOverrideSchema)))
672+
.returns(PublicSimulationOutput.schema),
663673

664674
isValidTx: z
665675
.function()

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,4 @@ export * from './aztec-node-debug.js';
44
export * from './private_kernel_prover.js';
55
export * from './get_logs_response.js';
66
export * from './api_limit.js';
7+
export * from './public_data_tree_override.js';

0 commit comments

Comments
 (0)