Skip to main content

ccstat_core/
string_pool.rs

1//! String interning for memory optimization
2//!
3//! This module provides string interning capabilities to reduce memory usage
4//! when dealing with repeated strings like model names and session IDs.
5
6use once_cell::sync::Lazy;
7use std::sync::RwLock;
8use string_interner::{DefaultBackend, DefaultSymbol, StringInterner};
9
10/// Global string interner for model names
11static MODEL_INTERNER: Lazy<RwLock<StringInterner<DefaultBackend>>> =
12    Lazy::new(|| RwLock::new(StringInterner::default()));
13
14/// Global string interner for session IDs
15static SESSION_INTERNER: Lazy<RwLock<StringInterner<DefaultBackend>>> =
16    Lazy::new(|| RwLock::new(StringInterner::default()));
17
18/// Interned model name
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub struct InternedModel(DefaultSymbol);
21
22impl InternedModel {
23    /// Create or get an interned model name
24    pub fn new(model: &str) -> Self {
25        let mut interner = MODEL_INTERNER.write().unwrap();
26        let symbol = interner.get_or_intern(model);
27        Self(symbol)
28    }
29
30    /// Get the string value
31    pub fn as_str(&self) -> String {
32        let interner = MODEL_INTERNER.read().unwrap();
33        interner.resolve(self.0).unwrap_or("unknown").to_string()
34    }
35}
36
37/// Interned session ID
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
39pub struct InternedSession(DefaultSymbol);
40
41impl InternedSession {
42    /// Create or get an interned session ID
43    pub fn new(session: &str) -> Self {
44        let mut interner = SESSION_INTERNER.write().unwrap();
45        let symbol = interner.get_or_intern(session);
46        Self(symbol)
47    }
48
49    /// Get the string value
50    pub fn as_str(&self) -> String {
51        let interner = SESSION_INTERNER.read().unwrap();
52        interner.resolve(self.0).unwrap_or("unknown").to_string()
53    }
54}
55
56/// Statistics about string interning
57pub struct InternerStats {
58    pub model_count: usize,
59    pub session_count: usize,
60    pub model_memory_saved: usize,
61    pub session_memory_saved: usize,
62}
63
64impl InternerStats {
65    /// Get current interner statistics
66    pub fn current() -> Self {
67        let model_interner = MODEL_INTERNER.read().unwrap();
68        let session_interner = SESSION_INTERNER.read().unwrap();
69
70        // Estimate memory savings (rough calculation)
71        let avg_model_len = 20; // Average model name length
72        let avg_session_len = 36; // Average session ID length
73
74        let model_count = model_interner.len();
75        let session_count = session_interner.len();
76
77        // Memory saved = (number of duplicates) * (average string size)
78        // This is a rough estimate
79        let model_memory_saved = model_count.saturating_sub(10) * avg_model_len;
80        let session_memory_saved = session_count.saturating_sub(100) * avg_session_len;
81
82        Self {
83            model_count,
84            session_count,
85            model_memory_saved,
86            session_memory_saved,
87        }
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn test_model_interning() {
97        let model1 = InternedModel::new("claude-3-opus");
98        let model2 = InternedModel::new("claude-3-opus");
99        let model3 = InternedModel::new("claude-3-sonnet");
100
101        // Same strings should have same symbol
102        assert_eq!(model1, model2);
103        assert_ne!(model1, model3);
104
105        // Should be able to retrieve strings
106        assert_eq!(model1.as_str(), "claude-3-opus");
107        assert_eq!(model3.as_str(), "claude-3-sonnet");
108    }
109
110    #[test]
111    fn test_session_interning() {
112        let session1 = InternedSession::new("session-123");
113        let session2 = InternedSession::new("session-123");
114        let session3 = InternedSession::new("session-456");
115
116        assert_eq!(session1, session2);
117        assert_ne!(session1, session3);
118
119        assert_eq!(session1.as_str(), "session-123");
120        assert_eq!(session3.as_str(), "session-456");
121    }
122
123    #[test]
124    fn test_interner_stats() {
125        // Create some interned strings
126        for i in 0..5 {
127            InternedModel::new(&format!("model-{i}"));
128            InternedSession::new(&format!("session-{i}"));
129        }
130
131        let stats = InternerStats::current();
132        assert!(stats.model_count >= 5);
133        assert!(stats.session_count >= 5);
134    }
135}