diff --git a/code2llm/analysis/native_bridge.py b/code2llm/analysis/native_bridge.py index e1254ee..bcc9fc5 100644 --- a/code2llm/analysis/native_bridge.py +++ b/code2llm/analysis/native_bridge.py @@ -242,3 +242,16 @@ def native_format_toon_header( return None +def native_resolve_call_graph( + functions_with_calls: List[Tuple[str, List[str]]], +) -> Optional[Tuple[Dict[str, List[str]], Dict[str, List[str]], List[str]]]: + """Resolve call graph edges and find entry points natively.""" + if _HAS_NATIVE_RUST and hasattr(code2llm_rust, "resolve_call_graph"): + try: + return code2llm_rust.resolve_call_graph(functions_with_calls) + except Exception as e: + logger.debug("code2llm_rust.resolve_call_graph failed: %s", e) + return None + + + diff --git a/code2llm/core/analyzer.py b/code2llm/core/analyzer.py index 96f26f5..95b6fee 100644 --- a/code2llm/core/analyzer.py +++ b/code2llm/core/analyzer.py @@ -65,7 +65,10 @@ from .file_analyzer import FileAnalyzer, _analyze_single_file from .persistent_cache import PersistentCache from .refactoring import RefactoringAnalyzer -from code2llm.analysis.native_bridge import native_walk_project_files +from code2llm.analysis.native_bridge import ( + native_walk_project_files, + native_resolve_call_graph, +) logger = logging.getLogger(__name__) @@ -530,14 +533,30 @@ def _build_call_graph(self, result: AnalysisResult) -> None: flush=True, ) - # Build lookup maps for O(1) resolution - simple_to_full = self._build_simple_name_map(result) - - # Map calls between functions - self._collect_call_edges(result, simple_to_full) - - # Find entry points - self._find_entry_points(result) + # Fast path: native Rust resolution + func_calls = [ + (func_name, list(getattr(func, "calls", []))) + for func_name, func in result.functions.items() + ] + native_res = native_resolve_call_graph(func_calls) + if native_res is not None: + resolved_calls, called_by, entry_points = native_res + for func_name, calls in resolved_calls.items(): + if func_name in result.functions: + result.functions[func_name].calls = calls + for func_name, callers in called_by.items(): + if func_name in result.functions: + result.functions[func_name].called_by = callers + result.entry_points = entry_points + else: + # Build lookup maps for O(1) resolution + simple_to_full = self._build_simple_name_map(result) + + # Map calls between functions + self._collect_call_edges(result, simple_to_full) + + # Find entry points + self._find_entry_points(result) if self.config.verbose: print( diff --git a/docs/analysis/rust-optimization-breakdown.md b/docs/analysis/rust-optimization-breakdown.md index f3cdb4a..2990e52 100644 --- a/docs/analysis/rust-optimization-breakdown.md +++ b/docs/analysis/rust-optimization-breakdown.md @@ -61,6 +61,10 @@ Through systematic profiling and code analysis of `code2llm` during multi-langua - **Problem**: Generating and formatting formatted header rows and metrics summaries in Python across thousands of functions and modules causes heavy string/list allocations. - **Native Rust Solution**: Native string buffer formatter in `toon.rs` emitting pre-formatted diagnostic lines. +### Bottleneck K: Cross-Module Call Graph Edge Resolution & Entry Points +- **Problem**: In `ProjectAnalyzer._build_call_graph`, iterating through every function's call list in Python, executing string `rsplit` operations to find module matches, and building `called_by` with quadratic member checks adds significant runtime on large repositories. +- **Native Rust Solution**: `calls::resolve_call_graph` resolving simple-to-full name candidate lookups with same-module affinity in a single linear pass in native Rust, returning resolved forward calls, backward caller maps, and root entry points. + ## 3. Architecture & Refactoring Strategy To prepare `code2llm` for clean separation and extraction into the `code2llm-rust` package, the following refactoring was implemented: @@ -71,6 +75,7 @@ To prepare `code2llm` for clean separation and extraction into the `code2llm-rus - Provides transparent fallback to pure Python when `code2llm-rust` is absent, guaranteeing 100% backward compatibility and portability. 2. **`code2llm.core.analyzer`**: - Refactored `_collect_files` to delegate to `native_bridge.native_walk_project_files` with Python fallback. + - Refactored `_build_call_graph` to delegate to `native_bridge.native_resolve_call_graph` with Python fallback. 3. **`code2llm.analysis.complexity`**: - Refactored `estimate_function_complexity`, `compute_cyclomatic_complexity`, and `_rust_batch_complexity` to delegate to `native_bridge`. 4. **`code2llm.analysis.call_graph_engine`**: @@ -95,7 +100,7 @@ To prepare `code2llm` for clean separation and extraction into the `code2llm-rus - **PyO3 Bindings**: CPython extension module `code2llm_rust` - **Engine Components**: - `src/complexity.rs`: Function boundary extraction and McCabe complexity. - - `src/calls.rs`: Token-level call extraction. + - `src/calls.rs`: Token-level call extraction and cross-module call graph resolution. - `src/centrality.rs`: Parallel Brandes' betweenness centrality. - `src/cycles.rs`: Linear-time Tarjan SCC cycle detection. - `src/reachability.rs`: Fast in-memory graph reachability. @@ -110,3 +115,4 @@ To prepare `code2llm` for clean separation and extraction into the `code2llm-rus - Python optional dependency: `code2llm[native]` in `pyproject.toml`. + diff --git a/packages/code2llm-rust/src/calls.rs b/packages/code2llm-rust/src/calls.rs index 9c12eda..a5821d1 100644 --- a/packages/code2llm-rust/src/calls.rs +++ b/packages/code2llm-rust/src/calls.rs @@ -72,3 +72,113 @@ pub fn batch_calls( }) .collect() } + +use std::collections::{HashMap, HashSet}; + +/// Result of full call graph resolution. +pub struct CallGraphResolution { + pub resolved_calls: HashMap>, + pub called_by: HashMap>, + pub entry_points: Vec, +} + +/// Fast native call graph edge resolution matching ProjectAnalyzer._build_call_graph. +pub fn resolve_call_graph( + functions_with_calls: &[(String, Vec)], +) -> CallGraphResolution { + let known_functions: HashSet<&str> = functions_with_calls + .iter() + .map(|(name, _)| name.as_str()) + .collect(); + + let mut simple_to_full: HashMap<&str, Vec<&str>> = HashMap::new(); + for (name, _) in functions_with_calls { + let simple = match name.rfind('.') { + Some(idx) => &name[idx + 1..], + None => name.as_str(), + }; + simple_to_full.entry(simple).or_default().push(name.as_str()); + } + + let mut resolved_calls: HashMap> = HashMap::new(); + let mut called_by: HashMap> = HashMap::new(); + + for (name, _) in functions_with_calls { + called_by.entry(name.clone()).or_default(); + } + + for (func_name, calls) in functions_with_calls { + let func_module = match func_name.rfind('.') { + Some(idx) => &func_name[..idx], + None => "", + }; + + let mut res_calls = Vec::with_capacity(calls.len()); + for called in calls { + let called_str = called.as_str(); + let resolved_opt = if known_functions.contains(called_str) { + Some(called_str) + } else if let Some(candidates) = simple_to_full.get(called_str) { + let mut best = candidates[0]; + for &cand in candidates { + let cand_module = match cand.rfind('.') { + Some(idx) => &cand[..idx], + None => "", + }; + if cand_module == func_module { + best = cand; + break; + } + } + Some(best) + } else { + None + }; + + if let Some(resolved) = resolved_opt { + res_calls.push(resolved.to_string()); + called_by + .entry(resolved.to_string()) + .or_default() + .push(func_name.clone()); + } + } + resolved_calls.insert(func_name.clone(), res_calls); + } + + let mut entry_points: Vec = Vec::new(); + for (name, callers) in &called_by { + if callers.is_empty() { + entry_points.push(name.clone()); + } + } + entry_points.sort(); + + CallGraphResolution { + resolved_calls, + called_by, + entry_points, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_resolve_call_graph() { + let input = vec![ + ("app.main".to_string(), vec!["helper".to_string()]), + ("app.helper".to_string(), vec!["db.query".to_string()]), + ("db.query".to_string(), vec![]), + ]; + + let res = resolve_call_graph(&input); + assert_eq!(res.resolved_calls["app.main"], vec!["app.helper"]); + assert_eq!(res.resolved_calls["app.helper"], vec!["db.query"]); + assert_eq!(res.called_by["app.helper"], vec!["app.main"]); + assert_eq!(res.called_by["db.query"], vec!["app.helper"]); + assert_eq!(res.entry_points, vec!["app.main"]); + } +} + diff --git a/packages/code2llm-rust/src/lib.rs b/packages/code2llm-rust/src/lib.rs index b2d07d2..5426c69 100644 --- a/packages/code2llm-rust/src/lib.rs +++ b/packages/code2llm-rust/src/lib.rs @@ -222,6 +222,23 @@ fn format_toon_header( ) } +/// Resolve call graph edges and find entry points. +#[pyfunction] +#[pyo3(signature = (functions_with_calls))] +fn resolve_call_graph( + py: Python<'_>, + functions_with_calls: Vec<(String, Vec)>, +) -> ( + HashMap>, + HashMap>, + Vec, +) { + py.allow_threads(|| { + let res = calls::resolve_call_graph(&functions_with_calls); + (res.resolved_calls, res.called_by, res.entry_points) + }) +} + /// Module initialization #[pymodule] fn code2llm_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -240,8 +257,10 @@ fn code2llm_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(compute_file_content_hash, m)?)?; m.add_function(wrap_pyfunction!(check_changed_files, m)?)?; m.add_function(wrap_pyfunction!(format_toon_header, m)?)?; + m.add_function(wrap_pyfunction!(resolve_call_graph, m)?)?; m.add("__version__", env!("CARGO_PKG_VERSION"))?; Ok(()) } + diff --git a/tests/test_native_rust_bridge.py b/tests/test_native_rust_bridge.py index 90cfc97..dc8994d 100644 --- a/tests/test_native_rust_bridge.py +++ b/tests/test_native_rust_bridge.py @@ -298,3 +298,29 @@ def test_native_format_toon_header(): assert "critical:5/100" in lines[1] +def test_native_resolve_call_graph(): + """Verify native call graph resolution and entry points finding.""" + from code2llm.analysis.native_bridge import native_resolve_call_graph + + func_calls = [ + ("app.cli.run", ["parse_args", "app.core.execute"]), + ("app.cli.parse_args", []), + ("app.core.execute", ["db.connect"]), + ("db.connect", []), + ] + + res = native_resolve_call_graph(func_calls) + assert res is not None + resolved_calls, called_by, entry_points = res + + # parse_args should be resolved to app.cli.parse_args + assert resolved_calls["app.cli.run"] == ["app.cli.parse_args", "app.core.execute"] + assert "app.cli.run" in called_by["app.cli.parse_args"] + assert "app.cli.run" in called_by["app.core.execute"] + assert "app.core.execute" in called_by["db.connect"] + + # Entry point is only app.cli.run + assert entry_points == ["app.cli.run"] + + +