lean_ctx/
heatmap.rs

1use crate::core::graph_index::{self, ProjectIndex};
2use std::collections::HashMap;
3
4struct HeatEntry {
5    path: String,
6    token_count: usize,
7    connections: usize,
8    heat_score: f64,
9}
10
11pub fn cmd_heatmap(args: &[String]) {
12    let project_root = std::env::current_dir()
13        .ok()
14        .and_then(|d| d.to_str().map(String::from))
15        .unwrap_or_else(|| ".".to_string());
16
17    let top_n: usize = args
18        .iter()
19        .find_map(|a| a.strip_prefix("--top="))
20        .and_then(|v| v.parse().ok())
21        .unwrap_or(20);
22
23    let dir_filter: Option<&str> = args
24        .iter()
25        .find_map(|a| a.strip_prefix("--dir="))
26        .map(|s| s.trim_end_matches('/'));
27
28    let sort_by = if args.iter().any(|a| a == "--by=connections") {
29        SortBy::Connections
30    } else if args.iter().any(|a| a == "--by=tokens") {
31        SortBy::Tokens
32    } else {
33        SortBy::Heat
34    };
35
36    let json_output = args.iter().any(|a| a == "--json");
37
38    let index = graph_index::load_or_build(&project_root);
39
40    let entries = build_heat_entries(&index, dir_filter);
41
42    if entries.is_empty() {
43        eprintln!("No files found in project graph.");
44        eprintln!("  Run: lean-ctx setup  (to build the project graph)");
45        return;
46    }
47
48    let mut sorted = entries;
49    match sort_by {
50        SortBy::Heat => sorted.sort_by(|a, b| {
51            b.heat_score
52                .partial_cmp(&a.heat_score)
53                .unwrap_or(std::cmp::Ordering::Equal)
54        }),
55        SortBy::Tokens => sorted.sort_by_key(|x| std::cmp::Reverse(x.token_count)),
56        SortBy::Connections => sorted.sort_by_key(|x| std::cmp::Reverse(x.connections)),
57    }
58
59    let top = &sorted[..sorted.len().min(top_n)];
60
61    if json_output {
62        print_json(top);
63    } else {
64        print_heatmap(&project_root, top, &sorted);
65    }
66}
67
68enum SortBy {
69    Heat,
70    Tokens,
71    Connections,
72}
73
74fn build_heat_entries(index: &ProjectIndex, dir_filter: Option<&str>) -> Vec<HeatEntry> {
75    let mut connection_counts: HashMap<String, usize> = HashMap::new();
76    for edge in &index.edges {
77        *connection_counts.entry(edge.from.clone()).or_default() += 1;
78        *connection_counts.entry(edge.to.clone()).or_default() += 1;
79    }
80
81    let max_tokens = index
82        .files
83        .values()
84        .map(|f| f.token_count)
85        .max()
86        .unwrap_or(1) as f64;
87    let max_connections = connection_counts.values().max().copied().unwrap_or(1) as f64;
88
89    index
90        .files
91        .values()
92        .filter(|f| {
93            if let Some(dir) = dir_filter {
94                f.path.starts_with(dir) || f.path.starts_with(&format!("./{dir}"))
95            } else {
96                true
97            }
98        })
99        .map(|f| {
100            let connections = connection_counts.get(&f.path).copied().unwrap_or(0);
101            let token_norm = f.token_count as f64 / max_tokens;
102            let conn_norm = connections as f64 / max_connections;
103            let heat_score = token_norm * 0.4 + conn_norm * 0.6;
104
105            HeatEntry {
106                path: f.path.clone(),
107                token_count: f.token_count,
108                connections,
109                heat_score,
110            }
111        })
112        .collect()
113}
114
115fn heat_color(score: f64) -> &'static str {
116    if score > 0.8 {
117        "\x1b[91m" // bright red
118    } else if score > 0.6 {
119        "\x1b[31m" // red
120    } else if score > 0.4 {
121        "\x1b[33m" // yellow
122    } else if score > 0.2 {
123        "\x1b[36m" // cyan
124    } else {
125        "\x1b[34m" // blue
126    }
127}
128
129fn heat_bar(score: f64, width: usize) -> String {
130    let filled = (score * width as f64).round() as usize;
131    let blocks = "█".repeat(filled);
132    let empty = "░".repeat(width.saturating_sub(filled));
133    format!("{}{blocks}\x1b[38;5;239m{empty}\x1b[0m", heat_color(score))
134}
135
136fn print_heatmap(project_root: &str, entries: &[HeatEntry], all: &[HeatEntry]) {
137    let total_files = all.len();
138    let total_tokens: usize = all.iter().map(|e| e.token_count).sum();
139    let total_connections: usize = all.iter().map(|e| e.connections).sum();
140
141    let project_name = std::path::Path::new(project_root).file_name().map_or_else(
142        || project_root.to_string(),
143        |n| n.to_string_lossy().to_string(),
144    );
145
146    println!();
147    println!("\x1b[1;37m  Context Heat Map\x1b[0m  \x1b[38;5;239m{project_name}\x1b[0m");
148    println!(
149        "\x1b[38;5;239m  {total_files} files · {total_tokens} tokens · {total_connections} connections\x1b[0m"
150    );
151    println!();
152
153    let max_path_len = entries.iter().map(|e| e.path.len()).max().unwrap_or(30);
154    let path_width = max_path_len.min(50);
155
156    println!(
157        "  \x1b[38;5;239m{:<width$}  {:>6}  {:>5}  HEAT\x1b[0m",
158        "FILE",
159        "TOKENS",
160        "CONNS",
161        width = path_width
162    );
163    println!("  \x1b[38;5;239m{}\x1b[0m", "─".repeat(path_width + 32));
164
165    for entry in entries {
166        let display_path = if entry.path.len() > path_width {
167            let skip = entry.path.len() - path_width + 3;
168            format!("...{}", &entry.path[skip..])
169        } else {
170            entry.path.clone()
171        };
172
173        let bar = heat_bar(entry.heat_score, 16);
174
175        println!(
176            "  {color}{:<width$}\x1b[0m  \x1b[38;5;245m{:>6}\x1b[0m  \x1b[38;5;245m{:>5}\x1b[0m  {bar}  {color}{:.0}%\x1b[0m",
177            display_path,
178            entry.token_count,
179            entry.connections,
180            entry.heat_score * 100.0,
181            color = heat_color(entry.heat_score),
182            width = path_width,
183        );
184    }
185
186    println!();
187    println!(
188        "  \x1b[38;5;239mLegend: \x1b[91m█\x1b[38;5;239m hot  \x1b[33m█\x1b[38;5;239m warm  \x1b[36m█\x1b[38;5;239m cool  \x1b[34m█\x1b[38;5;239m cold\x1b[0m"
189    );
190    println!(
191        "  \x1b[38;5;239mOptions: --top=N  --dir=path  --by=tokens|connections  --json\x1b[0m"
192    );
193    println!();
194}
195
196fn print_json(entries: &[HeatEntry]) {
197    let items: Vec<serde_json::Value> = entries
198        .iter()
199        .map(|e| {
200            serde_json::json!({
201                "path": e.path,
202                "token_count": e.token_count,
203                "connections": e.connections,
204                "heat_score": (e.heat_score * 100.0).round() / 100.0,
205            })
206        })
207        .collect();
208
209    println!(
210        "{}",
211        serde_json::to_string_pretty(&items).unwrap_or_else(|_| "[]".to_string())
212    );
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    #[test]
220    fn test_heat_color_ranges() {
221        assert_eq!(heat_color(0.9), "\x1b[91m");
222        assert_eq!(heat_color(0.7), "\x1b[31m");
223        assert_eq!(heat_color(0.5), "\x1b[33m");
224        assert_eq!(heat_color(0.3), "\x1b[36m");
225        assert_eq!(heat_color(0.1), "\x1b[34m");
226    }
227
228    #[test]
229    fn test_heat_bar_length() {
230        let bar = heat_bar(0.5, 10);
231        assert!(bar.contains("█████"));
232    }
233
234    #[test]
235    fn test_build_heat_entries_empty() {
236        let index = ProjectIndex::new(".");
237        let entries = build_heat_entries(&index, None);
238        assert!(entries.is_empty());
239    }
240}
lean_ctx/heatmap.rs

lean_ctx/
heatmap.rs
