feat: contract class/instance overrides for simulation

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

@@ -88,7 +88,7 @@ Use these 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
-- Simulate a contract instance as if it had been upgraded
+- Simulate a contract instance as if it had been upgraded (see `fastForwardContractUpdate` below)
 - Test branches that depend on rare values without orchestrating the contract calls that produce them
 
 ### Fast-forwarding a contract update

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

@@ -1552,7 +1552,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
         }),
       });
       const contractsDB = new PublicContractsDB(this.contractDataSource, this.log.getBindings());
-      contractsDB.addContracts(contractOverrides);
+      contractsDB.addContracts(contractOverrides, stateOverrides?.contractClasses);
       const processor = publicProcessorFactory.create(merkleTreeFork, newGlobalVariables, config, contractsDB);
 
       // REFACTOR: Consider merging ProcessReturnValues into ProcessedTx

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(contractInstances?: ContractInstanceWithAddress[]): void {
+  /** Inserts typed contract instances and classes directly into the current checkpoint. */
+  public addContracts(
+    contractInstances?: ContractInstanceWithAddress[],
+    contractClasses?: ContractClassPublic[],
+  ): void {
     const currentState = this.getCurrentState();
+    for (const contractClass of contractClasses ?? []) {
+      currentState.addClass(contractClass.id, contractClass);
+    }
     for (const instance of contractInstances ?? []) {
       currentState.addInstance(instance.address, instance);
     }

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

@@ -1,19 +1,24 @@
 import { z } from 'zod';
 
+import { type ContractClassPublic, ContractClassPublicSchema } from '../contract/interfaces/contract_class.js';
 import { type PublicStorageOverride, PublicStorageOverrideSchema } from './public_storage_override.js';
 
 /**
  * Pre-simulation state-tree overrides applied to the ephemeral world-state fork before running the tx.
  *
  * - `publicStorage`: write specific (contract, slot, value) entries in the public-data tree.
+ * - `contractClasses`: override contract classes in the contract DB.
  *
  * Contract instance overrides live separately as the top-level `contractOverrides` simulation option.
  */
 export type StateOverrides = {
   /** Public-storage writes to apply before simulation. */
   publicStorage?: PublicStorageOverride[];
+  /** Contract classes to override in the contract DB for the duration of the simulation. */
+  contractClasses?: ContractClassPublic[];
 };
 
 export const StateOverridesSchema = z.object({
   publicStorage: z.array(PublicStorageOverrideSchema).optional(),
+  contractClasses: z.array(ContractClassPublicSchema).optional(),
 });