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

lean_ctx/core/
heatmap.rs
