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

lean_ctx/core/
heatmap.rs
