Skip to main content

scope/
tokens.rs

1//! # Token Alias Storage
2//!
3//! This module provides storage and retrieval of token aliases,
4//! allowing users to reference tokens by friendly names instead
5//! of full contract addresses.
6//!
7//! ## Storage Location
8//!
9//! Token aliases are stored in `~/.local/share/scope/tokens.yaml`
10//!
11//! ## Usage
12//!
13//! ```rust,no_run
14//! use scope::tokens::TokenAliases;
15//!
16//! // Load existing aliases
17//! let mut aliases = TokenAliases::load();
18//!
19//! // Add an alias
20//! aliases.add("USDC", "ethereum", "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48", "USD Coin");
21//!
22//! // Look up an alias
23//! if let Some(info) = aliases.get("USDC", Some("ethereum")) {
24//!     println!("USDC address: {}", info.address);
25//! }
26//!
27//! // Save aliases
28//! aliases.save().unwrap();
29//! ```
30
31use crate::error::{Result, ScopeError};
32use serde::{Deserialize, Serialize};
33use std::collections::HashMap;
34use std::path::PathBuf;
35
36/// A saved token alias with its address and metadata.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct TokenInfo {
39    /// Token contract address.
40    pub address: String,
41
42    /// Token symbol.
43    pub symbol: String,
44
45    /// Token name.
46    pub name: String,
47
48    /// Blockchain network.
49    pub chain: String,
50
51    /// When this alias was last used.
52    #[serde(default)]
53    pub last_used: Option<i64>,
54}
55
56/// Collection of token aliases.
57#[derive(Debug, Clone, Default, Serialize, Deserialize)]
58pub struct TokenAliases {
59    /// Map of alias -> chain -> token info.
60    /// Using nested maps to support same symbol on different chains.
61    #[serde(default)]
62    aliases: HashMap<String, HashMap<String, TokenInfo>>,
63
64    /// Recent tokens for quick access.
65    #[serde(default)]
66    recent: Vec<TokenInfo>,
67}
68
69impl TokenAliases {
70    /// Returns the path to the token aliases file.
71    pub fn aliases_path() -> Option<PathBuf> {
72        dirs::data_dir().map(|p| p.join("scope").join("tokens.yaml"))
73    }
74
75    /// Loads token aliases from disk.
76    pub fn load() -> Self {
77        Self::aliases_path()
78            .and_then(|path| std::fs::read_to_string(&path).ok())
79            .and_then(|contents| serde_yaml::from_str(&contents).ok())
80            .unwrap_or_default()
81    }
82
83    /// Saves token aliases to disk.
84    pub fn save(&self) -> Result<()> {
85        if let Some(path) = Self::aliases_path() {
86            if let Some(parent) = path.parent() {
87                std::fs::create_dir_all(parent).map_err(|e| {
88                    ScopeError::Io(format!("Failed to create token aliases directory: {}", e))
89                })?;
90            }
91            let contents = serde_yaml::to_string(self)
92                .map_err(|e| ScopeError::Export(format!("Failed to serialize aliases: {}", e)))?;
93            std::fs::write(&path, contents)
94                .map_err(|e| ScopeError::Io(format!("Failed to write token aliases: {}", e)))?;
95        }
96        Ok(())
97    }
98
99    /// Adds or updates a token alias.
100    ///
101    /// # Arguments
102    ///
103    /// * `alias` - The friendly name/symbol to use (case-insensitive)
104    /// * `chain` - The blockchain network
105    /// * `address` - The token contract address
106    /// * `name` - The full token name
107    pub fn add(&mut self, alias: &str, chain: &str, address: &str, name: &str) {
108        let alias_key = alias.to_uppercase();
109        let chain_key = chain.to_lowercase();
110
111        let info = TokenInfo {
112            address: address.to_string(),
113            symbol: alias.to_uppercase(),
114            name: name.to_string(),
115            chain: chain_key.clone(),
116            last_used: Some(chrono::Utc::now().timestamp()),
117        };
118
119        // Add to aliases map
120        self.aliases
121            .entry(alias_key)
122            .or_default()
123            .insert(chain_key, info.clone());
124
125        // Add to recent (remove existing first, then add to front)
126        self.recent
127            .retain(|t| !(t.symbol == info.symbol && t.chain == info.chain));
128        self.recent.insert(0, info);
129
130        // Keep only last 20 recent
131        self.recent.truncate(20);
132    }
133
134    /// Looks up a token alias.
135    ///
136    /// # Arguments
137    ///
138    /// * `alias` - The alias to look up (case-insensitive)
139    /// * `chain` - Optional chain filter. If None, returns first match.
140    ///
141    /// # Returns
142    ///
143    /// Returns the token info if found.
144    pub fn get(&self, alias: &str, chain: Option<&str>) -> Option<&TokenInfo> {
145        let alias_key = alias.to_uppercase();
146
147        if let Some(chain_map) = self.aliases.get(&alias_key) {
148            if let Some(chain) = chain {
149                let chain_key = chain.to_lowercase();
150                chain_map.get(&chain_key)
151            } else {
152                // Return the first one (or prefer ethereum if available)
153                chain_map
154                    .get("ethereum")
155                    .or_else(|| chain_map.values().next())
156            }
157        } else {
158            None
159        }
160    }
161
162    /// Gets all chains that have this alias defined.
163    pub fn get_chains_for_alias(&self, alias: &str) -> Vec<&str> {
164        let alias_key = alias.to_uppercase();
165        self.aliases
166            .get(&alias_key)
167            .map(|chain_map| chain_map.keys().map(|s| s.as_str()).collect())
168            .unwrap_or_default()
169    }
170
171    /// Returns recent tokens.
172    pub fn recent(&self) -> &[TokenInfo] {
173        &self.recent
174    }
175
176    /// Removes an alias.
177    pub fn remove(&mut self, alias: &str, chain: Option<&str>) {
178        let alias_key = alias.to_uppercase();
179
180        if let Some(chain) = chain {
181            let chain_key = chain.to_lowercase();
182            if let Some(chain_map) = self.aliases.get_mut(&alias_key) {
183                chain_map.remove(&chain_key);
184                if chain_map.is_empty() {
185                    self.aliases.remove(&alias_key);
186                }
187            }
188            self.recent
189                .retain(|t| !(t.symbol == alias_key && t.chain == chain_key));
190        } else {
191            self.aliases.remove(&alias_key);
192            self.recent.retain(|t| t.symbol != alias_key);
193        }
194    }
195
196    /// Lists all saved aliases.
197    pub fn list(&self) -> Vec<&TokenInfo> {
198        self.aliases
199            .values()
200            .flat_map(|chain_map| chain_map.values())
201            .collect()
202    }
203
204    /// Checks if the input looks like a token address or a name.
205    pub fn is_address(input: &str) -> bool {
206        // EVM address: 0x + 40 hex chars
207        if input.starts_with("0x") && input.len() == 42 {
208            return input[2..].chars().all(|c| c.is_ascii_hexdigit());
209        }
210
211        // Solana address: base58, 32-44 chars
212        if input.len() >= 32
213            && input.len() <= 44
214            && let Ok(decoded) = bs58::decode(input).into_vec()
215            && decoded.len() == 32
216        {
217            return true;
218        }
219
220        // Tron address: T + 33 chars
221        if input.starts_with('T') && input.len() == 34 {
222            return bs58::decode(input).into_vec().is_ok();
223        }
224
225        false
226    }
227}
228
229// ============================================================================
230// Unit Tests
231// ============================================================================
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn test_token_aliases_default() {
239        let aliases = TokenAliases::default();
240        assert!(aliases.aliases.is_empty());
241        assert!(aliases.recent.is_empty());
242    }
243
244    #[test]
245    fn test_add_and_get_alias() {
246        let mut aliases = TokenAliases::default();
247        aliases.add(
248            "USDC",
249            "ethereum",
250            "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48",
251            "USD Coin",
252        );
253
254        let info = aliases.get("USDC", Some("ethereum")).unwrap();
255        assert_eq!(info.symbol, "USDC");
256        assert_eq!(info.chain, "ethereum");
257        assert!(info.address.starts_with("0x"));
258
259        // Case insensitive lookup
260        let info2 = aliases.get("usdc", Some("ethereum")).unwrap();
261        assert_eq!(info2.symbol, "USDC");
262    }
263
264    #[test]
265    fn test_get_without_chain() {
266        let mut aliases = TokenAliases::default();
267        aliases.add("USDC", "ethereum", "0xETH...", "USD Coin");
268        aliases.add("USDC", "polygon", "0xPOLY...", "USD Coin");
269
270        // Should prefer ethereum
271        let info = aliases.get("USDC", None).unwrap();
272        assert_eq!(info.chain, "ethereum");
273    }
274
275    #[test]
276    fn test_remove_alias() {
277        let mut aliases = TokenAliases::default();
278        aliases.add("USDC", "ethereum", "0x...", "USD Coin");
279        aliases.add("USDC", "polygon", "0x...", "USD Coin");
280
281        // Remove specific chain
282        aliases.remove("USDC", Some("ethereum"));
283        assert!(aliases.get("USDC", Some("ethereum")).is_none());
284        assert!(aliases.get("USDC", Some("polygon")).is_some());
285
286        // Remove all chains
287        aliases.remove("USDC", None);
288        assert!(aliases.get("USDC", None).is_none());
289    }
290
291    #[test]
292    fn test_is_address() {
293        // EVM addresses
294        assert!(TokenAliases::is_address(
295            "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"
296        ));
297        assert!(TokenAliases::is_address(
298            "0x0000000000000000000000000000000000000000"
299        ));
300
301        // Not addresses
302        assert!(!TokenAliases::is_address("USDC"));
303        assert!(!TokenAliases::is_address("ethereum"));
304        assert!(!TokenAliases::is_address("0x123")); // Too short
305    }
306
307    #[test]
308    fn test_recent_tokens() {
309        let mut aliases = TokenAliases::default();
310        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
311        aliases.add("WETH", "ethereum", "0x2...", "Wrapped Ether");
312
313        assert_eq!(aliases.recent().len(), 2);
314        // Most recent first
315        assert_eq!(aliases.recent()[0].symbol, "WETH");
316    }
317
318    #[test]
319    fn test_list_aliases() {
320        let mut aliases = TokenAliases::default();
321        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
322        aliases.add("USDC", "polygon", "0x2...", "USD Coin");
323        aliases.add("WETH", "ethereum", "0x3...", "Wrapped Ether");
324
325        let list = aliases.list();
326        assert_eq!(list.len(), 3);
327    }
328
329    #[test]
330    fn test_get_chains_for_alias() {
331        let mut aliases = TokenAliases::default();
332        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
333        aliases.add("USDC", "polygon", "0x2...", "USD Coin");
334
335        let chains = aliases.get_chains_for_alias("USDC");
336        assert_eq!(chains.len(), 2);
337        assert!(chains.contains(&"ethereum"));
338        assert!(chains.contains(&"polygon"));
339    }
340
341    #[test]
342    fn test_get_chains_for_missing_alias() {
343        let aliases = TokenAliases::default();
344        let chains = aliases.get_chains_for_alias("NONEXISTENT");
345        assert!(chains.is_empty());
346    }
347
348    #[test]
349    fn test_is_address_solana() {
350        // Valid Solana address (base58, 32-44 chars, decodes to 32 bytes)
351        assert!(TokenAliases::is_address(
352            "DRpbCBMxVnDK7maPM5tGv6MvB3v1sRMC86PZ8okm21hy"
353        ));
354        // System program address
355        assert!(TokenAliases::is_address("11111111111111111111111111111111"));
356    }
357
358    #[test]
359    fn test_is_address_tron() {
360        // Valid Tron address (starts with T, 34 chars)
361        assert!(TokenAliases::is_address(
362            "TDqSquXBgUCLYvYC4XZgrprLK589dkhSCf"
363        ));
364        assert!(TokenAliases::is_address(
365            "TR7NHqjeKQxGTCi8q8ZY4pL8otSzgjLj6t"
366        ));
367    }
368
369    #[test]
370    fn test_is_address_edge_cases() {
371        assert!(!TokenAliases::is_address("")); // Empty
372        assert!(!TokenAliases::is_address("0x")); // Incomplete EVM prefix
373        assert!(!TokenAliases::is_address("T123")); // Too short for Tron
374        assert!(!TokenAliases::is_address("hello world")); // Random text
375    }
376
377    #[test]
378    fn test_remove_specific_chain() {
379        let mut aliases = TokenAliases::default();
380        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
381        aliases.add("USDC", "polygon", "0x2...", "USD Coin");
382
383        // Remove only polygon
384        aliases.remove("USDC", Some("polygon"));
385        assert!(aliases.get("USDC", Some("polygon")).is_none());
386        assert!(aliases.get("USDC", Some("ethereum")).is_some());
387    }
388
389    #[test]
390    fn test_remove_last_chain_cleans_up() {
391        let mut aliases = TokenAliases::default();
392        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
393
394        // Removing the only chain should clean up the alias entirely
395        aliases.remove("USDC", Some("ethereum"));
396        assert!(aliases.get("USDC", None).is_none());
397        let chains = aliases.get_chains_for_alias("USDC");
398        assert!(chains.is_empty());
399    }
400
401    #[test]
402    fn test_remove_cleans_recent() {
403        let mut aliases = TokenAliases::default();
404        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
405        assert_eq!(aliases.recent().len(), 1);
406
407        aliases.remove("USDC", None);
408        assert!(aliases.recent().is_empty());
409    }
410
411    #[test]
412    fn test_add_updates_existing() {
413        let mut aliases = TokenAliases::default();
414        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
415        aliases.add("USDC", "ethereum", "0x2...", "USD Coin V2");
416
417        let info = aliases.get("USDC", Some("ethereum")).unwrap();
418        assert_eq!(info.address, "0x2...");
419        assert_eq!(info.name, "USD Coin V2");
420    }
421
422    #[test]
423    fn test_recent_truncation() {
424        let mut aliases = TokenAliases::default();
425        // Add 25 tokens, recent should be capped at 20
426        for i in 0..25 {
427            aliases.add(
428                &format!("T{}", i),
429                "ethereum",
430                &format!("0x{}...", i),
431                &format!("Token {}", i),
432            );
433        }
434        assert_eq!(aliases.recent().len(), 20);
435        // The most recent should be T24
436        assert_eq!(aliases.recent()[0].symbol, "T24");
437    }
438
439    #[test]
440    fn test_case_insensitive_operations() {
441        let mut aliases = TokenAliases::default();
442        aliases.add("usdc", "Ethereum", "0x1...", "USD Coin");
443
444        // Alias stored uppercase, chain stored lowercase
445        let info = aliases.get("USDC", Some("ethereum")).unwrap();
446        assert_eq!(info.symbol, "USDC");
447        assert_eq!(info.chain, "ethereum");
448    }
449
450    #[test]
451    fn test_token_info_has_last_used() {
452        let mut aliases = TokenAliases::default();
453        aliases.add("USDC", "ethereum", "0x1...", "USD Coin");
454        let info = aliases.get("USDC", Some("ethereum")).unwrap();
455        assert!(info.last_used.is_some());
456    }
457
458    #[test]
459    fn test_save_and_load_roundtrip() {
460        let mut aliases = TokenAliases::default();
461        aliases.add("SAVE_TEST", "ethereum", "0xsave...", "Save Test Token");
462
463        // save() writes to the standard location which should be writable in test env
464        let result = aliases.save();
465        assert!(result.is_ok());
466
467        // Load it back
468        let loaded = TokenAliases::load();
469        let info = loaded.get("SAVE_TEST", Some("ethereum"));
470        assert!(info.is_some());
471        assert_eq!(info.unwrap().address, "0xsave...");
472
473        // Cleanup: remove the test entry and save again
474        aliases.remove("SAVE_TEST", None);
475        let _ = aliases.save();
476    }
477}