Skip to main content

lean_ctx/core/
graph_provider.rs

1use std::path::Path;
2use std::sync::atomic::{AtomicBool, Ordering};
3
4use super::graph_index::ProjectIndex;
5use super::property_graph::CodeGraph;
6
7static GRAPH_BUILD_TRIGGERED: AtomicBool = AtomicBool::new(false);
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum GraphProviderSource {
11    PropertyGraph,
12    GraphIndex,
13}
14
15pub enum GraphProvider {
16    PropertyGraph(CodeGraph),
17    GraphIndex(ProjectIndex),
18}
19
20pub struct OpenGraphProvider {
21    pub source: GraphProviderSource,
22    pub provider: GraphProvider,
23}
24
25impl GraphProvider {
26    pub fn node_count(&self) -> Option<usize> {
27        match self {
28            GraphProvider::PropertyGraph(g) => g.node_count().ok(),
29            GraphProvider::GraphIndex(i) => Some(i.file_count()),
30        }
31    }
32
33    pub fn edge_count(&self) -> Option<usize> {
34        match self {
35            GraphProvider::PropertyGraph(g) => g.edge_count().ok(),
36            GraphProvider::GraphIndex(i) => Some(i.edge_count()),
37        }
38    }
39
40    pub fn dependencies(&self, file_path: &str) -> Vec<String> {
41        match self {
42            GraphProvider::PropertyGraph(g) => g.dependencies(file_path).unwrap_or_default(),
43            GraphProvider::GraphIndex(i) => i
44                .edges
45                .iter()
46                .filter(|e| e.kind == "import" && e.from == file_path)
47                .map(|e| e.to.clone())
48                .collect(),
49        }
50    }
51
52    pub fn dependents(&self, file_path: &str) -> Vec<String> {
53        match self {
54            GraphProvider::PropertyGraph(g) => g.dependents(file_path).unwrap_or_default(),
55            GraphProvider::GraphIndex(i) => i
56                .edges
57                .iter()
58                .filter(|e| e.kind == "import" && e.to == file_path)
59                .map(|e| e.from.clone())
60                .collect(),
61        }
62    }
63
64    pub fn related(&self, file_path: &str, depth: usize) -> Vec<String> {
65        match self {
66            GraphProvider::PropertyGraph(g) => g
67                .impact_analysis(file_path, depth)
68                .map(|r| r.affected_files)
69                .unwrap_or_default(),
70            GraphProvider::GraphIndex(i) => i.get_related(file_path, depth),
71        }
72    }
73
74    /// Scored related files using multi-edge weights.
75    /// Falls back to unscored deps/dependents for GraphIndex backend.
76    pub fn related_files_scored(&self, file_path: &str, limit: usize) -> Vec<(String, f64)> {
77        match self {
78            GraphProvider::PropertyGraph(g) => {
79                g.related_files(file_path, limit).unwrap_or_default()
80            }
81            GraphProvider::GraphIndex(_) => {
82                let mut result: Vec<(String, f64)> = Vec::new();
83                for dep in self.dependencies(file_path) {
84                    result.push((dep, 1.0));
85                }
86                for dep in self.dependents(file_path) {
87                    if !result.iter().any(|(p, _)| *p == dep) {
88                        result.push((dep, 0.5));
89                    }
90                }
91                result.truncate(limit);
92                result
93            }
94        }
95    }
96}
97
98pub fn open_best_effort(project_root: &str) -> Option<OpenGraphProvider> {
99    let root = Path::new(project_root);
100
101    let mut pg_provider = None;
102    let mut pg_populated = false;
103    if let Ok(pg) = CodeGraph::open(root) {
104        let nodes = pg.node_count().unwrap_or(0);
105        let edges = pg.edge_count().unwrap_or(0);
106        pg_populated = nodes > 0 && edges > 0;
107        if pg_populated {
108            return Some(OpenGraphProvider {
109                source: GraphProviderSource::PropertyGraph,
110                provider: GraphProvider::PropertyGraph(pg),
111            });
112        }
113        if nodes > 0 {
114            pg_provider = Some(pg);
115        }
116    }
117
118    // Trigger lazy SQLite graph build if PropertyGraph is empty,
119    // even when the JSON graph index provides a fallback.
120    if !pg_populated {
121        trigger_lazy_graph_build(project_root);
122    }
123
124    if let Some(idx) = super::index_orchestrator::try_load_graph_index(project_root) {
125        if !idx.edges.is_empty() {
126            return Some(OpenGraphProvider {
127                source: GraphProviderSource::GraphIndex,
128                provider: GraphProvider::GraphIndex(idx),
129            });
130        }
131        if !idx.files.is_empty() {
132            return Some(OpenGraphProvider {
133                source: GraphProviderSource::GraphIndex,
134                provider: GraphProvider::GraphIndex(idx),
135            });
136        }
137    }
138
139    if let Some(pg) = pg_provider {
140        return Some(OpenGraphProvider {
141            source: GraphProviderSource::PropertyGraph,
142            provider: GraphProvider::PropertyGraph(pg),
143        });
144    }
145
146    None
147}
148
149/// Triggers a background graph build once per process when the graph is empty.
150fn trigger_lazy_graph_build(project_root: &str) {
151    if GRAPH_BUILD_TRIGGERED.swap(true, Ordering::SeqCst) {
152        return;
153    }
154    let root = Path::new(project_root);
155    let is_project = root.is_dir()
156        && (root.join(".git").exists()
157            || root.join("Cargo.toml").exists()
158            || root.join("package.json").exists()
159            || root.join("go.mod").exists());
160    if !is_project {
161        return;
162    }
163    let root_owned = project_root.to_string();
164    std::thread::spawn(move || {
165        // TODO(arch): calls into tools::ctx_impact -- should use a trait/callback
166        // to decouple core from tools layer.
167        let _ = crate::tools::ctx_impact::handle("build", None, &root_owned, None, None);
168    });
169}
170
171pub fn open_or_build(project_root: &str) -> Option<OpenGraphProvider> {
172    if let Some(p) = open_best_effort(project_root) {
173        return Some(p);
174    }
175    let idx = super::graph_index::load_or_build(project_root);
176    if idx.files.is_empty() {
177        return None;
178    }
179    Some(OpenGraphProvider {
180        source: GraphProviderSource::GraphIndex,
181        provider: GraphProvider::GraphIndex(idx),
182    })
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn best_effort_prefers_graph_index_when_property_graph_empty() {
191        let _lock = crate::core::data_dir::test_env_lock();
192        let tmp = tempfile::tempdir().expect("tempdir");
193        let data = tmp.path().join("data");
194        std::fs::create_dir_all(&data).expect("mkdir data");
195        std::env::set_var("LEAN_CTX_DATA_DIR", data.to_string_lossy().to_string());
196
197        let project_root = tmp.path().join("proj");
198        std::fs::create_dir_all(&project_root).expect("mkdir proj");
199        let root = project_root.to_string_lossy().to_string();
200
201        let mut idx = ProjectIndex::new(&root);
202        idx.files.insert(
203            "src/main.rs".to_string(),
204            super::super::graph_index::FileEntry {
205                path: "src/main.rs".to_string(),
206                hash: "h".to_string(),
207                language: "rs".to_string(),
208                line_count: 1,
209                token_count: 1,
210                exports: vec![],
211                summary: String::new(),
212            },
213        );
214        idx.save().expect("save index");
215
216        let open = open_best_effort(&root).expect("open");
217        assert_eq!(open.source, GraphProviderSource::GraphIndex);
218
219        std::env::remove_var("LEAN_CTX_DATA_DIR");
220    }
221
222    #[test]
223    fn best_effort_none_when_no_graphs() {
224        let _lock = crate::core::data_dir::test_env_lock();
225        let tmp = tempfile::tempdir().expect("tempdir");
226        let data = tmp.path().join("data");
227        std::fs::create_dir_all(&data).expect("mkdir data");
228        std::env::set_var("LEAN_CTX_DATA_DIR", data.to_string_lossy().to_string());
229
230        let project_root = tmp.path().join("proj");
231        std::fs::create_dir_all(&project_root).expect("mkdir proj");
232        let root = project_root.to_string_lossy().to_string();
233
234        let open = open_best_effort(&root);
235        assert!(open.is_none());
236
237        std::env::remove_var("LEAN_CTX_DATA_DIR");
238    }
239}