Skip auto-funding for accounts that were explicitly dealt

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

@@ -5,7 +5,7 @@ use std::{
 };
 
 use alloy_primitives::{Address, Bytes, map::foldhash::HashMap, ruint::aliases::U256};
-use foundry_cheatcodes::{Ecx, MockCallDataContext, MockCallReturnData};
+use foundry_cheatcodes::{DealRecord, Ecx, MockCallDataContext, MockCallReturnData};
 use foundry_evm::constants::CHEATCODE_ADDRESS;
 use polkadot_sdk::{
     frame_system,
@@ -15,7 +15,7 @@ use polkadot_sdk::{
     },
     pallet_revive_uapi::ReturnFlags,
     polkadot_sdk_frame::prelude::OriginFor,
-    sp_core::H160,
+    sp_core::{H160, U256 as SpU256},
 };
 use revive_env::Runtime;
 
@@ -60,17 +60,36 @@ impl MockHandlerImpl {
         state.mocked_functions = mock_inner.mocked_functions.clone();
     }
 
-    pub(crate) fn fund_pranked_accounts(&self, account: Address) {
-        // Fuzzed prank addresses have no balance, so they won't exist in revive, and
-        // calls will fail, this is not a problem when running in REVM.
-        // TODO: Figure it out why this is still needed.
-        let balance = Pallet::<Runtime>::evm_balance(&H160::from_slice(account.as_slice()));
-        if balance == 0.into() {
-            Pallet::<Runtime>::set_evm_balance(
-                &H160::from_slice(account.as_slice()),
-                u128::MAX.into(),
-            )
-            .expect("Could not fund pranked account");
+    /// Syncs balances for pranked accounts between REVM and pallet-revive.
+    ///
+    /// If the account was explicitly dealt to via vm.deal(), sync that balance to pallet-revive.
+    /// This handles cases where vm.deal() was called in a callback and pallet-revive's balance
+    /// diverged from REVM's balance.
+    ///
+    /// If the account was NOT dealt to and has 0 balance, fund with u128::MAX so fuzzed
+    /// prank addresses can make calls in pallet-revive.
+    pub(crate) fn fund_pranked_accounts(account: Address, eth_deals: &[DealRecord]) {
+        let account_h160 = H160::from_slice(account.as_slice());
+
+        // Check if account was explicitly dealt to via vm.deal()
+        // Use the most recent deal record for this account
+        if let Some(deal) = eth_deals.iter().rev().find(|d| d.address == account) {
+            // Sync the dealt balance to pallet-revive
+            let target_balance =
+                SpU256::from_little_endian(&deal.new_balance.as_le_bytes()).min(u128::MAX.into());
+            let pvm_balance = Pallet::<Runtime>::evm_balance(&account_h160);
+            if pvm_balance != target_balance {
+                Pallet::<Runtime>::set_evm_balance(&account_h160, target_balance)
+                    .expect("Could not sync dealt account balance");
+            }
+            return;
+        }
+
+        // Account was not dealt to - fund with u128::MAX if balance is 0
+        let pvm_balance = Pallet::<Runtime>::evm_balance(&account_h160);
+        if pvm_balance == 0.into() {
+            Pallet::<Runtime>::set_evm_balance(&account_h160, u128::MAX.into())
+                .expect("Could not fund pranked account");
         }
     }
 }

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

@@ -1001,6 +1001,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
             sp_core::U256::from_little_endian(&U256::from(ecx.tx.gas_price).as_le_bytes());
         let mut tracer = Tracer::new(state.expected_calls.clone(), state.expected_creates.clone());
         let caller_h160 = H160::from_slice(input.caller().as_slice());
+        let eth_deals = &state.eth_deals;
 
         let res = ctx.externalities.execute_with(|| {
             tracer.watch_address(&caller_h160);
@@ -1011,7 +1012,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
                     AccountId32Mapper::<Runtime>::to_fallback_account_id(&caller_h160);
                 let origin = OriginFor::<Runtime>::signed(origin_account_id.clone());
                 let evm_value = sp_core::U256::from_little_endian(&input.value().as_le_bytes());
-                mock_handler.fund_pranked_accounts(input.caller());
+                MockHandlerImpl::fund_pranked_accounts(input.caller(), eth_deals);
                 if !exists {
                     let nonce = ecx
                         .journaled_state
@@ -1185,6 +1186,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
         let caller_h160 = H160::from_slice(call.caller.as_slice());
 
         let mut tracer = Tracer::new(state.expected_calls.clone(), state.expected_creates.clone());
+        let eth_deals = &state.eth_deals;
         let res = ctx.externalities.execute_with(|| {
             // Watch the caller's address so its nonce changes get tracked in prestate trace
             tracer.watch_address(&caller_h160);
@@ -1193,7 +1195,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
                 let origin = OriginFor::<Runtime>::signed(
                     AccountId32Mapper::<Runtime>::to_fallback_account_id(&caller_h160),
                 );
-                mock_handler.fund_pranked_accounts(call.caller);
+                MockHandlerImpl::fund_pranked_accounts(call.caller, eth_deals);
 
                 let evm_value = sp_core::U256::from_little_endian(&call.call_value().as_le_bytes());
                 let target = H160::from_slice(call.target_address.as_slice());

testdata/default/revive/Prank.t.sol

@@ -675,3 +675,40 @@ contract Counter {
         number++;
     }
 }
+
+contract FundPrankedAccountsReproTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    address alice = address(0xA11CE);
+
+    function test_PrankDealtAccountWithZeroBalance() public {
+        vm.deal(alice, 3 ether);
+        assertEq(alice.balance, 3 ether, "alice should have 3 ether after deal");
+
+        vm.deal(alice, 0);
+        assertEq(alice.balance, 0, "alice should have 0 balance after second deal");
+
+        vm.startPrank(alice);
+        assertEq(alice.balance, 0, "alice balance should remain 0, not be reset to u128::MAX");
+        vm.stopPrank();
+    }
+
+    function test_PrankFuzzedAddressGetsFunded() public {
+        address fuzzedAddr = address(0xF022ED);
+        uint256 initialBalance = fuzzedAddr.balance;
+
+        vm.prank(fuzzedAddr);
+
+        assertTrue(fuzzedAddr.balance >= initialBalance, "fuzzed address should be auto-funded");
+    }
+
+    function test_PrankDealtAccountWithBalance() public {
+        vm.deal(alice, 5 ether);
+        assertEq(alice.balance, 5 ether, "alice should have 5 ether");
+
+        vm.startPrank(alice);
+        assertEq(alice.balance, 5 ether, "alice balance should remain 5 ether");
+        vm.stopPrank();
+    }
+
+}