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

lean_ctx/core/
heatmap.rs
