@@ -11,31 +11,34 @@ use alloy_sol_types::SolValue;
1111use foundry_cheatcodes:: {
1212 Broadcast , BroadcastableTransactions , CheatcodeInspectorStrategy ,
1313 CheatcodeInspectorStrategyContext , CheatcodeInspectorStrategyRunner , CheatsConfig , CheatsCtxt ,
14- CommonCreateInput , DealRecord , Ecx , EvmCheatcodeInspectorStrategyRunner , Result ,
15- Vm :: { dealCall, getNonce_0Call, pvmCall, rollCall, setNonceCall, setNonceUnsafeCall, warpCall} ,
14+ CommonCreateInput , DealRecord , Ecx , Error , EvmCheatcodeInspectorStrategyRunner , Result ,
15+ Vm :: {
16+ dealCall, etchCall, getNonce_0Call, pvmCall, rollCall, setNonceCall, setNonceUnsafeCall,
17+ warpCall,
18+ } ,
1619} ;
1720
1821use foundry_common:: sh_err;
1922use foundry_compilers:: resolc:: dual_compiled_contracts:: DualCompiledContracts ;
2023use revive_env:: { AccountId , Runtime , System , Timestamp } ;
2124
25+ use crate :: {
26+ cheatcodes:: mock_handler:: MockHandlerImpl , execute_with_externalities, trace,
27+ tracing:: apply_prestate_trace,
28+ } ;
29+ use alloy_eips:: eip7702:: SignedAuthorization ;
2230use polkadot_sdk:: {
2331 frame_support:: traits:: { Currency , fungible:: Mutate } ,
2432 pallet_balances,
2533 pallet_revive:: {
26- self , AddressMapper , BalanceOf , BalanceWithDust , Code , Config , ExecConfig , Pallet ,
34+ self , AccountInfo , AddressMapper , BalanceOf , BalanceWithDust , Code , Config , ContractInfo ,
35+ ExecConfig , Executable , Pallet ,
2736 } ,
2837 polkadot_runtime_common:: Bounded ,
2938 polkadot_sdk_frame:: prelude:: OriginFor ,
3039 sp_core:: { self , H160 } ,
3140 sp_weights:: Weight ,
3241} ;
33-
34- use crate :: {
35- cheatcodes:: mock_handler:: MockHandlerImpl , execute_with_externalities, trace,
36- tracing:: apply_prestate_trace,
37- } ;
38- use alloy_eips:: eip7702:: SignedAuthorization ;
3942use revm:: {
4043 bytecode:: opcode as op,
4144 context:: { CreateScheme , JournalTr } ,
@@ -157,6 +160,44 @@ fn set_block_number(new_height: U256, ecx: Ecx<'_, '_, '_>) {
157160 } ) ;
158161}
159162
163+ // Implements the `etch` cheatcode for PVM.
164+ fn etch_call ( target : & Address , new_runtime_code : & Bytes , ecx : Ecx < ' _ , ' _ , ' _ > ) -> Result {
165+ let origin_address = H160 :: from_slice ( ecx. tx . caller . as_slice ( ) ) ;
166+ let origin_account = AccountId :: to_fallback_account_id ( & origin_address) ;
167+
168+ execute_with_externalities ( |externalities| {
169+ externalities. execute_with ( || {
170+ let contract_blob = Pallet :: < Runtime > :: try_upload_pvm_code (
171+ origin_account. clone ( ) ,
172+ new_runtime_code. to_vec ( ) ,
173+ BalanceOf :: < Runtime > :: max_value ( ) ,
174+ & ExecConfig :: new_substrate_tx ( ) ,
175+ )
176+ . map_err ( |_| <& str as Into < Error > >:: into ( "Could not upload PVM code" ) ) ?
177+ . 0 ;
178+ let mut contract_info = if let Some ( contract_info) =
179+ AccountInfo :: < Runtime > :: load_contract ( & H160 :: from_slice ( target. as_slice ( ) ) )
180+ {
181+ contract_info
182+ } else {
183+ ContractInfo :: < Runtime > :: new (
184+ & origin_address,
185+ System :: account_nonce ( origin_account) ,
186+ * contract_blob. code_hash ( ) ,
187+ )
188+ . map_err ( |_| <& str as Into < Error > >:: into ( "Could not create contract info" ) ) ?
189+ } ;
190+ contract_info. code_hash = * contract_blob. code_hash ( ) ;
191+ AccountInfo :: < Runtime > :: insert_contract (
192+ & H160 :: from_slice ( target. as_slice ( ) ) ,
193+ contract_info,
194+ ) ;
195+ Ok :: < ( ) , Error > ( ( ) )
196+ } )
197+ } ) ?;
198+ Ok ( Default :: default ( ) )
199+ }
200+
160201fn set_timestamp ( new_timestamp : U256 , ecx : Ecx < ' _ , ' _ , ' _ > ) {
161202 // Set timestamp in EVM context.
162203 ecx. block . timestamp = new_timestamp;
@@ -251,6 +292,12 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
251292
252293 Ok ( Default :: default ( ) )
253294 }
295+ t if using_pvm && is :: < etchCall > ( t) => {
296+ let & etchCall { ref target, ref newRuntimeBytecode } =
297+ cheatcode. as_any ( ) . downcast_ref ( ) . unwrap ( ) ;
298+ etch_call ( target, newRuntimeBytecode, ccx. ecx ) ?;
299+ cheatcode. dyn_apply ( ccx, executor)
300+ }
254301 // Not custom, just invoke the default behavior
255302 _ => cheatcode. dyn_apply ( ccx, executor) ,
256303 }
0 commit comments