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    std::fs::write(&path, json)
142}
143
144pub fn record_file_access(file_path: &str, original_tokens: usize, saved_tokens: usize) {
145    let mut guard = HEATMAP_BUFFER
146        .lock()
147        .unwrap_or_else(std::sync::PoisonError::into_inner);
148    let hm = guard.get_or_insert_with(load_from_disk);
149    hm.record_access(file_path, original_tokens, saved_tokens);
150
151    // Enforce bounded retention.
152    if hm.entries.len() > HEATMAP_MAX_ENTRIES {
153        let mut items: Vec<(String, u32)> = hm
154            .entries
155            .values()
156            .map(|e| (e.path.clone(), e.access_count))
157            .collect();
158        items.sort_by_key(|x| x.1);
159        let drop_n = hm.entries.len().saturating_sub(HEATMAP_MAX_ENTRIES);
160        for (path, _) in items.into_iter().take(drop_n) {
161            hm.entries.remove(&path);
162        }
163    }
164
165    let n = HEATMAP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
166    if n.is_multiple_of(HEATMAP_FLUSH_EVERY) && save_to_disk(hm).is_ok() {
167        hm.dirty = false;
168    }
169}
170
171pub fn format_heatmap_status(heatmap: &HeatMap, limit: usize) -> String {
172    let top = heatmap.top_files(limit);
173    if top.is_empty() {
174        return "No file access data recorded yet.".to_string();
175    }
176    let mut lines = vec![format!(
177        "File Access Heat Map ({} tracked files):",
178        heatmap.entries.len()
179    )];
180    lines.push(String::new());
181    for (i, entry) in top.iter().enumerate() {
182        let short = short_path(&entry.path);
183        let heat = heat_indicator(entry.access_count);
184        lines.push(format!(
185            "  {heat} #{} {} — {} accesses, {:.0}% compression, {} tok saved",
186            i + 1,
187            short,
188            entry.access_count,
189            entry.avg_compression_ratio * 100.0,
190            entry.total_tokens_saved
191        ));
192    }
193    lines.join("\n")
194}
195
196pub fn format_directory_summary(heatmap: &HeatMap) -> String {
197    let dirs = heatmap.directory_summary();
198    if dirs.is_empty() {
199        return "No directory data.".to_string();
200    }
201    let mut lines = vec!["Directory Heat Map:".to_string(), String::new()];
202    for (dir, count, saved) in dirs.iter().take(15) {
203        let heat = heat_indicator(*count);
204        lines.push(format!(
205            "  {heat} {dir}/ — {count} accesses, {saved} tok saved"
206        ));
207    }
208    lines.join("\n")
209}
210
211fn heat_indicator(count: u32) -> &'static str {
212    match count {
213        0 => "  ",
214        1..=3 => "▁▁",
215        4..=8 => "▃▃",
216        9..=15 => "▅▅",
217        16..=30 => "▇▇",
218        _ => "██",
219    }
220}
221
222fn short_path(path: &str) -> &str {
223    let parts: Vec<&str> = path.rsplitn(3, '/').collect();
224    if parts.len() >= 2 {
225        let start = path.len() - parts[0].len() - parts[1].len() - 1;
226        &path[start..]
227    } else {
228        path
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn record_and_query() {
238        let mut hm = HeatMap::default();
239        hm.record_access("src/main.rs", 100, 80);
240        hm.record_access("src/main.rs", 100, 90);
241        hm.record_access("src/lib.rs", 200, 50);
242
243        assert_eq!(hm.entries.len(), 2);
244        assert_eq!(hm.entries["src/main.rs"].access_count, 2);
245        assert_eq!(hm.entries["src/lib.rs"].total_tokens_saved, 50);
246    }
247
248    #[test]
249    fn top_files_sorted() {
250        let mut hm = HeatMap::default();
251        hm.record_access("a.rs", 100, 50);
252        hm.record_access("b.rs", 100, 50);
253        hm.record_access("b.rs", 100, 50);
254        hm.record_access("c.rs", 100, 50);
255        hm.record_access("c.rs", 100, 50);
256        hm.record_access("c.rs", 100, 50);
257
258        let top = hm.top_files(2);
259        assert_eq!(top.len(), 2);
260        assert_eq!(top[0].path, "c.rs");
261        assert_eq!(top[1].path, "b.rs");
262    }
263
264    #[test]
265    fn directory_summary_works() {
266        let mut hm = HeatMap::default();
267        hm.record_access("src/a.rs", 100, 50);
268        hm.record_access("src/b.rs", 100, 50);
269        hm.record_access("tests/t.rs", 200, 100);
270
271        let dirs = hm.directory_summary();
272        assert!(dirs.len() >= 2);
273    }
274
275    #[test]
276    fn cold_files_detection() {
277        let mut hm = HeatMap::default();
278        hm.record_access("src/a.rs", 100, 50);
279
280        let all = vec![
281            "src/a.rs".to_string(),
282            "src/b.rs".to_string(),
283            "src/c.rs".to_string(),
284        ];
285        let cold = hm.cold_files(&all, 10);
286        assert_eq!(cold.len(), 2);
287        assert!(cold.contains(&"src/b.rs".to_string()));
288    }
289
290    #[test]
291    fn heat_indicators() {
292        assert_eq!(heat_indicator(0), "  ");
293        assert_eq!(heat_indicator(1), "▁▁");
294        assert_eq!(heat_indicator(10), "▅▅");
295        assert_eq!(heat_indicator(50), "██");
296    }
297
298    #[test]
299    fn compression_ratio() {
300        let mut hm = HeatMap::default();
301        hm.record_access("a.rs", 1000, 800);
302        let entry = &hm.entries["a.rs"];
303        assert!((entry.avg_compression_ratio - 0.8).abs() < 0.01);
304    }
305}
lean_ctx/core/heatmap.rs

lean_ctx/core/
heatmap.rs
