Add support for etch cheatcode

crates/anvil-polkadot/Cargo.toml

@@ -20,7 +20,7 @@ path = "bin/main.rs"
 [dependencies]
 # foundry internal
 substrate-runtime = { path = "substrate-runtime" }
-polkadot-sdk = {git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", default-features = false, features = [
+polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/etch_dependencies", default-features = false, features = [
     "sc-allocator",
 	"sc-basic-authorship",
 	"sc-block-builder",

crates/anvil-polkadot/substrate-runtime/Cargo.toml

@@ -12,7 +12,7 @@ license.workspace = true
 [dependencies]
 array-bytes = { version = "6.2.2", default-features = false }
 codec = { version = "3.7.5", default-features = false, package = "parity-scale-codec" }
-polkadot-sdk = {git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", default-features = false, features = [
+polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/etch_dependencies", default-features = false, features = [
 	"pallet-balances",
 	"pallet-revive",
 	"pallet-sudo",
@@ -27,7 +27,7 @@ scale-info = { version = "2.11.6", default-features = false }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 
 [build-dependencies]
-polkadot-sdk = {git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", default-features = false, optional = true, features = ["substrate-wasm-builder"] }
+polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/etch_dependencies", default-features = false, optional = true, features = ["substrate-wasm-builder"] }
 
 [features]
 default = ["std"]

crates/forge/tests/cli/main.rs

@@ -27,6 +27,7 @@ mod revive_cmd;
 mod revive_compiler;
 mod revive_config;
 mod revive_create;
+mod revive_etch;
 mod revive_inspect;
 mod revive_mock_call;
 mod revive_mock_calls;

crates/forge/tests/cli/revive_etch.rs

@@ -0,0 +1,117 @@
+// Just a copy of cheatcodes Prank.t.sol adapted to work with pvm backend.
+// The adaptions are only to switch back and forth between evm and pvm.
+forgetest!(mock_call, |prj, cmd| {
+    prj.insert_ds_test();
+    prj.insert_vm();
+    prj.insert_console();
+    prj.add_source(
+        "MockCall.t.sol",
+        r#"
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+
+import "./test.sol";
+import "./Vm.sol";
+import {console} from "./console.sol";
+
+contract Adder {
+    function add(uint256 a, uint256 b) public pure returns (uint256) {
+        return a * b; // Intentional bug to verify etch works
+    }
+}
+
+contract NestedAdder {
+    uint256 public inner_a;
+    uint256 public inner_b;
+    constructor(uint256 a, uint256 b) {
+        inner_a = a;
+        inner_b = b;
+    }
+
+    function nested_call(address target) public returns (uint256) {
+        // Perform the add call on the target contract address
+        (bool success, bytes memory data) = target.call(abi.encodeWithSignature("add(uint256,uint256)", inner_a, inner_b));
+        require(success, "Nested call failed");
+        uint256 result = abi.decode(data, (uint256));
+        assert(success);
+        return result;
+    }
+}
+
+
+contract EtchTest is DSTest {
+    Vm constant vm = Vm(address(bytes20(uint160(uint256(keccak256("hevm cheat code"))))));
+    // This is the bytecode for the  correct Adder contract above compiled with resolc.
+    bytes constant code = hex"50564d00008c0c000000000000010700c13000c00080047106000000000e0000001c0000002a000000350000004700000063616c6c5f646174615f636f707963616c6c5f646174615f6c6f616463616c6c5f646174615f73697a657365616c5f72657475726e7365745f696d6d757461626c655f6461746176616c75655f7472616e7366657272656405110287920463616c6c879e066465706c6f79068bec25028a531a001f004b007c00a500ae00c100cd00fa005e01630192019701bb0212036003b103c0033c048d04b2046c05ec05f5050106180627063d0660061a0769079e07ab07c607ca07080850099511f07b10087b158475010a02013d0700000251050750100209501004b3009511807b10787b15707b1668951580008411e04911384911304911208000330740951820491128501006e6084911584911504911484911408317400a0501821750821858821948821a40d49808d4a707d4870752072e6417501008ec008217188218108219088216d49707d48609d47909989920d48707977720d497075107090050100a61056467330850100c75098377330833090a283e076467330850100e62098378330733093300100a03019511a07b10587b15507b16489515608411e064164916384916304916208000330740956820491628501012370839070000025317045f9517e08477e07b67186471837733080a010182671880771c977720140800000000f7021677ab8736330014951120fe7b10d8017b15d0017b16c8019515e0018411e04921b8014921b0014921a8014921a0018317a0010a05288d02501016a3049511c07b10387b15307b16289515408411f0647664173300189511f07b10087b156475330820330740330048288f0850121a3a063200828910828a18828b088288d4ba0ad4980bd4ab0b98bb20d4a909979920d4b9095209449511c07b10387b15307b16289515408411e06476838883170a01821718821810821908821a7b67187b68107b69087b6a9551c08210388215308216289511403200009511e0fe7b1018017b1510017b160801951520018411f0828310828208829a08828c829b7b1b20829410d32a067b1638d82a06d8cb00821b38dab006828b188298187b1838c93409c969087b1828d869087b1330d8340664b4821938c9b909c96909c98909c92a08c908087b1c38821a20c9ca0a8e8b88aa4085aa01db8b0a8f98821c288ecbdb980bd49c08db8b0a510ace0064768217387b17c0007b127b12c8008217307b17d0009517e0009518c0007b14087b14d80050101ce3fe8217e8007b17288217e0007b17208217f8007b17188217f0007b1710821838958720d88708821ac88a09d8a90ada880a821830c88a0ad88a08821b08c88b0b7b1a707b19687b1760951780009518607b1b7850101e8cfe821790008218980082198000821a88007b67307b68387b69207b6a288217107b67108217187b67188217207b678217287b67089551e0fe821018018215100182160801951120013200501020a602828a10828b18828c088289d4cb0bd4a908d4b808988820d4ba0a97aa20d4a80852083f9511d07b10287b15209515308411f0827a18827810827b0882777b177b1b087b181064187b1a18649750102280059551d08210288215209511303200008218108217087b87088217187b877b861082177b87189551808210788215708216689511800032008217b0018218b8018219a801821aa001d49808d4a707d487075207e70138070000024921380149213001492128017b1720014921580149215001492140010495176001951840019519200149214801501026c7fd82126001821768018218700182197801821b8001821c880182169001821a98017b1ad8007b16d0007b1cc8007b1bc0007b19f8007b18f0007b17e800951700019518e0009519c0007b12e000501028d801821700017b1738821708017b1730821710017b1728821718017b17209517a00050102ab6fc8217b8007b17188216b0008218a8007b1810821aa0007b1a088219207b19588219287b19508219307b19488219387b19407b17787b16707b1868951780009518609519407b1a6033002c9511807b10787b15707b1668951580008411f07b171082878292828b08829308957a207b1a18d87a06c86b0a7b1a08d8ba0cda660c828a10828818829410829918c8ca06d8a60cc88c0c7b1c7b18387b1a307b1b287b17207b19587b14507b13489517409518207b124033002428f7fd821918821b10821008d49b07d46008d47808988820d46707977720d4870752075d646482178800821898007b183882138000821a9000d3b706d8b70cd80308da680cc94a06c9c602d8c606d84a0a821c38c99c0cc9ac0cc96c0cc9b707c98707c90306d4c707d42608d47808d42707988820977720d487075107080050102e19640764685010302e048378836933073300320a03019511f87b103308100002838833070133093300340a03013308491718491710491708490732004911184911104901113307046418491108501038f402390804000256183f0b200304000240013308100002838833070133092433003a0a03019511e87b10107b15087b16828b188294188282828c08829a088295828610829810c8ca09c82503d85305c85909d3a900d8a90ada050ac86805c85a0ad85a00c8b404d88508c84808c88000d86a05d3b008d8b00bda850bd3a606d46808d3c906d8c90cd82305db6c05db8b0552051b7b737b79087b7a107b70188210108215088216951118320033003c9511b07b10487b15409515508411f0491130491128491120140700000000717b484e9518207b173833073300362815029511807b10787b15707b1668951580008411f08282828308828410828818829a829b08829c1082991864767b19187b1c107b1b087b1a7b18387b14307b132895174095182064197b122050103efcfe821750821858821940821a487b67107b68187b697b6a089551808210788215708216689511800032009511f87b10330750104067f89511f87b103307015010425af89511c07b10387b15307b16289515408411f064766417501044a2f95012460632008217108218188219821a087b67107b68187b697b6a089551c0821038821530821628951140320239080800025108c0f8330730000383770a0428b3f87c78017c797c7a027c7b03978808d4980897aa1097bb18d4ba0ad4a8087c79057c7a047c7b067c7c07979908d4a90997bb1097cc18d4cb0bd4b909979920d489027c79097c7a087c7b0a7c7c0b979908d4a90997bb1097cc18d4cb0bd4b9097c7a0d7c7b0c7c7c0e7c780f97aa08d4ba0a97cc10978818d4c808d4a808978820d498037c78117c7a107c7b127c7c13978808d4a80897bb1097cc18d4cb0bd4b8087c7a157c7b147c7c167c791797aa08d4ba0a97cc10979918d4c909d4a909979920d4890a7c78197c79187c7b1a7c7c1b978808d4980897bb1097cc18d4cb0bd4b8087c791d7c7b1c7c7c1e7c771f979908d4b90997cc10977718d4c707d49707977720d487076f776fa86f396f2a7b5a187b59107b58087b57821008821595111032009511d87b10207b15187b161082897b19088289087b1982851082861833082050104ad3006f686f59821a6faa821b086fbb787b18787a10787908787898bc38787c1f98bc30787c1e98bc28787c1d98bc20787c1c98bc18787c1b98bc10787c1a98bb08787b1998ab38787b1798ab30787b1698ab28787b1598ab20787b1498ab18787b1398ab10787b1298aa08787a11989a38787a0f989a30787a0e989a28787a0d989a20787a0c989a18787a0b989a10787a0a98990878790998893878790798893078790698892878790598892078790498891878790398891078790298880878780182102082151882161095112832008b7910520931c8780883881f8488e05638000001253309040002390a040002ae8a093d080400020133081000028377c887073200004969488424892421494892344992a490a4244992a1423515aa3449929290248448523549004944442422224a4892a4925492849294244992244955929294246944442449254992244992a424499294a494244a1489884488442291242549922489244444244992a42422220a494a9224499224a9a41492a449524a924411114922444444848888101111111111111111112112894422442212499224491249922491549224499224494992244992244992a4884424499224499224499251a14a5328a525a9242921443094a4494a4a4a922449922449922449929224a52449920490942449a4242529494992129224294912a5424a4892a4429494a424490a112a294992244992244992244992244992244992244992244992244992244992244992aa24a5942425250955294949922449922449922449922449922449922449922449529224082161280500";
+
+    // Test etching code into an existing contract instance works correctly.
+    function testEtchExistingContract() public {
+        vm.pvm(true);
+        Adder adder = new Adder();
+
+        // Without etch, the add function is broken
+        uint256 buggy_result = adder.add(1, 2);
+        assertEq(buggy_result, 2);
+
+        // Etch the correct bytecode into the existing contract
+        vm.etch(address(adder), code);
+        uint256 result = adder.add(1, 2);
+        assertEq(result, 3);
+
+        // Verify that nested calls also work correctly after etch
+        uint256 nested_call_result = (new NestedAdder(1, 2)).nested_call(address(adder));
+        assertEq(nested_call_result, 3);
+    }
+
+    // Test etching code into any arbitrary address works correctly.
+    function testEtchAnyContract() public {
+        vm.pvm(true);
+        // Etch the correct bytecode into an arbitrary address
+        address target = address(7070707);
+        vm.etch(target, code);
+         (bool success, bytes memory output) = target.call(abi.encodeWithSignature("add(uint256,uint256)", 1, 2));
+        uint256 result1 = abi.decode(output, (uint256));
+
+        assert(success);
+        assertEq(result1, 3);
+
+        uint256 nested_call_result = (new NestedAdder(1, 2)).nested_call(address(target));
+        assertEq(nested_call_result, 3);
+
+        // Etch into the zero address as well to verify it works for reserved addresses
+        address target2 = address(0);
+        vm.etch(target2, code);
+         (bool success2, bytes memory output2) = target2.call(abi.encodeWithSignature("add(uint256,uint256)", 1, 2));
+        uint256 result2 = abi.decode(output2, (uint256));
+
+        assert(success2);
+        assertEq(result2, 3);
+
+        uint256 nested_call_result2 = (new NestedAdder(1, 2)).nested_call(address(target2));
+        assertEq(nested_call_result2, 3);
+    }
+}
+"#,
+    )
+    .unwrap();
+
+    let res = cmd.args(["test", "--resolc", "--resolc-startup"]).assert_success();
+
+    res.stderr_eq(str![""]).stdout_eq(str![[r#"
+[COMPILING_FILES] with [SOLC_VERSION]
+[SOLC_VERSION] [ELAPSED]
+Compiler run successful!
+[COMPILING_FILES] with [RESOLC_VERSION]
+[RESOLC_VERSION] [ELAPSED]
+Compiler run successful!
+
+Ran 2 tests for src/MockCall.t.sol:EtchTest
+[PASS] testEtchAnyContract() ([GAS])
+[PASS] testEtchExistingContract() ([GAS])
+Suite result: ok. 2 passed; 0 failed; 0 skipped; [ELAPSED]
+
+Ran 1 test suite [ELAPSED]: 2 tests passed, 0 failed, 0 skipped (2 total tests)
+
+"#]]);
+});

crates/revive-env/Cargo.toml

@@ -16,7 +16,7 @@ std = ["polkadot-sdk/std"]
 codec = { version = "3.7.5", default-features = false, package = "parity-scale-codec" }
 scale-info = { version = "2.11.6", default-features = false }
 
-polkadot-sdk = {git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", features = [
+polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/etch_dependencies", features = [
     "experimental",
     "runtime",
     "polkadot-runtime-common",

crates/revive-strategy/Cargo.toml

@@ -35,8 +35,9 @@ semver.workspace = true
 
 codec = { version = "3.7.5", default-features = false, package = "parity-scale-codec" }
 scale-info = { version = "2.11.6", default-features = false }
+hex = "0.4"
 
-polkadot-sdk = {git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", features = [
+polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/etch_dependencies", features = [
     "experimental",
     "runtime",
     "polkadot-runtime-common",

crates/revive-strategy/src/cheatcodes/mod.rs

@@ -11,31 +11,34 @@ use alloy_sol_types::SolValue;
 use foundry_cheatcodes::{
     Broadcast, BroadcastableTransactions, CheatcodeInspectorStrategy,
     CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyRunner, CheatsConfig, CheatsCtxt,
-    CommonCreateInput, DealRecord, Ecx, EvmCheatcodeInspectorStrategyRunner, Result,
-    Vm::{dealCall, getNonce_0Call, pvmCall, rollCall, setNonceCall, setNonceUnsafeCall, warpCall},
+    CommonCreateInput, DealRecord, Ecx, Error, EvmCheatcodeInspectorStrategyRunner, Result,
+    Vm::{
+        dealCall, etchCall, getNonce_0Call, pvmCall, rollCall, setNonceCall, setNonceUnsafeCall,
+        warpCall,
+    },
 };
 
 use foundry_common::sh_err;
 use foundry_compilers::resolc::dual_compiled_contracts::DualCompiledContracts;
 use revive_env::{AccountId, Runtime, System, Timestamp};
 
+use crate::{
+    cheatcodes::mock_handler::MockHandlerImpl, execute_with_externalities, trace,
+    tracing::apply_prestate_trace,
+};
+use alloy_eips::eip7702::SignedAuthorization;
 use polkadot_sdk::{
     frame_support::traits::{Currency, fungible::Mutate},
     pallet_balances,
     pallet_revive::{
-        self, AddressMapper, BalanceOf, BalanceWithDust, Code, Config, ExecConfig, Pallet,
+        self, AccountInfo, AddressMapper, BalanceOf, BalanceWithDust, Code, Config, ContractInfo,
+        ExecConfig, Executable, Pallet,
     },
     polkadot_runtime_common::Bounded,
     polkadot_sdk_frame::prelude::OriginFor,
     sp_core::{self, H160},
     sp_weights::Weight,
 };
-
-use crate::{
-    cheatcodes::mock_handler::MockHandlerImpl, execute_with_externalities, trace,
-    tracing::apply_prestate_trace,
-};
-use alloy_eips::eip7702::SignedAuthorization;
 use revm::{
     bytecode::opcode as op,
     context::{CreateScheme, JournalTr},
@@ -157,6 +160,44 @@ fn set_block_number(new_height: U256, ecx: Ecx<'_, '_, '_>) {
     });
 }
 
+// Implements the `etch` cheatcode for PVM.
+fn etch_call(target: &Address, new_runtime_code: &Bytes, ecx: Ecx<'_, '_, '_>) -> Result {
+    let origin_address = H160::from_slice(ecx.tx.caller.as_slice());
+    let origin_account = AccountId::to_fallback_account_id(&origin_address);
+
+    execute_with_externalities(|externalities| {
+        externalities.execute_with(|| {
+            let contract_blob = Pallet::<Runtime>::try_upload_pvm_code(
+                origin_account.clone(),
+                new_runtime_code.to_vec(),
+                BalanceOf::<Runtime>::max_value(),
+                &ExecConfig::new_substrate_tx(),
+            )
+            .map_err(|_| <&str as Into<Error>>::into("Could not upload PVM code"))?
+            .0;
+            let mut contract_info = if let Some(contract_info) =
+                AccountInfo::<Runtime>::load_contract(&H160::from_slice(target.as_slice()))
+            {
+                contract_info
+            } else {
+                ContractInfo::<Runtime>::new(
+                    &origin_address,
+                    System::account_nonce(origin_account),
+                    *contract_blob.code_hash(),
+                )
+                .map_err(|_| <&str as Into<Error>>::into("Could not create contract info"))?
+            };
+            contract_info.code_hash = *contract_blob.code_hash();
+            AccountInfo::<Runtime>::insert_contract(
+                &H160::from_slice(target.as_slice()),
+                contract_info,
+            );
+            Ok::<(), Error>(())
+        })
+    })?;
+    Ok(Default::default())
+}
+
 fn set_timestamp(new_timestamp: U256, ecx: Ecx<'_, '_, '_>) {
     // Set timestamp in EVM context.
     ecx.block.timestamp = new_timestamp;
@@ -251,6 +292,12 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
 
                 Ok(Default::default())
             }
+            t if using_pvm && is::<etchCall>(t) => {
+                let &etchCall { ref target, ref newRuntimeBytecode } =
+                    cheatcode.as_any().downcast_ref().unwrap();
+                etch_call(target, newRuntimeBytecode, ccx.ecx)?;
+                cheatcode.dyn_apply(ccx, executor)
+            }
             // Not custom, just invoke the default behavior
             _ => cheatcode.dyn_apply(ccx, executor),
         }