wallet.rs
1 /// Wallet management for tracking balances across Alpha and Delta chains 2 /// 3 /// Supports: 4 /// - AX (Alpha native token) 5 /// - sAX (Shielded AX on Delta) 6 /// - DX (Delta native token) 7 8 use crate::{BotError, Result}; 9 use serde::{Deserialize, Serialize}; 10 use std::collections::HashMap; 11 12 /// Chain identifier 13 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] 14 pub enum ChainId { 15 Alpha, 16 Delta, 17 } 18 19 impl std::fmt::Display for ChainId { 20 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 21 match self { 22 ChainId::Alpha => write!(f, "alpha"), 23 ChainId::Delta => write!(f, "delta"), 24 } 25 } 26 } 27 28 /// Token identifier 29 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] 30 pub enum Token { 31 /// Alpha native token 32 AX, 33 /// Shielded AX on Delta 34 SAX, 35 /// Delta native token 36 DX, 37 } 38 39 impl std::fmt::Display for Token { 40 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 41 match self { 42 Token::AX => write!(f, "AX"), 43 Token::SAX => write!(f, "sAX"), 44 Token::DX => write!(f, "DX"), 45 } 46 } 47 } 48 49 /// Balance amount (using u128 for large values) 50 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] 51 pub struct Balance { 52 amount: u128, 53 } 54 55 impl Balance { 56 pub fn new(amount: u128) -> Self { 57 Self { amount } 58 } 59 60 pub fn zero() -> Self { 61 Self { amount: 0 } 62 } 63 64 pub fn amount(&self) -> u128 { 65 self.amount 66 } 67 68 pub fn add(&self, other: Balance) -> Result<Balance> { 69 self.amount 70 .checked_add(other.amount) 71 .map(Balance::new) 72 .ok_or_else(|| BotError::WalletError("Balance overflow".to_string())) 73 } 74 75 pub fn sub(&self, other: Balance) -> Result<Balance> { 76 self.amount 77 .checked_sub(other.amount) 78 .map(Balance::new) 79 .ok_or_else(|| BotError::WalletError("Insufficient balance".to_string())) 80 } 81 82 pub fn is_zero(&self) -> bool { 83 self.amount == 0 84 } 85 } 86 87 impl std::fmt::Display for Balance { 88 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 89 write!(f, "{}", self.amount) 90 } 91 } 92 93 /// Multi-chain wallet for bot accounts 94 #[derive(Debug, Clone, Serialize, Deserialize)] 95 pub struct Wallet { 96 /// Owner's bot ID 97 pub owner_id: String, 98 99 /// Balances per token 100 balances: HashMap<Token, Balance>, 101 102 /// Pending operations (transaction hashes) 103 pending_ops: Vec<String>, 104 } 105 106 impl Wallet { 107 /// Create a new wallet 108 pub fn new(owner_id: String) -> Self { 109 let mut balances = HashMap::new(); 110 balances.insert(Token::AX, Balance::zero()); 111 balances.insert(Token::SAX, Balance::zero()); 112 balances.insert(Token::DX, Balance::zero()); 113 114 Self { 115 owner_id, 116 balances, 117 pending_ops: Vec::new(), 118 } 119 } 120 121 /// Create a wallet with initial balances 122 pub fn with_balances(owner_id: String, initial: HashMap<Token, Balance>) -> Self { 123 let mut wallet = Self::new(owner_id); 124 for (token, balance) in initial { 125 wallet.balances.insert(token, balance); 126 } 127 wallet 128 } 129 130 /// Get balance for a token 131 pub fn balance(&self, token: &Token) -> Balance { 132 self.balances.get(token).copied().unwrap_or_else(Balance::zero) 133 } 134 135 /// Credit (add) to balance 136 pub fn credit(&mut self, token: Token, amount: Balance) -> Result<()> { 137 let current = self.balance(&token); 138 let new_balance = current.add(amount)?; 139 self.balances.insert(token, new_balance); 140 Ok(()) 141 } 142 143 /// Debit (subtract) from balance 144 pub fn debit(&mut self, token: Token, amount: Balance) -> Result<()> { 145 let current = self.balance(&token); 146 let new_balance = current.sub(amount)?; 147 self.balances.insert(token, new_balance); 148 Ok(()) 149 } 150 151 /// Check if wallet has sufficient balance 152 pub fn has_balance(&self, token: &Token, amount: Balance) -> bool { 153 self.balance(token).amount() >= amount.amount() 154 } 155 156 /// Add a pending operation 157 pub fn add_pending_op(&mut self, tx_hash: String) { 158 self.pending_ops.push(tx_hash); 159 } 160 161 /// Clear pending operations 162 pub fn clear_pending_ops(&mut self) { 163 self.pending_ops.clear(); 164 } 165 166 /// Get all pending operations 167 pub fn pending_ops(&self) -> &[String] { 168 &self.pending_ops 169 } 170 171 /// Get total value across all tokens (simplified) 172 /// Returns value in smallest units 173 pub fn total_value(&self) -> u128 { 174 self.balances.values().map(|b| b.amount()).sum() 175 } 176 177 /// Snapshot of all balances 178 pub fn snapshot(&self) -> HashMap<Token, Balance> { 179 self.balances.clone() 180 } 181 } 182 183 #[cfg(test)] 184 mod tests { 185 use super::*; 186 187 #[test] 188 fn test_wallet_creation() { 189 let wallet = Wallet::new("test-bot".to_string()); 190 assert_eq!(wallet.balance(&Token::AX), Balance::zero()); 191 assert_eq!(wallet.balance(&Token::SAX), Balance::zero()); 192 assert_eq!(wallet.balance(&Token::DX), Balance::zero()); 193 } 194 195 #[test] 196 fn test_credit_debit() { 197 let mut wallet = Wallet::new("test-bot".to_string()); 198 199 // Credit 1000 AX 200 wallet.credit(Token::AX, Balance::new(1000)).unwrap(); 201 assert_eq!(wallet.balance(&Token::AX).amount(), 1000); 202 203 // Debit 300 AX 204 wallet.debit(Token::AX, Balance::new(300)).unwrap(); 205 assert_eq!(wallet.balance(&Token::AX).amount(), 700); 206 } 207 208 #[test] 209 fn test_insufficient_balance() { 210 let mut wallet = Wallet::new("test-bot".to_string()); 211 wallet.credit(Token::AX, Balance::new(100)).unwrap(); 212 213 // Try to debit more than available 214 let result = wallet.debit(Token::AX, Balance::new(200)); 215 assert!(result.is_err()); 216 } 217 218 #[test] 219 fn test_has_balance() { 220 let mut wallet = Wallet::new("test-bot".to_string()); 221 wallet.credit(Token::AX, Balance::new(1000)).unwrap(); 222 223 assert!(wallet.has_balance(&Token::AX, Balance::new(500))); 224 assert!(wallet.has_balance(&Token::AX, Balance::new(1000))); 225 assert!(!wallet.has_balance(&Token::AX, Balance::new(1001))); 226 } 227 228 #[test] 229 fn test_pending_ops() { 230 let mut wallet = Wallet::new("test-bot".to_string()); 231 232 wallet.add_pending_op("tx_hash_1".to_string()); 233 wallet.add_pending_op("tx_hash_2".to_string()); 234 235 assert_eq!(wallet.pending_ops().len(), 2); 236 237 wallet.clear_pending_ops(); 238 assert_eq!(wallet.pending_ops().len(), 0); 239 } 240 241 #[test] 242 fn test_balance_operations() { 243 let b1 = Balance::new(100); 244 let b2 = Balance::new(50); 245 246 let sum = b1.add(b2).unwrap(); 247 assert_eq!(sum.amount(), 150); 248 249 let diff = b1.sub(b2).unwrap(); 250 assert_eq!(diff.amount(), 50); 251 } 252 }