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 ff4bb092fdf5..6a2d05315222 100644 --- a/docs/docs-developers/docs/aztec-js/how_to_test.md +++ b/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 diff --git a/yarn-project/aztec-node/src/aztec-node/server.ts b/yarn-project/aztec-node/src/aztec-node/server.ts index 1951944e7cb0..e94a0c2ae7cd 100644 --- a/yarn-project/aztec-node/src/aztec-node/server.ts +++ b/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 diff --git a/yarn-project/simulator/src/public/public_db_sources.ts b/yarn-project/simulator/src/public/public_db_sources.ts index 6852bf23473d..26f8d8ae9997 100644 --- a/yarn-project/simulator/src/public/public_db_sources.ts +++ b/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); } diff --git a/yarn-project/stdlib/src/interfaces/state_overrides.ts b/yarn-project/stdlib/src/interfaces/state_overrides.ts index 508bece18106..7730ed3953cb 100644 --- a/yarn-project/stdlib/src/interfaces/state_overrides.ts +++ b/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(), });