kabu_evm_utils/
revm_balances.rs

1use crate::remv_db_direct_access::calc_hashmap_cell;
2use crate::{nweth, NWETH};
3use alloy::providers::ext::AnvilApi;
4use alloy::{network::Network, primitives::Address, providers::Provider, sol_types::private::U256};
5use eyre::{eyre, Result};
6use kabu_defi_abi::IERC20::IERC20Instance;
7use kabu_defi_address_book::TokenAddressEth;
8use kabu_evm_db::KabuDBType;
9use kabu_node_debug_provider::DebugProviderExt;
10use tracing::error;
11
12pub struct BalanceCheater {}
13
14#[allow(dead_code)]
15impl BalanceCheater {
16    pub fn get_balance_cell(token: Address, owner: Address) -> Result<U256> {
17        match token {
18            TokenAddressEth::WETH => Ok(calc_hashmap_cell(U256::from(3u32), U256::from_be_slice(owner.as_slice()))),
19            TokenAddressEth::USDT => Ok(calc_hashmap_cell(U256::from(2u32), U256::from_be_slice(owner.as_slice()))),
20            TokenAddressEth::USDC => Ok(calc_hashmap_cell(U256::from(3u32), U256::from_be_slice(owner.as_slice()))),
21            TokenAddressEth::WSTETH => Ok(calc_hashmap_cell(U256::from(0u32), U256::from_be_slice(owner.as_slice()))),
22            TokenAddressEth::STETH => Ok(calc_hashmap_cell(U256::from(3u32), U256::from_be_slice(owner.as_slice()))),
23            _ => Err(eyre!("ADDRESS_CELL_UNKNOWN")),
24        }
25    }
26
27    pub async fn get_anvil_token_balance<P, N>(client: P, token: Address, owner: Address) -> eyre::Result<U256>
28    where
29        N: Network,
30        P: Provider<N> + DebugProviderExt<N> + Send + Sync + Clone + 'static,
31    {
32        let value = client.get_storage_at(token, Self::get_balance_cell(token, owner)?).await?;
33
34        Ok(value)
35    }
36
37    pub async fn set_anvil_token_balance<P, N>(client: P, token: Address, owner: Address, balance: U256) -> eyre::Result<()>
38    where
39        N: Network,
40        P: Provider<N> + Send + Sync + Clone + 'static,
41    {
42        let balance_cell = Self::get_balance_cell(token, owner)?;
43
44        if let Err(e) = client.anvil_set_storage_at(token, balance_cell, balance.into()).await {
45            error!("{e}");
46            return Err(eyre!(e));
47        }
48
49        let new_storage = client.get_storage_at(token, balance_cell).await?;
50
51        if balance != new_storage {
52            error!("{balance} != {new_storage}");
53            return Err(eyre!("STORAGE_NOT_SET"));
54        }
55
56        let token_instance = IERC20Instance::new(token, client.clone());
57
58        let new_balance = token_instance.balanceOf(owner).call_raw().await?;
59        println!("new_balance : {new_balance:?}");
60        if U256::from_be_slice(new_balance.as_ref()) != balance {
61            return Err(eyre!("BALANCE_NOT_SET"));
62        }
63        Ok(())
64    }
65    pub async fn set_anvil_token_balance_float<P, N>(client: P, token: Address, owner: Address, balance: f64) -> eyre::Result<()>
66    where
67        N: Network,
68        P: Provider<N> + Send + Sync + Clone + 'static,
69    {
70        let balance = nweth::NWETH::from_float(balance);
71        Self::set_anvil_token_balance(client, token, owner, balance).await
72    }
73
74    pub fn set_evm_token_balance(db: &mut KabuDBType, token: Address, owner: Address, balance: U256) -> eyre::Result<()> {
75        let balance_cell = calc_hashmap_cell(U256::from(3), U256::from_be_slice(owner.as_slice()));
76
77        db.insert_account_storage(token, balance_cell, balance).map_err(|_| eyre!("ERROR_INSERTING_ACCOUNT_STORAGE"))
78    }
79
80    pub fn set_evm_token_balance_float(db: &mut KabuDBType, token: Address, owner: Address, balance: f64) -> eyre::Result<()> {
81        Self::set_evm_token_balance(db, token, owner, NWETH::from_float(balance))
82    }
83}
kabu_evm_utils/revm_balances.rs

kabu_evm_utils/
revm_balances.rs
