lean_ctx/core/
heatmap.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::path::PathBuf;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::Mutex;
6
7const HEATMAP_FLUSH_EVERY: usize = 25;
8const HEATMAP_MAX_ENTRIES: usize = 10_000;
9
10static HEATMAP_BUFFER: Mutex<Option<HeatMap>> = Mutex::new(None);
11static HEATMAP_CALLS: AtomicUsize = AtomicUsize::new(0);
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct HeatEntry {
15    pub path: String,
16    pub access_count: u32,
17    pub last_access: String,
18    pub total_tokens_saved: u64,
19    pub total_original_tokens: u64,
20    pub avg_compression_ratio: f32,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize, Default)]
24pub struct HeatMap {
25    pub entries: HashMap<String, HeatEntry>,
26    #[serde(skip)]
27    dirty: bool,
28}
29
30impl HeatMap {
31    pub fn load() -> Self {
32        let mut guard = HEATMAP_BUFFER
33            .lock()
34            .unwrap_or_else(std::sync::PoisonError::into_inner);
35        if let Some(ref hm) = *guard {
36            return hm.clone();
37        }
38        let hm = load_from_disk();
39        *guard = Some(hm.clone());
40        hm
41    }
42
43    pub fn record_access(&mut self, file_path: &str, original_tokens: usize, saved_tokens: usize) {
44        let now = chrono::Utc::now().to_rfc3339();
45        let entry = self
46            .entries
47            .entry(file_path.to_string())
48            .or_insert_with(|| HeatEntry {
49                path: file_path.to_string(),
50                access_count: 0,
51                last_access: now.clone(),
52                total_tokens_saved: 0,
53                total_original_tokens: 0,
54                avg_compression_ratio: 0.0,
55            });
56        entry.access_count += 1;
57        entry.last_access = now;
58        entry.total_tokens_saved += saved_tokens as u64;
59        entry.total_original_tokens += original_tokens as u64;
60        if entry.total_original_tokens > 0 {
61            entry.avg_compression_ratio = 1.0
62                - (entry.total_original_tokens - entry.total_tokens_saved) as f32
63                    / entry.total_original_tokens as f32;
64        }
65        self.dirty = true;
66    }
67
68    pub fn save(&self) -> std::io::Result<()> {
69        if !self.dirty && !self.entries.is_empty() {
70            return Ok(());
71        }
72        save_to_disk(self)?;
73        let mut guard = HEATMAP_BUFFER
74            .lock()
75            .unwrap_or_else(std::sync::PoisonError::into_inner);
76        *guard = Some(self.clone());
77        Ok(())
78    }
79
80    pub fn top_files(&self, limit: usize) -> Vec<&HeatEntry> {
81        let mut sorted: Vec<&HeatEntry> = self.entries.values().collect();
82        sorted.sort_by_key(|x| std::cmp::Reverse(x.access_count));
83        sorted.truncate(limit);
84        sorted
85    }
86
87    pub fn directory_summary(&self) -> Vec<(String, u32, u64)> {
88        let mut dirs: HashMap<String, (u32, u64)> = HashMap::new();
89        for entry in self.entries.values() {
90            let dir = std::path::Path::new(&entry.path)
91                .parent()
92                .map_or_else(|| ".".to_string(), |p| p.to_string_lossy().to_string());
93            let stat = dirs.entry(dir).or_insert((0, 0));
94            stat.0 += entry.access_count;
95            stat.1 += entry.total_tokens_saved;
96        }
97        let mut result: Vec<(String, u32, u64)> = dirs
98            .into_iter()
99            .map(|(dir, (count, saved))| (dir, count, saved))
100            .collect();
101        result.sort_by_key(|x| std::cmp::Reverse(x.1));
102        result
103    }
104
105    pub fn cold_files(&self, all_files: &[String], limit: usize) -> Vec<String> {
106        let hot: std::collections::HashSet<&str> = self
107            .entries
108            .keys()
109            .map(std::string::String::as_str)
110            .collect();
111        let mut cold: Vec<String> = all_files
112            .iter()
113            .filter(|f| !hot.contains(f.as_str()))
114            .cloned()
115            .collect();
116        cold.truncate(limit);
117        cold
118    }
119
120    fn storage_path() -> PathBuf {
121        crate::core::data_dir::lean_ctx_data_dir()
122            .unwrap_or_else(|_| PathBuf::from("."))
123            .join("heatmap.json")
124    }
125}
126
127fn load_from_disk() -> HeatMap {
128    let path = HeatMap::storage_path();
129    match std::fs::read_to_string(&path) {
130        Ok(content) => serde_json::from_str(&content).unwrap_or_default(),
131        Err(_) => HeatMap::default(),
132    }
133}
134
135fn save_to_disk(hm: &HeatMap) -> std::io::Result<()> {
136    let path = HeatMap::storage_path();
137    if let Some(parent) = path.parent() {
138        std::fs::create_dir_all(parent)?;
139    }
140    let json = serde_json::to_string_pretty(hm)?;
141    let tmp = path.with_extension("json.tmp");
142    std::fs::write(&tmp, &json)?;
143    std::fs::rename(&tmp, &path)
144}
145
146pub fn record_file_access(file_path: &str, original_tokens: usize, saved_tokens: usize) {
147    let file_path = std::fs::canonicalize(file_path).map_or_else(
148        |_| file_path.to_string(),
149        |p| p.to_string_lossy().into_owned(),
150    );
151    let file_path = file_path.as_str();
152
153    let mut guard = HEATMAP_BUFFER
154        .lock()
155        .unwrap_or_else(std::sync::PoisonError::into_inner);
156    let hm = guard.get_or_insert_with(load_from_disk);
157    hm.record_access(file_path, original_tokens, saved_tokens);
158
159    // Enforce bounded retention.
160    if hm.entries.len() > HEATMAP_MAX_ENTRIES {
161        let mut items: Vec<(String, u32)> = hm
162            .entries
163            .values()
164            .map(|e| (e.path.clone(), e.access_count))
165            .collect();
166        items.sort_by_key(|x| x.1);
167        let drop_n = hm.entries.len().saturating_sub(HEATMAP_MAX_ENTRIES);
168        for (path, _) in items.into_iter().take(drop_n) {
169            hm.entries.remove(&path);
170        }
171    }
172
173    let n = HEATMAP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
174    if n.is_multiple_of(HEATMAP_FLUSH_EVERY) && save_to_disk(hm).is_ok() {
175        hm.dirty = false;
176    }
177}
178
179pub fn flush() {
180    let guard = HEATMAP_BUFFER
181        .lock()
182        .unwrap_or_else(std::sync::PoisonError::into_inner);
183    if let Some(ref hm) = *guard {
184        if hm.dirty {
185            let _ = save_to_disk(hm);
186        }
187    }
188}
189
190pub fn reset() {
191    let mut guard = HEATMAP_BUFFER
192        .lock()
193        .unwrap_or_else(std::sync::PoisonError::into_inner);
194    *guard = Some(HeatMap::default());
195    if let Some(hm) = guard.as_ref() {
196        let _ = save_to_disk(hm);
197    }
198}
199
200pub fn format_heatmap_status(heatmap: &HeatMap, limit: usize) -> String {
201    let top = heatmap.top_files(limit);
202    if top.is_empty() {
203        return "No file access data recorded yet.".to_string();
204    }
205    let mut lines = vec![format!(
206        "File Access Heat Map ({} tracked files):",
207        heatmap.entries.len()
208    )];
209    lines.push(String::new());
210    for (i, entry) in top.iter().enumerate() {
211        let short = short_path(&entry.path);
212        let heat = heat_indicator(entry.access_count);
213        lines.push(format!(
214            "  {heat} #{} {} — {} accesses, {:.0}% compression, {} tok saved",
215            i + 1,
216            short,
217            entry.access_count,
218            entry.avg_compression_ratio * 100.0,
219            entry.total_tokens_saved
220        ));
221    }
222    lines.join("\n")
223}
224
225pub fn format_directory_summary(heatmap: &HeatMap) -> String {
226    let dirs = heatmap.directory_summary();
227    if dirs.is_empty() {
228        return "No directory data.".to_string();
229    }
230    let mut lines = vec!["Directory Heat Map:".to_string(), String::new()];
231    for (dir, count, saved) in dirs.iter().take(15) {
232        let heat = heat_indicator(*count);
233        lines.push(format!(
234            "  {heat} {dir}/ — {count} accesses, {saved} tok saved"
235        ));
236    }
237    lines.join("\n")
238}
239
240fn heat_indicator(count: u32) -> &'static str {
241    match count {
242        0 => "  ",
243        1..=3 => "▁▁",
244        4..=8 => "▃▃",
245        9..=15 => "▅▅",
246        16..=30 => "▇▇",
247        _ => "██",
248    }
249}
250
251fn short_path(path: &str) -> &str {
252    let parts: Vec<&str> = path.rsplitn(3, '/').collect();
253    if parts.len() >= 2 {
254        let start = path.len() - parts[0].len() - parts[1].len() - 1;
255        &path[start..]
256    } else {
257        path
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    #[test]
266    fn record_and_query() {
267        let mut hm = HeatMap::default();
268        hm.record_access("src/main.rs", 100, 80);
269        hm.record_access("src/main.rs", 100, 90);
270        hm.record_access("src/lib.rs", 200, 50);
271
272        assert_eq!(hm.entries.len(), 2);
273        assert_eq!(hm.entries["src/main.rs"].access_count, 2);
274        assert_eq!(hm.entries["src/lib.rs"].total_tokens_saved, 50);
275    }
276
277    #[test]
278    fn top_files_sorted() {
279        let mut hm = HeatMap::default();
280        hm.record_access("a.rs", 100, 50);
281        hm.record_access("b.rs", 100, 50);
282        hm.record_access("b.rs", 100, 50);
283        hm.record_access("c.rs", 100, 50);
284        hm.record_access("c.rs", 100, 50);
285        hm.record_access("c.rs", 100, 50);
286
287        let top = hm.top_files(2);
288        assert_eq!(top.len(), 2);
289        assert_eq!(top[0].path, "c.rs");
290        assert_eq!(top[1].path, "b.rs");
291    }
292
293    #[test]
294    fn directory_summary_works() {
295        let mut hm = HeatMap::default();
296        hm.record_access("src/a.rs", 100, 50);
297        hm.record_access("src/b.rs", 100, 50);
298        hm.record_access("tests/t.rs", 200, 100);
299
300        let dirs = hm.directory_summary();
301        assert!(dirs.len() >= 2);
302    }
303
304    #[test]
305    fn cold_files_detection() {
306        let mut hm = HeatMap::default();
307        hm.record_access("src/a.rs", 100, 50);
308
309        let all = vec![
310            "src/a.rs".to_string(),
311            "src/b.rs".to_string(),
312            "src/c.rs".to_string(),
313        ];
314        let cold = hm.cold_files(&all, 10);
315        assert_eq!(cold.len(), 2);
316        assert!(cold.contains(&"src/b.rs".to_string()));
317    }
318
319    #[test]
320    fn heat_indicators() {
321        assert_eq!(heat_indicator(0), "  ");
322        assert_eq!(heat_indicator(1), "▁▁");
323        assert_eq!(heat_indicator(10), "▅▅");
324        assert_eq!(heat_indicator(50), "██");
325    }
326
327    #[test]
328    fn compression_ratio() {
329        let mut hm = HeatMap::default();
330        hm.record_access("a.rs", 1000, 800);
331        let entry = &hm.entries["a.rs"];
332        assert!((entry.avg_compression_ratio - 0.8).abs() < 0.01);
333    }
334}
lean_ctx/core/heatmap.rs

lean_ctx/core/
heatmap.rs
