nautilus_model/defi/
wallet.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use std::{collections::HashSet, fmt::Display};
17
18use alloy_primitives::{Address, U256};
19
20use crate::{
21    defi::Token,
22    enums::CurrencyType,
23    types::{AccountBalance, Currency, Money, Quantity},
24};
25
26/// Represents the balance of a specific ERC-20 token held in a wallet.
27///
28/// This struct tracks the raw token amount along with optional USD valuation
29/// and the token metadata.
30#[derive(Debug)]
31pub struct TokenBalance {
32    /// The raw token amount as a 256-bit unsigned integer.
33    pub amount: U256,
34    /// The optional USD equivalent value of the token balance.
35    pub amount_usd: Option<Quantity>,
36    /// The token metadata including chain, address, name, symbol, and decimals.
37    pub token: Token,
38}
39
40impl TokenBalance {
41    /// Creates a new [`TokenBalance`] instance.
42    #[must_use]
43    pub const fn new(amount: U256, token: Token) -> Self {
44        Self {
45            amount,
46            token,
47            amount_usd: None,
48        }
49    }
50
51    /// Converts the raw token amount to a human-readable [`Quantity`].
52    ///
53    /// # Errors
54    ///
55    /// Returns an error if the U256 amount cannot be converted to a `Quantity`.
56    pub fn as_quantity(&self) -> anyhow::Result<Quantity> {
57        Quantity::from_u256(self.amount, self.token.decimals).map_err(Into::into)
58    }
59
60    fn as_money(&self) -> anyhow::Result<Money> {
61        let currency = Currency::new_checked(
62            &self.token.symbol,
63            self.token.decimals,
64            0,
65            &self.token.name,
66            CurrencyType::Crypto,
67        )?;
68        Money::from_u256(self.amount, currency).map_err(Into::into)
69    }
70
71    /// Sets the USD equivalent value for this token balance.
72    pub fn set_amount_usd(&mut self, amount_usd: Quantity) {
73        self.amount_usd = Some(amount_usd);
74    }
75}
76
77impl Display for TokenBalance {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        let quantity = self.as_quantity().unwrap_or_default();
80
81        match &self.amount_usd {
82            Some(usd) => write!(
83                f,
84                "TokenBalance(token={}, amount={}, usd=${:.2})",
85                self.token.symbol,
86                quantity.as_decimal(),
87                usd.as_f64()
88            ),
89            None => write!(
90                f,
91                "TokenBalance(token={}, amount={})",
92                self.token.symbol,
93                quantity.as_decimal()
94            ),
95        }
96    }
97}
98
99/// Represents the complete balance state of a blockchain wallet.
100///
101/// Tracks both the native currency balance (e.g., ETH, ARB) and ERC-20 token
102/// balances for a wallet address. The `token_universe` defines which tokens
103/// should be tracked for balance fetching.
104#[derive(Debug)]
105pub struct WalletBalance {
106    /// The balance of the chain's native currency
107    pub native_currency: Option<Money>,
108    /// Collection of ERC-20 token balances held in the wallet.
109    pub token_balances: Vec<TokenBalance>,
110    /// Set of token addresses to track for balance updates.
111    pub token_universe: HashSet<Address>,
112}
113
114impl WalletBalance {
115    /// Creates a new [`WalletBalance`] with the specified token universe.
116    #[must_use]
117    pub const fn new(token_universe: HashSet<Address>) -> Self {
118        Self {
119            native_currency: None,
120            token_balances: vec![],
121            token_universe,
122        }
123    }
124
125    /// Returns `true` if the token universe has been initialized with token addresses.
126    #[must_use]
127    pub fn is_token_universe_initialized(&self) -> bool {
128        !self.token_universe.is_empty()
129    }
130
131    /// Sets the native currency balance for the wallet.
132    pub fn set_native_currency_balance(&mut self, balance: Money) {
133        self.native_currency = Some(balance);
134    }
135
136    /// Adds an ERC-20 token balance to the wallet.
137    pub fn add_token_balance(&mut self, token_balance: TokenBalance) {
138        self.token_balances.push(token_balance);
139    }
140
141    /// Replaces the complete native and token balance snapshot.
142    ///
143    /// # Errors
144    ///
145    /// Returns an error if the token balances do not exactly match the configured universe,
146    /// contain duplicate currencies, or cannot be represented as account balances. The existing
147    /// snapshot remains unchanged on error.
148    pub fn replace_balances(
149        &mut self,
150        native_currency: Money,
151        mut token_balances: Vec<TokenBalance>,
152    ) -> anyhow::Result<Vec<AccountBalance>> {
153        token_balances.sort_unstable_by_key(|balance| balance.token.address);
154        let replacement = Self {
155            native_currency: Some(native_currency),
156            token_balances,
157            token_universe: self.token_universe.clone(),
158        };
159        let balances = replacement.account_balances(replacement.token_balances.iter())?;
160        *self = replacement;
161        Ok(balances)
162    }
163
164    /// Returns the complete wallet snapshot as account balances.
165    ///
166    /// # Errors
167    ///
168    /// Returns an error if the snapshot is incomplete, contains duplicate currencies, or a token
169    /// amount cannot be represented as money.
170    pub fn as_account_balances(&self) -> anyhow::Result<Vec<AccountBalance>> {
171        let mut token_balances = self.token_balances.iter().collect::<Vec<_>>();
172        token_balances.sort_unstable_by_key(|balance| balance.token.address);
173        self.account_balances(token_balances)
174    }
175
176    fn account_balances<'a>(
177        &'a self,
178        token_balances: impl IntoIterator<Item = &'a TokenBalance>,
179    ) -> anyhow::Result<Vec<AccountBalance>> {
180        self.validate_token_addresses()?;
181
182        let native_currency = self
183            .native_currency
184            .ok_or_else(|| anyhow::anyhow!("Wallet balance snapshot has no native currency"))?;
185        let mut currencies = HashSet::new();
186        currencies.insert(native_currency.currency);
187
188        let mut balances = Vec::with_capacity(self.token_balances.len() + 1);
189        balances.push(AccountBalance::new_checked(
190            native_currency,
191            Money::zero(native_currency.currency),
192            native_currency,
193        )?);
194
195        for token_balance in token_balances {
196            let total = token_balance.as_money()?;
197            if !currencies.insert(total.currency) {
198                anyhow::bail!(
199                    "Wallet balance snapshot contains duplicate currency {}",
200                    total.currency
201                );
202            }
203            balances.push(AccountBalance::new_checked(
204                total,
205                Money::zero(total.currency),
206                total,
207            )?);
208        }
209
210        Ok(balances)
211    }
212
213    fn validate_token_addresses(&self) -> anyhow::Result<()> {
214        let mut token_addresses = HashSet::with_capacity(self.token_balances.len());
215        let mut duplicates = Vec::new();
216
217        for balance in &self.token_balances {
218            if !token_addresses.insert(balance.token.address) {
219                duplicates.push(balance.token.address);
220            }
221        }
222
223        if !duplicates.is_empty() {
224            duplicates.sort_unstable();
225            duplicates.dedup();
226            anyhow::bail!(
227                "Wallet balance snapshot contains duplicate token addresses: {}",
228                format_addresses(&duplicates)
229            );
230        }
231
232        let mut missing = self
233            .token_universe
234            .difference(&token_addresses)
235            .copied()
236            .collect::<Vec<_>>();
237        let mut unexpected = token_addresses
238            .difference(&self.token_universe)
239            .copied()
240            .collect::<Vec<_>>();
241        missing.sort_unstable();
242        unexpected.sort_unstable();
243
244        match (missing.is_empty(), unexpected.is_empty()) {
245            (false, false) => anyhow::bail!(
246                "Wallet balance snapshot is missing configured token addresses: {}; contains unexpected token addresses: {}",
247                format_addresses(&missing),
248                format_addresses(&unexpected)
249            ),
250            (false, true) => anyhow::bail!(
251                "Wallet balance snapshot is missing configured token addresses: {}",
252                format_addresses(&missing)
253            ),
254            (true, false) => anyhow::bail!(
255                "Wallet balance snapshot contains unexpected token addresses: {}",
256                format_addresses(&unexpected)
257            ),
258            (true, true) => Ok(()),
259        }
260    }
261}
262
263fn format_addresses(addresses: &[Address]) -> String {
264    addresses
265        .iter()
266        .map(ToString::to_string)
267        .collect::<Vec<_>>()
268        .join(", ")
269}
270
271#[cfg(test)]
272mod tests {
273    use std::sync::Arc;
274
275    use alloy_primitives::{U256, address};
276    use rstest::rstest;
277
278    use super::*;
279    use crate::defi::{
280        SharedChain, Token,
281        chain::chains,
282        stubs::{arbitrum, usdc, weth},
283    };
284
285    fn create_token(symbol: &str, decimals: u8) -> Token {
286        Token::new(
287            Arc::new(chains::ETHEREUM.clone()),
288            address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"),
289            format!("{symbol} Token"),
290            symbol.to_string(),
291            decimals,
292        )
293    }
294
295    #[rstest]
296    fn test_token_balance_as_quantity_18_decimals(#[from(arbitrum)] chain: SharedChain) {
297        // Test case: NU token with 18 decimals
298        // Raw amount: 10342000000000000000000 (10342 * 10^18)
299        // Expected: 10342.000000000000000000
300        let token = Token::new(
301            chain,
302            address!("0x4fE83213D56308330EC302a8BD641f1d0113A4Cc"),
303            "NuCypher".to_string(),
304            "NU".to_string(),
305            18,
306        );
307        let amount = U256::from(10342u64) * U256::from(10u64).pow(U256::from(18u64));
308        let balance = TokenBalance::new(amount, token);
309
310        let quantity = balance.as_quantity().unwrap();
311        assert_eq!(
312            quantity.as_decimal().to_string(),
313            "10342.000000000000000000"
314        );
315    }
316
317    #[rstest]
318    fn test_token_balance_as_quantity_6_decimals() {
319        // Test case: USDC with 6 decimals
320        // Raw amount: 92220728254 (92220.728254 * 10^6)
321        // Expected: 92220.728254
322        let token = create_token("USDC", 6);
323        let amount = U256::from(92_220_728_254_u64);
324        let balance = TokenBalance::new(amount, token);
325
326        let quantity = balance.as_quantity().unwrap();
327        assert_eq!(quantity.as_decimal().to_string(), "92220.728254");
328    }
329
330    #[rstest]
331    fn test_token_balance_as_quantity_fractional_18_decimals(#[from(arbitrum)] chain: SharedChain) {
332        // Test case: mETH with 18 decimals and fractional amount
333        // Raw amount: 758325512078001391
334        // Expected: 0.758325512078001391
335        let token = Token::new(
336            chain,
337            address!("0xd5F7838F5C461fefF7FE49ea5ebaF7728bB0ADfa"),
338            "mETH".to_string(),
339            "mETH".to_string(),
340            18,
341        );
342        let amount = U256::from(758_325_512_078_001_391_u64);
343        let balance = TokenBalance::new(amount, token);
344
345        let quantity = balance.as_quantity().unwrap();
346        assert_eq!(quantity.as_decimal().to_string(), "0.758325512078001391");
347    }
348
349    #[rstest]
350    fn test_token_balance_display_18_decimals(#[from(arbitrum)] chain: SharedChain) {
351        // Test Display implementation with 18 decimal token
352        let token = Token::new(
353            chain,
354            address!("0x912CE59144191C1204E64559FE8253a0e49E6548"),
355            "Arbitrum".to_string(),
356            "ARB".to_string(),
357            18,
358        );
359        // 7922.013795343949480329 ARB
360        let amount = U256::from_str_radix("7922013795343949480329", 10).unwrap();
361        let balance = TokenBalance::new(amount, token);
362
363        let display = balance.to_string();
364        assert!(display.contains("ARB"));
365        assert!(display.contains("7922.013795343949480329"));
366    }
367
368    #[rstest]
369    fn test_token_balance_display_6_decimals() {
370        // Test Display implementation with 6 decimal token (USDC)
371        let token = create_token("USDC", 6);
372        let amount = U256::from(92_220_728_254_u64); // 92220.728254 USDC
373        let balance = TokenBalance::new(amount, token);
374
375        let display = balance.to_string();
376        assert!(display.contains("USDC"));
377        assert!(display.contains("92220.728254"));
378    }
379
380    #[rstest]
381    fn test_token_balance_set_amount_usd(weth: Token) {
382        let amount = U256::from(1u64) * U256::from(10u64).pow(U256::from(18u64));
383        let mut balance = TokenBalance::new(amount, weth);
384
385        assert!(balance.amount_usd.is_none());
386
387        let usd_value = Quantity::from("3500.00");
388        balance.set_amount_usd(usd_value);
389
390        assert!(balance.amount_usd.is_some());
391        assert_eq!(
392            balance.amount_usd.unwrap().as_decimal().to_string(),
393            "3500.00"
394        );
395    }
396
397    #[rstest]
398    fn test_wallet_balance_new_empty() {
399        let wallet = WalletBalance::new(HashSet::new());
400
401        assert!(wallet.native_currency.is_none());
402        assert!(wallet.token_balances.is_empty());
403        assert!(!wallet.is_token_universe_initialized());
404    }
405
406    #[rstest]
407    fn test_wallet_balance_with_token_universe() {
408        let mut tokens = HashSet::new();
409        tokens.insert(address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")); // USDC
410        tokens.insert(address!("0x912CE59144191C1204E64559FE8253a0e49E6548")); // ARB
411
412        let wallet = WalletBalance::new(tokens);
413
414        assert!(wallet.is_token_universe_initialized());
415        assert_eq!(wallet.token_universe.len(), 2);
416    }
417
418    #[rstest]
419    fn test_wallet_balance_set_native_currency() {
420        let mut wallet = WalletBalance::new(HashSet::new());
421
422        assert!(wallet.native_currency.is_none());
423
424        let eth_balance = Money::new(50.936_054, crate::types::Currency::ETH());
425        wallet.set_native_currency_balance(eth_balance);
426
427        assert!(wallet.native_currency.is_some());
428    }
429
430    #[rstest]
431    fn test_wallet_balance_add_token_balance(usdc: Token, weth: Token) {
432        let mut wallet = WalletBalance::new(HashSet::new());
433
434        let usdc_balance = TokenBalance::new(U256::from(100_000_000u64), usdc); // 100 USDC
435        let weth_balance = TokenBalance::new(U256::from(10u64).pow(U256::from(18u64)), weth); // 1 WETH
436
437        wallet.add_token_balance(usdc_balance);
438        wallet.add_token_balance(weth_balance);
439
440        assert_eq!(wallet.token_balances.len(), 2);
441        assert_eq!(wallet.token_balances[0].token.symbol, "USDC");
442        assert_eq!(wallet.token_balances[1].token.symbol, "WETH");
443    }
444
445    #[rstest]
446    fn test_replace_balances_retains_snapshot_on_conversion_failure(weth: Token) {
447        let mut wallet = WalletBalance::new(HashSet::from([weth.address]));
448        let native_currency = weth.chain.native_currency();
449        let native = Money::from_wei(U256::from(1_000_000_000_000_000_000_u64), native_currency);
450        let token = TokenBalance::new(U256::from(2_000_000_000_000_000_000_u64), weth.clone());
451        wallet.replace_balances(native, vec![token]).unwrap();
452
453        let error = wallet
454            .replace_balances(
455                Money::from_wei(U256::from(3_000_000_000_000_000_000_u64), native_currency),
456                vec![TokenBalance::new(U256::MAX, weth)],
457            )
458            .unwrap_err();
459
460        assert!(
461            error.to_string().contains("exceeds QuantityRaw range"),
462            "was: {error}"
463        );
464        assert_eq!(wallet.native_currency.unwrap(), native);
465        assert_eq!(wallet.token_balances.len(), 1);
466        assert_eq!(
467            wallet.token_balances[0].amount,
468            U256::from(2_000_000_000_000_000_000_u64)
469        );
470    }
471
472    #[rstest]
473    fn test_account_balances_reports_missing_token_address(usdc: Token, weth: Token) {
474        let missing = weth.address;
475        let wallet = WalletBalance {
476            native_currency: None,
477            token_balances: vec![TokenBalance::new(U256::from(1_u64), usdc.clone())],
478            token_universe: HashSet::from([usdc.address, missing]),
479        };
480
481        let error = wallet.as_account_balances().unwrap_err();
482
483        assert_eq!(
484            error.to_string(),
485            format!("Wallet balance snapshot is missing configured token addresses: {missing}")
486        );
487    }
488
489    #[rstest]
490    fn test_account_balances_reports_unexpected_token_address(usdc: Token, weth: Token) {
491        let unexpected = weth.address;
492        let wallet = WalletBalance {
493            native_currency: None,
494            token_balances: vec![
495                TokenBalance::new(U256::from(1_u64), usdc.clone()),
496                TokenBalance::new(U256::from(2_u64), weth),
497            ],
498            token_universe: HashSet::from([usdc.address]),
499        };
500
501        let error = wallet.as_account_balances().unwrap_err();
502
503        assert_eq!(
504            error.to_string(),
505            format!("Wallet balance snapshot contains unexpected token addresses: {unexpected}")
506        );
507    }
508
509    #[rstest]
510    fn test_account_balances_reports_missing_and_unexpected_token_addresses(
511        usdc: Token,
512        weth: Token,
513    ) {
514        let missing = usdc.address;
515        let unexpected = weth.address;
516        let wallet = WalletBalance {
517            native_currency: None,
518            token_balances: vec![TokenBalance::new(U256::from(1_u64), weth)],
519            token_universe: HashSet::from([missing]),
520        };
521
522        let error = wallet.as_account_balances().unwrap_err();
523
524        assert_eq!(
525            error.to_string(),
526            format!(
527                "Wallet balance snapshot is missing configured token addresses: {missing}; contains unexpected token addresses: {unexpected}"
528            )
529        );
530    }
531
532    #[rstest]
533    fn test_account_balances_reports_duplicate_before_set_differences(usdc: Token, weth: Token) {
534        let duplicate = usdc.address;
535        let wallet = WalletBalance {
536            native_currency: None,
537            token_balances: vec![
538                TokenBalance::new(U256::from(1_u64), usdc.clone()),
539                TokenBalance::new(U256::from(2_u64), usdc),
540            ],
541            token_universe: HashSet::from([weth.address]),
542        };
543
544        let error = wallet.as_account_balances().unwrap_err();
545
546        assert_eq!(
547            error.to_string(),
548            format!("Wallet balance snapshot contains duplicate token addresses: {duplicate}")
549        );
550    }
551
552    #[rstest]
553    fn test_replace_balances_rejects_duplicate_currency_symbols(
554        usdc: Token,
555        #[from(arbitrum)] chain: SharedChain,
556    ) {
557        let duplicate = Token::new(
558            chain,
559            address!("0x0000000000000000000000000000000000000001"),
560            "Other USD Coin".to_string(),
561            "USDC".to_string(),
562            18,
563        );
564        let mut wallet = WalletBalance::new(HashSet::from([usdc.address, duplicate.address]));
565
566        let error = wallet
567            .replace_balances(
568                Money::from_wei(
569                    U256::from(1_000_000_000_000_000_000_u64),
570                    usdc.chain.native_currency(),
571                ),
572                vec![
573                    TokenBalance::new(U256::from(1_u64), usdc),
574                    TokenBalance::new(U256::from(2_u64), duplicate),
575                ],
576            )
577            .unwrap_err();
578
579        assert_eq!(
580            error.to_string(),
581            "Wallet balance snapshot contains duplicate currency USDC"
582        );
583        assert!(wallet.native_currency.is_none());
584        assert!(wallet.token_balances.is_empty());
585    }
586}
nautilus_model/defi/wallet.rs

nautilus_model/defi/
wallet.rs
