feat: merge-train/fairies (#23074)

BEGIN_COMMIT_OVERRIDE
refactor: drop artifact field from SimulationOverrides (#22957)
feat: public data tree overrides for simulation (#22830)
feat: expose wallet.registerContractClass (prep for simulating contract
upgrades) (#22932)
END_COMMIT_OVERRIDE

Author: AztecBot

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)

docs/docs-developers/docs/resources/migration_notes.md

@@ -151,6 +151,40 @@ If you relied on a bundled bare-name binary for general use:
 
 If you set `Noir: Nargo Path` in the VS Code Noir extension to `$HOME/.aztec/current/bin/nargo`, change it to `$HOME/.aztec/current/bin/aztec-nargo` (the symlink is a drop-in for `nargo`). See the [Noir VSCode Extension guide](../aztec-nr/installation.md) for details.
 
+### [Stdlib] `SimulationOverrides.contracts` entries no longer carry an artifact
+
+`ContractOverrides` entries are now `{ instance }` only. To override a contract's artifact, pre-register the target class via `pxe.registerContractClass(artifact)` and set the override instance's `currentContractClassId` to that class id:
+
+```diff
+- const instance = await getContractInstanceFromInstantiationParams(stubArtifact, { salt: Fr.random() });
++ const instance = await pxe.getContractInstance(addr);
++ await pxe.registerContractClass(stubArtifact);
++ const stubClassId = (await getContractClassFromArtifact(stubArtifact)).id;
+- overrides = { contracts: { [addr.toString()]: { instance, artifact: stubArtifact } } };
++ 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`.
+
+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:

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<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 { 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<void> {
+  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()),
+  );
+}

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<PublicSimulationOutput> {
+  public async simulatePublicCalls(
+    tx: Tx,
+    skipFeeEnforcement = false,
+    overrides?: SimulationOverrides,
+  ): Promise<PublicSimulationOutput> {
     // 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,

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';

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<SendInteractionOptions, 'fee'> & {
   /** 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;
 };
 
 /**

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

@@ -138,23 +138,25 @@ export interface ContractsCapability {
 export interface GrantedContractsCapability extends ContractsCapability {}
 
 /**
- * Contract class capability - for querying contract class metadata.
+ * Contract class capability - for querying contract class meatadata and registering contract classes.
  *
  * Maps to wallet methods:
- * - getContractClassMetadata
+ * - getContractClassMetadata (when canGetMetadata: true)
+ * - registerContractClass (when canRegister: true)
  *
  * Contract classes are identified by their class ID (Fr), not by contract address.
  * Multiple contract instances can share the same class. This capability grants
- * permission to query metadata for specific contract classes.
+ * permission to query metadata for, and register, specific contract classes.
  *
  * Apps typically acquire this permission automatically when registering a contract
  * with an artifact (the wallet auto-grants permission for that contract's class ID).
  *
  * @example
- * // Query specific contract classes
+ * // Register and query a specific contract class
  * \{
  *   type: 'contractClasses',
- *   classes: [classId1, classId2],
+ *   classes: [classId1],
+ *   canRegister: true,
  *   canGetMetadata: true
  * \}
  *
@@ -177,6 +179,9 @@ export interface ContractClassesCapability {
    */
   classes: '*' | Fr[];
 
+  /** Can register a contract class artifact in the local PXE. Maps to: registerContractClass */
+  canRegister?: boolean;
+
   /** Can query contract class metadata. Maps to: getContractClassMetadata */
   canGetMetadata: boolean;
 }

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

@@ -157,6 +157,19 @@ describe('WalletSchema', () => {
     });
   });
 
+  it('registerContractClass', async () => {
+    const mockArtifact: ContractArtifact = {
+      name: 'TestContract',
+      aztecVersion: DEV_VERSION,
+      functions: [],
+      nonDispatchPublicFunctions: [],
+      outputs: { structs: {}, globals: {} },
+      fileMap: {},
+      storageLayout: {},
+    };
+    await context.client.registerContractClass(mockArtifact);
+  });
+
   it('simulateTx', async () => {
     const exec: ExecutionPayload = {
       calls: [],
@@ -451,6 +464,8 @@ class MockWallet implements Wallet {
     };
   }
 
+  async registerContractClass(_artifact: any): Promise<void> {}
+
   async simulateTx(_exec: ExecutionPayload, _opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset> {
     return TxSimulationResultWithAppOffset.fromResultAndOffset(await TxSimulationResult.random(), 0);
   }

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,
@@ -271,6 +272,12 @@ export type Wallet = {
     artifact?: ContractArtifact,
     secretKey?: Fr,
   ): Promise<ContractInstanceWithAddress>;
+  /**
+   * Registers a contract class artifact in the local PXE without binding it to any instance.
+   * Useful for simulation flows that need the artifact available locally before any on-chain
+   * upgrade has taken effect. No chain check.
+   */
+  registerContractClass(artifact: ContractArtifact): Promise<void>;
   simulateTx(exec: ExecutionPayload, opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset>;
   executeUtility(call: FunctionCall, opts: ExecuteUtilityOptions): Promise<UtilityExecutionResult>;
   profileTx(exec: ExecutionPayload, opts: ProfileOptions): Promise<TxProfileResult>;
@@ -335,6 +342,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({
@@ -427,6 +435,7 @@ export const GrantedContractsCapabilitySchema = ContractsCapabilitySchema;
 export const ContractClassesCapabilitySchema = z.object({
   type: z.literal('contractClasses'),
   classes: z.union([z.literal('*'), z.array(schemas.Fr)]),
+  canRegister: optional(z.boolean()),
   canGetMetadata: z.boolean(),
 });
 
@@ -557,6 +566,7 @@ const WalletMethodSchemas = {
     .function()
     .args(ContractInstanceWithAddressSchema, optional(ContractArtifactSchema), optional(schemas.Fr))
     .returns(ContractInstanceWithAddressSchema),
+  registerContractClass: z.function().args(ContractArtifactSchema).returns(z.void()),
   simulateTx: z
     .function()
     .args(ExecutionPayloadSchema, SimulateOptionsSchema)