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| b.heat_score.partial_cmp(&a.heat_score).unwrap()),
51        SortBy::Tokens => sorted.sort_by(|a, b| b.token_count.cmp(&a.token_count)),
52        SortBy::Connections => sorted.sort_by(|a, b| b.connections.cmp(&a.connections)),
53    }
54
55    let top = &sorted[..sorted.len().min(top_n)];
56
57    if json_output {
58        print_json(top);
59    } else {
60        print_heatmap(&project_root, top, &sorted);
61    }
62}
63
64enum SortBy {
65    Heat,
66    Tokens,
67    Connections,
68}
69
70fn build_heat_entries(index: &ProjectIndex, dir_filter: Option<&str>) -> Vec<HeatEntry> {
71    let mut connection_counts: HashMap<String, usize> = HashMap::new();
72    for edge in &index.edges {
73        *connection_counts.entry(edge.from.clone()).or_default() += 1;
74        *connection_counts.entry(edge.to.clone()).or_default() += 1;
75    }
76
77    let max_tokens = index
78        .files
79        .values()
80        .map(|f| f.token_count)
81        .max()
82        .unwrap_or(1) as f64;
83    let max_connections = connection_counts.values().max().copied().unwrap_or(1) as f64;
84
85    index
86        .files
87        .values()
88        .filter(|f| {
89            if let Some(dir) = dir_filter {
90                f.path.starts_with(dir) || f.path.starts_with(&format!("./{dir}"))
91            } else {
92                true
93            }
94        })
95        .map(|f| {
96            let connections = connection_counts.get(&f.path).copied().unwrap_or(0);
97            let token_norm = f.token_count as f64 / max_tokens;
98            let conn_norm = connections as f64 / max_connections;
99            let heat_score = token_norm * 0.4 + conn_norm * 0.6;
100
101            HeatEntry {
102                path: f.path.clone(),
103                token_count: f.token_count,
104                connections,
105                heat_score,
106            }
107        })
108        .collect()
109}
110
111fn heat_color(score: f64) -> &'static str {
112    if score > 0.8 {
113        "\x1b[91m" // bright red
114    } else if score > 0.6 {
115        "\x1b[31m" // red
116    } else if score > 0.4 {
117        "\x1b[33m" // yellow
118    } else if score > 0.2 {
119        "\x1b[36m" // cyan
120    } else {
121        "\x1b[34m" // blue
122    }
123}
124
125fn heat_bar(score: f64, width: usize) -> String {
126    let filled = (score * width as f64).round() as usize;
127    let blocks = "█".repeat(filled);
128    let empty = "░".repeat(width.saturating_sub(filled));
129    format!("{}{blocks}\x1b[38;5;239m{empty}\x1b[0m", heat_color(score))
130}
131
132fn print_heatmap(project_root: &str, entries: &[HeatEntry], all: &[HeatEntry]) {
133    let total_files = all.len();
134    let total_tokens: usize = all.iter().map(|e| e.token_count).sum();
135    let total_connections: usize = all.iter().map(|e| e.connections).sum();
136
137    let project_name = std::path::Path::new(project_root)
138        .file_name()
139        .map(|n| n.to_string_lossy().to_string())
140        .unwrap_or_else(|| project_root.to_string());
141
142    println!();
143    println!(
144        "\x1b[1;37m  Context Heat Map\x1b[0m  \x1b[38;5;239m{}\x1b[0m",
145        project_name
146    );
147    println!(
148        "\x1b[38;5;239m  {} files · {} tokens · {} connections\x1b[0m",
149        total_files, total_tokens, total_connections
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
