Skip to content
Merged
4 changes: 4 additions & 0 deletions crates/evm/core/src/backend/cow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,10 @@ impl DatabaseExt for CowBackend<'_> {
self.backend.cached_accounts()
}

fn cached_storage(&self, address: Address) -> Option<Map<U256, U256>> {
self.backend.cached_storage(address)
}

fn set_blockhash(&mut self, block_number: U256, block_hash: B256) {
self.backend.to_mut().set_blockhash(block_number, block_hash);
}
Expand Down
13 changes: 13 additions & 0 deletions crates/evm/core/src/backend/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -356,6 +356,11 @@ pub trait DatabaseExt: Database<Error = DatabaseError> + DatabaseCommit + Debug
/// Returns all accounts in the memory database cache
fn cached_accounts(&self) -> Vec<Address>;

/// Returns the cached storage for an account from the memory database cache.
/// This returns storage that was written during contract construction before
/// startup migration, which may not yet be in the journaled state.
fn cached_storage(&self, address: Address) -> Option<Map<U256, U256>>;

/// Ensures that `account` is allowed to execute cheatcodes
///
/// Returns an error if [`Self::has_cheatcode_access`] returns `false`
Expand Down Expand Up @@ -1535,6 +1540,14 @@ impl DatabaseExt for Backend {
.collect()
}

fn cached_storage(&self, address: Address) -> Option<Map<U256, U256>> {
self.mem_db
.cache
.accounts
.get(&address)
.map(|acc| acc.storage.iter().map(|(k, v)| (*k, *v)).collect())
}

fn set_blockhash(&mut self, block_number: U256, block_hash: B256) {
if let Some(db) = self.active_fork_db_mut() {
db.cache.block_hashes.insert(block_number.saturating_to(), block_hash);
Expand Down
11 changes: 11 additions & 0 deletions crates/forge/tests/it/revive/migration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,3 +176,14 @@ async fn test_initial_contract_deployment(#[case] runtime_mode: ReviveRuntimeMod
Filter::new("testInitialContractMigration", "InitialMigrationTest", ".*/revive/.*");
TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
}

#[rstest]
#[case::pvm(ReviveRuntimeMode::Pvm)]
#[case::evm(ReviveRuntimeMode::Evm)]
#[tokio::test(flavor = "multi_thread")]
async fn test_contract_with_constructor_deployment(#[case] runtime_mode: ReviveRuntimeMode) {
let runner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
let filter =
Filter::new("testConstructorArgPreserved", "ContractArgsMigrationTest", ".*/revive/.*");
TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
}
90 changes: 69 additions & 21 deletions crates/revive-strategy/src/cheatcodes/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -354,15 +354,15 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
}
t if using_revive && is::<rollCall>(t) => {
let &rollCall { newHeight } = cheatcode.as_any().downcast_ref().unwrap();
let new_block_number: u64 = newHeight.try_into().expect("Block number exceeds u32");
let new_block_number: u64 = newHeight.saturating_to();

// blockhash should be the same on both revive and revm sides, so fetch it before
// changing the block number.
let prev_new_height_hash = ccx
.ecx
.journaled_state
.database
.block_hash(new_block_number - 1)
.block_hash(new_block_number.saturating_sub(1))
.expect("Should not fail");
let new_height_hash = ccx
.ecx
Expand Down Expand Up @@ -695,10 +695,11 @@ fn select_revive(
}

for address in accounts {
tracing::info!("Migrating account {:?} (is_test_contract: {})", address, test_contract_addr == Some(address));
tracing::info!("Migrating account {:?} (is_test_contract: {})", address, test_contract_addr == Some(address));
let acc = data.journaled_state.load_account(address).expect("failed to load account");
let amount = acc.data.info.balance;
let nonce = acc.data.info.nonce;

let amount = acc.info.balance;
let nonce = acc.info.nonce;
let account = H160::from_slice(address.as_slice());
let account_id =
AccountId32Mapper::<Runtime>::to_fallback_account_id(&account);
Expand All @@ -710,7 +711,7 @@ fn select_revive(
a.nonce = nonce.min(u32::MAX.into()).try_into().expect("shouldn't happen");
});

if let Some(bytecode) = acc.data.info.code.as_ref() {
if let Some(bytecode) = acc.info.code.as_ref() {
let account_h160 = H160::from_slice(address.as_slice());

// Skip if contract already exists in pallet-revive
Expand Down Expand Up @@ -825,26 +826,73 @@ fn select_revive(
}
}
}
if AccountInfo::<Runtime>::load_contract(&account_h160).is_some() {
// Migrate complete account state (storage) for newly created/existing contract
for (slot, storage_slot) in &acc.data.storage {
let slot_bytes = slot.to_be_bytes::<32>();
let value_bytes = storage_slot.present_value.to_be_bytes::<32>();

if !storage_slot.present_value.is_zero() {
let _ = Pallet::<Runtime>::set_storage(
account_h160,
slot_bytes,
Some(value_bytes.to_vec()),
);
}
}
if AccountInfo::<Runtime>::load_contract(&account_h160).is_some() {
migrate_contract_storage(data, address, account_h160);
}
}
}
});
}

/// Migrates contract storage from REVM state to pallet-revive.
///
/// Merges storage from two sources:
/// 1. Journaled state: most recent storage values from current execution
/// 2. Database cache: storage written before startup migration, which is run as separate
/// transaction, already committed to cache
///
/// The journaled state takes precedence - cache values are only used for slots
/// not present in the journaled state.
fn migrate_contract_storage(data: Ecx<'_, '_, '_>, address: Address, account_h160: H160) {
use std::collections::HashSet;

// Track which slots we've already migrated from the journaled state
let mut migrated_slots: HashSet<U256> = HashSet::new();

// First, migrate storage from the journaled state (most up-to-date values)
if let Some(account_state) = data.journaled_state.state.get(&address) {
for (slot, storage_slot) in &account_state.storage {
migrated_slots.insert(*slot);

let slot_bytes = slot.to_be_bytes::<32>();
let value = storage_slot.present_value;

if !value.is_zero() {
let _ = Pallet::<Runtime>::set_storage(
account_h160,
slot_bytes,
Some(value.to_be_bytes::<32>().to_vec()),
);
} else {
// Handle case where storage was explicitly cleared
let _ = Pallet::<Runtime>::set_storage(account_h160, slot_bytes, None);
}
}
}

// Then, migrate storage from the database cache for slots NOT in journaled state
if let Some(cached_storage) = data.journaled_state.database.cached_storage(address) {
for (slot, value) in cached_storage {
// Skip slots already migrated from the journaled state
if migrated_slots.contains(&slot) {
continue;
}

let slot_bytes = slot.to_be_bytes::<32>();
if !value.is_zero() {
let _ = Pallet::<Runtime>::set_storage(
account_h160,
slot_bytes,
Some(value.to_be_bytes::<32>().to_vec()),
);
} else {
// Handle case where storage was explicitly cleared
let _ = Pallet::<Runtime>::set_storage(account_h160, slot_bytes, None);
}
}
}
}

fn select_evm(ctx: &mut PvmCheatcodeInspectorStrategyContext, data: Ecx<'_, '_, '_>) {
if !ctx.using_revive {
tracing::info!("already using REVM");
Expand Down Expand Up @@ -991,7 +1039,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
(contract.resolc_bytecode.as_bytes().unwrap().to_vec(), constructor_args.to_vec())
}
crate::ReviveRuntimeMode::Evm => {
// EVM mode: use EVM bytecode directly
// EVM mode: use EVM bytecode directly (includes constructor args)
tracing::info!("running create in EVM mode with EVM bytecode");
(init_code.0.to_vec(), vec![])
}
Expand Down
2 changes: 1 addition & 1 deletion crates/revive-strategy/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ impl TestEnv {
) {
// Set block number in pallet-revive runtime.
self.0.lock().unwrap().externalities.execute_with(|| {
let new_block_number: u64 = new_height.try_into().expect("Block number exceeds u64");
let new_block_number: u64 = new_height.saturating_to();
let digest = System::digest();
if System::block_hash(new_block_number) == H256::zero() {
// First initialize and finalize the parent block to set up correct hashes.
Expand Down
27 changes: 27 additions & 0 deletions testdata/default/revive/EvmToReviveMigration.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -384,3 +384,30 @@ contract InitialMigrationTest is DSTest {
assertEq(storageContract.get(), 42);
}
}

// Simple contract that stores a constructor argument
contract SimpleAddrStorage {
address public storedAddr;

constructor(address _addr) {
storedAddr = _addr;
}
}

contract ContractArgsMigrationTest is DSTest {
Vm constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
address myAddr = vm.addr(1900);

// Expected address for private key 1900
address constant EXPECTED = 0x6Af741AA4Ff39CF3De9A0Cb02A8Bab387E41abFB;
SimpleAddrStorage store = new SimpleAddrStorage(myAddr);

function testConstructorArgPreserved() public {
emit log_named_address("myAddr (from vm.addr)", myAddr);
emit log_named_address("store.storedAddr()", store.storedAddr());
emit log_named_address("EXPECTED", EXPECTED);

assertEq(myAddr, EXPECTED, "vm.addr(1900) should return correct address");
assertEq(store.storedAddr(), myAddr, "Constructor arg should be preserved");
}
}
Loading