Skip to main content

lean_ctx/tools/
ctx_graph.rs

1use std::collections::HashMap;
2use std::path::Path;
3
4use crate::core::graph_index::{self, ProjectIndex};
5use crate::core::tokens::count_tokens;
6
7pub fn handle(
8    action: &str,
9    path: Option<&str>,
10    root: &str,
11    cache: &mut crate::core::cache::SessionCache,
12    crp_mode: crate::tools::CrpMode,
13) -> String {
14    match action {
15        "build" => handle_build(root),
16        "related" => handle_related(path, root),
17        "symbol" => handle_symbol(path, root, cache, crp_mode),
18        "impact" => handle_impact(path, root),
19        "status" => handle_status(root),
20        _ => "Unknown action. Use: build, related, symbol, impact, status".to_string(),
21    }
22}
23
24fn handle_build(root: &str) -> String {
25    let index = graph_index::scan(root);
26
27    let mut by_lang: HashMap<&str, (usize, usize)> = HashMap::new();
28    for entry in index.files.values() {
29        let e = by_lang.entry(&entry.language).or_insert((0, 0));
30        e.0 += 1;
31        e.1 += entry.token_count;
32    }
33
34    let mut result = Vec::new();
35    result.push(format!(
36        "Project Graph: {} files, {} symbols, {} edges",
37        index.file_count(),
38        index.symbol_count(),
39        index.edge_count()
40    ));
41
42    let mut langs: Vec<_> = by_lang.iter().collect();
43    langs.sort_by_key(|(_, v)| std::cmp::Reverse(v.1));
44    result.push("\nLanguages:".to_string());
45    for (lang, (count, tokens)) in &langs {
46        result.push(format!("  {lang}: {count} files, {tokens} tok"));
47    }
48
49    let mut import_counts: HashMap<&str, usize> = HashMap::new();
50    for edge in &index.edges {
51        if edge.kind == "import" {
52            *import_counts.entry(&edge.to).or_insert(0) += 1;
53        }
54    }
55    let mut hotspots: Vec<_> = import_counts.iter().collect();
56    hotspots.sort_by_key(|x| std::cmp::Reverse(*x.1));
57
58    if !hotspots.is_empty() {
59        result.push(format!("\nMost imported ({}):", hotspots.len().min(10)));
60        for (module, count) in hotspots.iter().take(10) {
61            result.push(format!("  {module}: imported by {count} files"));
62        }
63    }
64
65    if let Some(dir) = ProjectIndex::index_dir(root) {
66        result.push(format!(
67            "\nIndex saved: {}",
68            crate::core::protocol::shorten_path(&dir.to_string_lossy())
69        ));
70    }
71
72    let output = result.join("\n");
73    let tokens = count_tokens(&output);
74    format!("{output}\n[ctx_graph build: {tokens} tok]")
75}
76
77fn handle_related(path: Option<&str>, root: &str) -> String {
78    let target = match path {
79        Some(p) => p,
80        None => return "path is required for 'related' action".to_string(),
81    };
82
83    let index = match ProjectIndex::load(root) {
84        Some(idx) => idx,
85        None => {
86            return "No graph index found. Run ctx_graph with action='build' first.".to_string()
87        }
88    };
89
90    let rel_target = graph_index::graph_relative_key(target, root);
91
92    let related = index.get_related(&rel_target, 2);
93    if related.is_empty() {
94        return format!(
95            "No related files found for {}",
96            crate::core::protocol::shorten_path(target)
97        );
98    }
99
100    let mut result = format!(
101        "Files related to {} ({}):\n",
102        crate::core::protocol::shorten_path(target),
103        related.len()
104    );
105    for r in &related {
106        result.push_str(&format!("  {}\n", crate::core::protocol::shorten_path(r)));
107    }
108
109    let tokens = count_tokens(&result);
110    format!("{result}[ctx_graph related: {tokens} tok]")
111}
112
113fn handle_symbol(
114    path: Option<&str>,
115    root: &str,
116    cache: &mut crate::core::cache::SessionCache,
117    crp_mode: crate::tools::CrpMode,
118) -> String {
119    let spec = match path {
120        Some(p) => p,
121        None => {
122            return "path is required for 'symbol' action (format: file.rs::function_name)"
123                .to_string()
124        }
125    };
126
127    let (file_part, symbol_name) = match spec.split_once("::") {
128        Some((f, s)) => (f, s),
129        None => return format!("Invalid symbol spec '{spec}'. Use format: file.rs::function_name"),
130    };
131
132    let index = match ProjectIndex::load(root) {
133        Some(idx) => idx,
134        None => {
135            return "No graph index found. Run ctx_graph with action='build' first.".to_string()
136        }
137    };
138
139    let rel_file = graph_index::graph_relative_key(file_part, root);
140
141    let key = format!("{rel_file}::{symbol_name}");
142    let symbol = match index.get_symbol(&key) {
143        Some(s) => s,
144        None => {
145            let available: Vec<&str> = index
146                .symbols
147                .keys()
148                .filter(|k| k.starts_with(&rel_file))
149                .map(|k| k.as_str())
150                .take(10)
151                .collect();
152            if available.is_empty() {
153                return format!("Symbol '{symbol_name}' not found in {rel_file}. Run ctx_graph action='build' to update the index.");
154            }
155            return format!(
156                "Symbol '{symbol_name}' not found in {rel_file}.\nAvailable symbols:\n  {}",
157                available.join("\n  ")
158            );
159        }
160    };
161
162    let abs_path = if Path::new(file_part).is_absolute() {
163        file_part.to_string()
164    } else {
165        Path::new(root)
166            .join(rel_file.trim_start_matches(['/', '\\']))
167            .to_string_lossy()
168            .to_string()
169    };
170
171    let content = match std::fs::read_to_string(&abs_path) {
172        Ok(c) => c,
173        Err(e) => return format!("Cannot read {abs_path}: {e}"),
174    };
175
176    let lines: Vec<&str> = content.lines().collect();
177    let start = symbol.start_line.saturating_sub(1);
178    let end = symbol.end_line.min(lines.len());
179
180    if start >= lines.len() {
181        return crate::tools::ctx_read::handle(cache, &abs_path, "full", crp_mode);
182    }
183
184    let mut result = format!(
185        "{}::{} ({}:{}-{})\n",
186        crate::core::protocol::shorten_path(&rel_file),
187        symbol_name,
188        symbol.kind,
189        symbol.start_line,
190        symbol.end_line
191    );
192
193    for (i, line) in lines[start..end].iter().enumerate() {
194        result.push_str(&format!("{:>4}|{}\n", start + i + 1, line));
195    }
196
197    let tokens = count_tokens(&result);
198    let full_tokens = count_tokens(&content);
199    let saved = full_tokens.saturating_sub(tokens);
200    let pct = if full_tokens > 0 {
201        (saved as f64 / full_tokens as f64 * 100.0).round() as usize
202    } else {
203        0
204    };
205
206    format!("{result}[ctx_graph symbol: {tokens} tok (full file: {full_tokens} tok, -{pct}%)]")
207}
208
209fn file_path_to_module_prefixes(
210    rel_path: &str,
211    project_root: &str,
212    index: &ProjectIndex,
213) -> Vec<String> {
214    let rel_path_slash = graph_index::graph_match_key(rel_path);
215    let without_ext = rel_path_slash
216        .strip_suffix(".rs")
217        .or_else(|| rel_path_slash.strip_suffix(".ts"))
218        .or_else(|| rel_path_slash.strip_suffix(".tsx"))
219        .or_else(|| rel_path_slash.strip_suffix(".js"))
220        .or_else(|| rel_path_slash.strip_suffix(".py"))
221        .or_else(|| rel_path_slash.strip_suffix(".kt"))
222        .or_else(|| rel_path_slash.strip_suffix(".kts"))
223        .unwrap_or(&rel_path_slash);
224
225    let module_path = without_ext
226        .strip_prefix("src/")
227        .unwrap_or(without_ext)
228        .replace('/', "::");
229
230    let module_path = if module_path.ends_with("::mod") {
231        module_path
232            .strip_suffix("::mod")
233            .unwrap_or(&module_path)
234            .to_string()
235    } else {
236        module_path
237    };
238
239    let crate_name = std::fs::read_to_string(Path::new(project_root).join("Cargo.toml"))
240        .or_else(|_| std::fs::read_to_string(Path::new(project_root).join("package.json")))
241        .ok()
242        .and_then(|c| {
243            c.lines()
244                .find(|l| l.contains("\"name\"") || l.starts_with("name"))
245                .and_then(|l| l.split('"').nth(1))
246                .map(|n| n.replace('-', "_"))
247        })
248        .unwrap_or_default();
249
250    let mut prefixes = vec![
251        format!("crate::{module_path}"),
252        format!("super::{module_path}"),
253        module_path.clone(),
254    ];
255    if !crate_name.is_empty() {
256        prefixes.insert(0, format!("{crate_name}::{module_path}"));
257    }
258
259    let ext = Path::new(rel_path)
260        .extension()
261        .and_then(|e| e.to_str())
262        .unwrap_or("");
263    if matches!(ext, "kt" | "kts") {
264        let abs_path = Path::new(project_root).join(rel_path.trim_start_matches(['/', '\\']));
265        if let Ok(content) = std::fs::read_to_string(abs_path) {
266            if let Some(package_name) = content.lines().map(str::trim).find_map(|line| {
267                line.strip_prefix("package ")
268                    .map(|rest| rest.trim().trim_end_matches(';').to_string())
269            }) {
270                prefixes.push(package_name.clone());
271                if let Some(entry) = index.files.get(rel_path) {
272                    for export in &entry.exports {
273                        prefixes.push(format!("{package_name}.{export}"));
274                    }
275                }
276                if let Some(file_stem) = Path::new(rel_path).file_stem().and_then(|s| s.to_str()) {
277                    prefixes.push(format!("{package_name}.{file_stem}"));
278                }
279            }
280        }
281    }
282
283    prefixes.sort();
284    prefixes.dedup();
285    prefixes
286}
287
288fn edge_matches_file(edge_to: &str, module_prefixes: &[String]) -> bool {
289    module_prefixes.iter().any(|prefix| {
290        edge_to == *prefix
291            || edge_to.starts_with(&format!("{prefix}::"))
292            || edge_to.starts_with(&format!("{prefix},"))
293    })
294}
295
296fn handle_impact(path: Option<&str>, root: &str) -> String {
297    let target = match path {
298        Some(p) => p,
299        None => return "path is required for 'impact' action".to_string(),
300    };
301
302    let index = match ProjectIndex::load(root) {
303        Some(idx) => idx,
304        None => {
305            return "No graph index found. Run ctx_graph with action='build' first.".to_string()
306        }
307    };
308
309    let rel_target = graph_index::graph_relative_key(target, root);
310
311    let module_prefixes = file_path_to_module_prefixes(&rel_target, root, &index);
312
313    let direct: Vec<&str> = index
314        .edges
315        .iter()
316        .filter(|e| e.kind == "import" && edge_matches_file(&e.to, &module_prefixes))
317        .map(|e| e.from.as_str())
318        .collect();
319
320    let mut all_dependents: Vec<String> = direct.iter().map(|s| s.to_string()).collect();
321    for d in &direct {
322        for dep in index.get_reverse_deps(d, 1) {
323            if !all_dependents.contains(&dep) && dep != rel_target {
324                all_dependents.push(dep);
325            }
326        }
327    }
328
329    if all_dependents.is_empty() {
330        return format!(
331            "No files depend on {}",
332            crate::core::protocol::shorten_path(target)
333        );
334    }
335
336    let mut result = format!(
337        "Impact of {} ({} dependents):\n",
338        crate::core::protocol::shorten_path(target),
339        all_dependents.len()
340    );
341
342    if !direct.is_empty() {
343        result.push_str(&format!("\nDirect ({}):\n", direct.len()));
344        for d in &direct {
345            result.push_str(&format!("  {}\n", crate::core::protocol::shorten_path(d)));
346        }
347    }
348
349    let indirect: Vec<&String> = all_dependents
350        .iter()
351        .filter(|d| !direct.contains(&d.as_str()))
352        .collect();
353    if !indirect.is_empty() {
354        result.push_str(&format!("\nIndirect ({}):\n", indirect.len()));
355        for d in &indirect {
356            result.push_str(&format!("  {}\n", crate::core::protocol::shorten_path(d)));
357        }
358    }
359
360    let tokens = count_tokens(&result);
361    format!("{result}[ctx_graph impact: {tokens} tok]")
362}
363
364fn handle_status(root: &str) -> String {
365    let index = match ProjectIndex::load(root) {
366        Some(idx) => idx,
367        None => return "No graph index. Run ctx_graph action='build' to create one.".to_string(),
368    };
369
370    let mut by_lang: HashMap<&str, usize> = HashMap::new();
371    let mut total_tokens = 0usize;
372    for entry in index.files.values() {
373        *by_lang.entry(&entry.language).or_insert(0) += 1;
374        total_tokens += entry.token_count;
375    }
376
377    let mut langs: Vec<_> = by_lang.iter().collect();
378    langs.sort_by(|a, b| b.1.cmp(a.1));
379    let lang_summary: String = langs
380        .iter()
381        .take(5)
382        .map(|(l, c)| format!("{l}:{c}"))
383        .collect::<Vec<_>>()
384        .join(" ");
385
386    format!(
387        "Graph: {} files, {} symbols, {} edges | {} tok total\nLast scan: {}\nLanguages: {lang_summary}\nStored: {}",
388        index.file_count(),
389        index.symbol_count(),
390        index.edge_count(),
391        total_tokens,
392        index.last_scan,
393        ProjectIndex::index_dir(root)
394            .map(|d| d.to_string_lossy().to_string())
395            .unwrap_or_default()
396    )
397}
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402
403    #[test]
404    fn test_edge_matches_file_crate_prefix() {
405        let prefixes = vec![
406            "lean_ctx::core::cache".to_string(),
407            "crate::core::cache".to_string(),
408            "super::core::cache".to_string(),
409            "core::cache".to_string(),
410        ];
411        assert!(edge_matches_file(
412            "lean_ctx::core::cache::SessionCache",
413            &prefixes
414        ));
415        assert!(edge_matches_file(
416            "crate::core::cache::SessionCache",
417            &prefixes
418        ));
419        assert!(edge_matches_file("crate::core::cache", &prefixes));
420        assert!(!edge_matches_file(
421            "lean_ctx::core::config::Config",
422            &prefixes
423        ));
424        assert!(!edge_matches_file("crate::core::cached_reader", &prefixes));
425    }
426
427    #[test]
428    fn test_file_path_to_module_prefixes_rust() {
429        let index = ProjectIndex::new("/nonexistent");
430        let prefixes = file_path_to_module_prefixes("src/core/cache.rs", "/nonexistent", &index);
431        assert!(prefixes.contains(&"crate::core::cache".to_string()));
432        assert!(prefixes.contains(&"core::cache".to_string()));
433    }
434
435    #[test]
436    fn test_file_path_to_module_prefixes_mod_rs() {
437        let index = ProjectIndex::new("/nonexistent");
438        let prefixes = file_path_to_module_prefixes("src/core/mod.rs", "/nonexistent", &index);
439        assert!(prefixes.contains(&"crate::core".to_string()));
440        assert!(!prefixes.iter().any(|p| p.contains("mod")));
441    }
442}