feat: contract class overrides for simulation

Adds StateOverrides.contractClasses for shadowing contract classes in the
simulator's contract DB - the building block for simulating against mock
implementations or unregistered classes.

Author: dbanks12

docs/docs-developers/docs/aztec-js/how_to_test.md

@@ -71,6 +71,12 @@ Use `.simulate()` to test reverts without spending gas. The simulation will thro
 
 `.simulate()` accepts a `stateOverrides` option that injects values into the simulator's ephemeral world-state fork before the call runs. The override is scoped to that single simulation; the real chain state is untouched.
 
+Three override flavors are supported:
+
+- `publicStorage`: write `(contract, slot, value)` into the public-data tree.
+- `contractClasses`: shadow contract classes in the simulator's contract DB (as if they had been published on the class registry).
+- `contractInstances`: shadow contract instances at specific addresses (as if they had been deployed, possibly pointing at one of the shadowed classes).
+
 Override a public-storage slot:
 
 ```typescript
@@ -81,10 +87,23 @@ const result = await contract.methods.read_balance(account).simulate({
 });
 ```
 
+Override a contract's class — useful for testing against a mock implementation:
+
+```typescript
+const result = await contract.methods.foo().simulate({
+  stateOverrides: {
+    contractClasses: [mockClass],
+    contractInstances: [{ ...existingInstance, currentContractClassId: mockClass.id }],
+  },
+});
+```
+
 Use this to:
 
 - Set up state preconditions without running a full setup transaction
 - Reproduce a bug from production by pinning storage to the values seen at a specific block
+- Swap a contract's bytecode for a mock implementation in tests
+- Simulate calls against contracts that aren't yet deployed
 - Test branches that depend on rare values without orchestrating the contract calls that produce them
 
 ### Fast-forwarding a contract update

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

@@ -23,7 +23,7 @@ const result = await contract.methods.read_balance(account).simulate({
 
 The same option flows through `wallet.simulateTx` and eventually to `simulatePublicCalls` RPC on `AztecNode`.
 
-A second override flavor, `contractInstances`, shadows contract instances in the simulator's contract DB - useful for simulating a contract being on a different class than the one it was deployed with. To simulate a complete on-chain upgrade flow, use the `fastForwardContractUpdate` helper which produces a coherent set of `publicStorage` and `contractInstances` overrides:
+Additional override flavors `contractInstances` and `contractClasses` shadow contract instances and classes in the simulator's contract DB - useful for simulating mock implementations or contracts being on a different class than the one they were deployed with. To simulate a complete on-chain upgrade flow, use the `fastForwardContractUpdate` helper which produces a coherent set of `publicStorage` and `contractInstances` overrides:
 
 ```typescript
 import { fastForwardContractUpdate } from '@aztec/aztec.js';

yarn-project/simulator/src/public/public_db_sources.ts

@@ -76,9 +76,15 @@ export class PublicContractsDB implements PublicContractsDBInterface {
     this.addContractsFromLogs(contractDeploymentData.getRevertibleContractDeploymentData());
   }
 
-  /** Inserts typed contract instances directly into the current checkpoint. */
-  public addContracts(contracts: { contractInstances?: ContractInstanceWithAddress[] }): void {
+  /** Inserts typed contract classes and instances directly into the current checkpoint. */
+  public addContracts(contracts: {
+    contractClasses?: ContractClassPublic[];
+    contractInstances?: ContractInstanceWithAddress[];
+  }): void {
     const currentState = this.getCurrentState();
+    for (const contractClass of contracts.contractClasses ?? []) {
+      currentState.addClass(contractClass.id, contractClass);
+    }
     for (const instance of contracts.contractInstances ?? []) {
       currentState.addInstance(instance.address, instance);
     }

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

@@ -1,5 +1,6 @@
 import { z } from 'zod';
 
+import { type ContractClassPublic, ContractClassPublicSchema } from '../contract/interfaces/contract_class.js';
 import {
   type ContractInstanceWithAddress,
   ContractInstanceWithAddressSchema,
@@ -11,16 +12,20 @@ import { type PublicStorageOverride, PublicStorageOverrideSchema } from './publi
  * provided overrides to the ephemeral world-state fork (and contract DB) before running the tx.
  *
  * - `publicStorage`: write specific (contract, slot, value) entries in the public-data tree
+ * - `contractClasses`: shadow contract classes in the contract DB
  * - `contractInstances`: shadow contract instances in the contract DB
  */
 export type StateOverrides = {
   /** Public-storage writes to apply before simulation. */
   publicStorage?: PublicStorageOverride[];
+  /** Contract classes to shadow in the contract DB for the duration of the simulation. */
+  contractClasses?: ContractClassPublic[];
   /** Contract instances to shadow in the contract DB for the duration of the simulation. */
   contractInstances?: ContractInstanceWithAddress[];
 };
 
 export const StateOverridesSchema = z.object({
   publicStorage: z.array(PublicStorageOverrideSchema).optional(),
+  contractClasses: z.array(ContractClassPublicSchema).optional(),
   contractInstances: z.array(ContractInstanceWithAddressSchema).optional(),
 });