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