kabu_types_entities/
account_nonce_balance.rs

1use alloy_primitives::{Address, U256};
2use std::collections::HashMap;
3
4#[derive(Debug, Clone, Default)]
5pub struct AccountNonceAndBalances {
6    nonce: u64,
7    balance: HashMap<Address, U256>,
8}
9
10impl AccountNonceAndBalances {
11    pub fn new() -> Self {
12        Self { nonce: 0, balance: HashMap::new() }
13    }
14
15    pub fn get_nonce(&self) -> u64 {
16        self.nonce
17    }
18
19    pub fn set_nonce(&mut self, nonce: u64) -> &mut Self {
20        self.nonce = nonce;
21        self
22    }
23
24    pub fn set_balance(&mut self, token: Address, balance: U256) -> &mut Self {
25        let entry = self.balance.entry(token).or_default();
26        *entry = balance;
27        self
28    }
29
30    pub fn add_balance(&mut self, token: Address, balance: U256) -> &mut Self {
31        let entry = self.balance.entry(token).or_default();
32        if let Some(value) = entry.checked_add(balance) {
33            *entry = value
34        }
35        self
36    }
37
38    pub fn sub_balance(&mut self, token: Address, balance: U256) -> &mut Self {
39        let entry = self.balance.entry(token).or_default();
40        if let Some(value) = entry.checked_sub(balance) {
41            *entry = value
42        }
43        self
44    }
45
46    pub fn get_eth_balance(&self) -> U256 {
47        self.balance.get(&Address::default()).cloned().unwrap_or_default()
48    }
49    pub fn get_balance(&self, token_address: &Address) -> U256 {
50        self.balance.get(token_address).cloned().unwrap_or_default()
51    }
52}
53
54#[derive(Debug, Clone, Default)]
55pub struct AccountNonceAndBalanceState {
56    accounts: HashMap<Address, AccountNonceAndBalances>,
57}
58
59impl AccountNonceAndBalanceState {
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    pub fn add_account(&mut self, account: Address) -> &mut AccountNonceAndBalances {
65        self.accounts.entry(account).or_default()
66    }
67
68    pub fn get_account(&self, account: &Address) -> Option<&AccountNonceAndBalances> {
69        self.accounts.get(account)
70    }
71
72    pub fn get_mut_account(&mut self, account: &Address) -> Option<&mut AccountNonceAndBalances> {
73        self.accounts.get_mut(account)
74    }
75
76    pub fn get_accounts_vec(&self) -> Vec<Address> {
77        self.accounts.keys().copied().collect()
78    }
79
80    pub fn is_monitored(&self, account: &Address) -> bool {
81        self.accounts.contains_key(account)
82    }
83
84    pub fn get_entry_or_default(&mut self, account: Address) -> &mut AccountNonceAndBalances {
85        self.accounts.entry(account).or_default()
86    }
87}
kabu_types_entities/account_nonce_balance.rs

kabu_types_entities/
account_nonce_balance.rs
