diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 3d55d29922e8e..d8dec93d041ae 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -1817,6 +1817,11 @@ impl Inspector> for Cheatcodes { outcome.result.output = Error::encode(msg); } } + + // Sync REVM state back to pallet-revive if this call was executed in REVM + if outcome.result.is_ok() { + self.strategy.runner.revive_call_end(self, ecx, call); + } } fn create(&mut self, ecx: Ecx, mut input: &mut CreateInputs) -> Option { diff --git a/crates/cheatcodes/src/strategy.rs b/crates/cheatcodes/src/strategy.rs index e286d66c0fc1b..b1633e72474de 100644 --- a/crates/cheatcodes/src/strategy.rs +++ b/crates/cheatcodes/src/strategy.rs @@ -261,6 +261,9 @@ pub trait CheatcodeInspectorStrategyExt { ) -> Option { None } + + fn revive_call_end(&self, _state: &mut crate::Cheatcodes, _ecx: Ecx, _call: &CallInputs) {} + // Remove duplicate accesses in storage_recorder fn revive_remove_duplicate_account_access(&self, _state: &mut crate::Cheatcodes) {} } diff --git a/crates/revive-strategy/src/cheatcodes/mod.rs b/crates/revive-strategy/src/cheatcodes/mod.rs index 1d70b7399a4a3..fad54eb2ca43f 100644 --- a/crates/revive-strategy/src/cheatcodes/mod.rs +++ b/crates/revive-strategy/src/cheatcodes/mod.rs @@ -699,6 +699,19 @@ fn select_revive(ctx: &mut PvmCheatcodeInspectorStrategyContext, data: Ecx<'_, ' ); } } + // Migrate complete account state (storage) for newly created 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::::set_storage( + account_h160, + slot_bytes, + Some(value_bytes.to_vec()), + ); + } + } } } } @@ -1131,6 +1144,29 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector } } } + + fn revive_call_end( + &self, + state: &mut foundry_cheatcodes::Cheatcodes, + ecx: Ecx<'_, '_, '_>, + call: &CallInputs, + ) { + let ctx = get_context_ref_mut(state.strategy.context.as_mut()); + + // Skip storage sync if: in PVM mode AND no test contract + if ctx.using_pvm + && ecx + .journaled_state + .database + .get_test_contract_address() + .map(|addr| call.bytecode_address != addr && call.target_address != addr) + .unwrap_or(true) + { + return; + } + + apply_revm_storage_diff(ecx, call.target_address); + } } fn post_exec( @@ -1188,3 +1224,44 @@ fn get_context_ref_mut( ) -> &mut PvmCheatcodeInspectorStrategyContext { ctx.as_any_mut().downcast_mut().expect("expected PvmCheatcodeInspectorStrategyContext") } + +/// Applies REVM storage diffs to pallet-revive (REVM → pallet-revive sync) +/// Note: Balance/nonce are NOT synced here as they're handled by migration in select_revive() +fn apply_revm_storage_diff(ecx: Ecx<'_, '_, '_>, address: Address) { + let Some(account_state) = ecx.journaled_state.state.get(&address) else { + return; + }; + + let h160_address = H160::from_slice(address.as_slice()); + + // Check if contract exists in pallet-revive before applying storage diffs + let contract_exists = execute_with_externalities(|externalities| { + externalities + .execute_with(|| AccountInfo::::load_contract(&h160_address).is_some()) + }); + + if !contract_exists { + return; + } + + execute_with_externalities(|externalities| { + externalities.execute_with(|| { + for (slot, storage_slot) in &account_state.storage { + if storage_slot.is_changed() { + let slot_bytes = slot.to_be_bytes::<32>(); + let new_value = storage_slot.present_value; + + if !new_value.is_zero() { + let _ = Pallet::::set_storage( + h160_address, + slot_bytes, + Some(new_value.to_be_bytes::<32>().to_vec()), + ); + } else { + let _ = Pallet::::set_storage(h160_address, slot_bytes, None); + } + } + } + }) + }); +} diff --git a/testdata/default/revive/EvmToReviveMigration.t.sol b/testdata/default/revive/EvmToReviveMigration.t.sol index 8e7cd8c3e06a6..15610458f80dd 100644 --- a/testdata/default/revive/EvmToReviveMigration.t.sol +++ b/testdata/default/revive/EvmToReviveMigration.t.sol @@ -170,8 +170,7 @@ contract EvmReviveMigrationTest is DSTest { vm.pvm(true); - // TODO: Enable after contract storage migration is on place - // assertEq(storageContract.get(), 42); + assertEq(storageContract.get(), 42); storageContract.set(100); assertEq(storageContract.get(), 100);