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    // Helper to create a token with specific decimals
286    fn create_token(symbol: &str, decimals: u8) -> Token {
287        Token::new(
288            Arc::new(chains::ETHEREUM.clone()),
289            address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"),
290            format!("{symbol} Token"),
291            symbol.to_string(),
292            decimals,
293        )
294    }
295
296    #[rstest]
297    fn test_token_balance_as_quantity_18_decimals(#[from(arbitrum)] chain: SharedChain) {
298        // Test case: NU token with 18 decimals
299        // Raw amount: 10342000000000000000000 (10342 * 10^18)
300        // Expected: 10342.000000000000000000
301        let token = Token::new(
302            chain,
303            address!("0x4fE83213D56308330EC302a8BD641f1d0113A4Cc"),
304            "NuCypher".to_string(),
305            "NU".to_string(),
306            18,
307        );
308        let amount = U256::from(10342u64) * U256::from(10u64).pow(U256::from(18u64));
309        let balance = TokenBalance::new(amount, token);
310
311        let quantity = balance.as_quantity().unwrap();
312        assert_eq!(
313            quantity.as_decimal().to_string(),
314            "10342.000000000000000000"
315        );
316    }
317
318    #[rstest]
319    fn test_token_balance_as_quantity_6_decimals() {
320        // Test case: USDC with 6 decimals
321        // Raw amount: 92220728254 (92220.728254 * 10^6)
322        // Expected: 92220.728254
323        let token = create_token("USDC", 6);
324        let amount = U256::from(92_220_728_254_u64);
325        let balance = TokenBalance::new(amount, token);
326
327        let quantity = balance.as_quantity().unwrap();
328        assert_eq!(quantity.as_decimal().to_string(), "92220.728254");
329    }
330
331    #[rstest]
332    fn test_token_balance_as_quantity_fractional_18_decimals(#[from(arbitrum)] chain: SharedChain) {
333        // Test case: mETH with 18 decimals and fractional amount
334        // Raw amount: 758325512078001391
335        // Expected: 0.758325512078001391
336        let token = Token::new(
337            chain,
338            address!("0xd5F7838F5C461fefF7FE49ea5ebaF7728bB0ADfa"),
339            "mETH".to_string(),
340            "mETH".to_string(),
341            18,
342        );
343        let amount = U256::from(758_325_512_078_001_391_u64);
344        let balance = TokenBalance::new(amount, token);
345
346        let quantity = balance.as_quantity().unwrap();
347        assert_eq!(quantity.as_decimal().to_string(), "0.758325512078001391");
348    }
349
350    #[rstest]
351    fn test_token_balance_display_18_decimals(#[from(arbitrum)] chain: SharedChain) {
352        // Test Display implementation with 18 decimal token
353        let token = Token::new(
354            chain,
355            address!("0x912CE59144191C1204E64559FE8253a0e49E6548"),
356            "Arbitrum".to_string(),
357            "ARB".to_string(),
358            18,
359        );
360        // 7922.013795343949480329 ARB
361        let amount = U256::from_str_radix("7922013795343949480329", 10).unwrap();
362        let balance = TokenBalance::new(amount, token);
363
364        let display = balance.to_string();
365        assert!(display.contains("ARB"));
366        assert!(display.contains("7922.013795343949480329"));
367    }
368
369    #[rstest]
370    fn test_token_balance_display_6_decimals() {
371        // Test Display implementation with 6 decimal token (USDC)
372        let token = create_token("USDC", 6);
373        let amount = U256::from(92_220_728_254_u64); // 92220.728254 USDC
374        let balance = TokenBalance::new(amount, token);
375
376        let display = balance.to_string();
377        assert!(display.contains("USDC"));
378        assert!(display.contains("92220.728254"));
379    }
380
381    #[rstest]
382    fn test_token_balance_set_amount_usd(weth: Token) {
383        let amount = U256::from(1u64) * U256::from(10u64).pow(U256::from(18u64));
384        let mut balance = TokenBalance::new(amount, weth);
385
386        assert!(balance.amount_usd.is_none());
387
388        let usd_value = Quantity::from("3500.00");
389        balance.set_amount_usd(usd_value);
390
391        assert!(balance.amount_usd.is_some());
392        assert_eq!(
393            balance.amount_usd.unwrap().as_decimal().to_string(),
394            "3500.00"
395        );
396    }
397
398    #[rstest]
399    fn test_wallet_balance_new_empty() {
400        let wallet = WalletBalance::new(HashSet::new());
401
402        assert!(wallet.native_currency.is_none());
403        assert!(wallet.token_balances.is_empty());
404        assert!(!wallet.is_token_universe_initialized());
405    }
406
407    #[rstest]
408    fn test_wallet_balance_with_token_universe() {
409        let mut tokens = HashSet::new();
410        tokens.insert(address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")); // USDC
411        tokens.insert(address!("0x912CE59144191C1204E64559FE8253a0e49E6548")); // ARB
412
413        let wallet = WalletBalance::new(tokens);
414
415        assert!(wallet.is_token_universe_initialized());
416        assert_eq!(wallet.token_universe.len(), 2);
417    }
418
419    #[rstest]
420    fn test_wallet_balance_set_native_currency() {
421        let mut wallet = WalletBalance::new(HashSet::new());
422
423        assert!(wallet.native_currency.is_none());
424
425        let eth_balance = Money::new(50.936_054, crate::types::Currency::ETH());
426        wallet.set_native_currency_balance(eth_balance);
427
428        assert!(wallet.native_currency.is_some());
429    }
430
431    #[rstest]
432    fn test_wallet_balance_add_token_balance(usdc: Token, weth: Token) {
433        let mut wallet = WalletBalance::new(HashSet::new());
434
435        let usdc_balance = TokenBalance::new(U256::from(100_000_000u64), usdc); // 100 USDC
436        let weth_balance = TokenBalance::new(U256::from(10u64).pow(U256::from(18u64)), weth); // 1 WETH
437
438        wallet.add_token_balance(usdc_balance);
439        wallet.add_token_balance(weth_balance);
440
441        assert_eq!(wallet.token_balances.len(), 2);
442        assert_eq!(wallet.token_balances[0].token.symbol, "USDC");
443        assert_eq!(wallet.token_balances[1].token.symbol, "WETH");
444    }
445
446    #[rstest]
447    fn test_replace_balances_retains_snapshot_on_conversion_failure(weth: Token) {
448        let mut wallet = WalletBalance::new(HashSet::from([weth.address]));
449        let native_currency = weth.chain.native_currency();
450        let native = Money::from_wei(U256::from(1_000_000_000_000_000_000_u64), native_currency);
451        let token = TokenBalance::new(U256::from(2_000_000_000_000_000_000_u64), weth.clone());
452        wallet.replace_balances(native, vec![token]).unwrap();
453
454        let error = wallet
455            .replace_balances(
456                Money::from_wei(U256::from(3_000_000_000_000_000_000_u64), native_currency),
457                vec![TokenBalance::new(U256::MAX, weth)],
458            )
459            .unwrap_err();
460
461        assert!(
462            error.to_string().contains("exceeds QuantityRaw range"),
463            "was: {error}"
464        );
465        assert_eq!(wallet.native_currency.unwrap(), native);
466        assert_eq!(wallet.token_balances.len(), 1);
467        assert_eq!(
468            wallet.token_balances[0].amount,
469            U256::from(2_000_000_000_000_000_000_u64)
470        );
471    }
472
473    #[rstest]
474    fn test_account_balances_reports_missing_token_address(usdc: Token, weth: Token) {
475        let missing = weth.address;
476        let wallet = WalletBalance {
477            native_currency: None,
478            token_balances: vec![TokenBalance::new(U256::from(1_u64), usdc.clone())],
479            token_universe: HashSet::from([usdc.address, missing]),
480        };
481
482        let error = wallet.as_account_balances().unwrap_err();
483
484        assert_eq!(
485            error.to_string(),
486            format!("Wallet balance snapshot is missing configured token addresses: {missing}")
487        );
488    }
489
490    #[rstest]
491    fn test_account_balances_reports_unexpected_token_address(usdc: Token, weth: Token) {
492        let unexpected = weth.address;
493        let wallet = WalletBalance {
494            native_currency: None,
495            token_balances: vec![
496                TokenBalance::new(U256::from(1_u64), usdc.clone()),
497                TokenBalance::new(U256::from(2_u64), weth),
498            ],
499            token_universe: HashSet::from([usdc.address]),
500        };
501
502        let error = wallet.as_account_balances().unwrap_err();
503
504        assert_eq!(
505            error.to_string(),
506            format!("Wallet balance snapshot contains unexpected token addresses: {unexpected}")
507        );
508    }
509
510    #[rstest]
511    fn test_account_balances_reports_missing_and_unexpected_token_addresses(
512        usdc: Token,
513        weth: Token,
514    ) {
515        let missing = usdc.address;
516        let unexpected = weth.address;
517        let wallet = WalletBalance {
518            native_currency: None,
519            token_balances: vec![TokenBalance::new(U256::from(1_u64), weth)],
520            token_universe: HashSet::from([missing]),
521        };
522
523        let error = wallet.as_account_balances().unwrap_err();
524
525        assert_eq!(
526            error.to_string(),
527            format!(
528                "Wallet balance snapshot is missing configured token addresses: {missing}; contains unexpected token addresses: {unexpected}"
529            )
530        );
531    }
532
533    #[rstest]
534    fn test_account_balances_reports_duplicate_before_set_differences(usdc: Token, weth: Token) {
535        let duplicate = usdc.address;
536        let wallet = WalletBalance {
537            native_currency: None,
538            token_balances: vec![
539                TokenBalance::new(U256::from(1_u64), usdc.clone()),
540                TokenBalance::new(U256::from(2_u64), usdc),
541            ],
542            token_universe: HashSet::from([weth.address]),
543        };
544
545        let error = wallet.as_account_balances().unwrap_err();
546
547        assert_eq!(
548            error.to_string(),
549            format!("Wallet balance snapshot contains duplicate token addresses: {duplicate}")
550        );
551    }
552
553    #[rstest]
554    fn test_replace_balances_rejects_duplicate_currency_symbols(
555        usdc: Token,
556        #[from(arbitrum)] chain: SharedChain,
557    ) {
558        let duplicate = Token::new(
559            chain,
560            address!("0x0000000000000000000000000000000000000001"),
561            "Other USD Coin".to_string(),
562            "USDC".to_string(),
563            18,
564        );
565        let mut wallet = WalletBalance::new(HashSet::from([usdc.address, duplicate.address]));
566
567        let error = wallet
568            .replace_balances(
569                Money::from_wei(
570                    U256::from(1_000_000_000_000_000_000_u64),
571                    usdc.chain.native_currency(),
572                ),
573                vec![
574                    TokenBalance::new(U256::from(1_u64), usdc),
575                    TokenBalance::new(U256::from(2_u64), duplicate),
576                ],
577            )
578            .unwrap_err();
579
580        assert_eq!(
581            error.to_string(),
582            "Wallet balance snapshot contains duplicate currency USDC"
583        );
584        assert!(wallet.native_currency.is_none());
585        assert!(wallet.token_balances.is_empty());
586    }
587}
nautilus_model/defi/wallet.rs

nautilus_model/defi/
wallet.rs
