1use std::cell::{OnceCell, RefCell};
34use std::collections::{BTreeSet, HashMap, HashSet};
35use std::path::Path;
36
37use rusqlite::Connection;
38
39use super::{CodeGraph, Node};
40use crate::core::bm25_index::source_gate::SourceGate;
41
42const GRAPH_TOOL: &str = "ctx_graph";
44
45const PROTECTED_FILES: usize = crate::core::bm25_index::source_gate::MAX_SOURCES;
49
50const MAX_FILES: usize = 65_536;
53
54pub(crate) struct GraphGate {
56 project_root: String,
57 protected: OnceCell<bool>,
58 present: RefCell<HashMap<String, bool>>,
60 sources: RefCell<SourceGate>,
61}
62
63impl GraphGate {
64 pub(crate) fn for_root(project_root: &str) -> Self {
65 Self {
66 project_root: project_root.to_owned(),
67 protected: OnceCell::new(),
68 present: RefCell::new(HashMap::new()),
69 sources: RefCell::new(SourceGate::new(Path::new(project_root), GRAPH_TOOL)),
70 }
71 }
72
73 pub(crate) fn file(&self, path: &str) -> bool {
75 self.reference(path, "")
76 }
77
78 pub(crate) fn symbol(&self, path: &str, name: &str) -> bool {
80 self.reference(path, identifier(name))
81 }
82
83 fn protected(&self) -> bool {
85 *self
86 .protected
87 .get_or_init(crate::core::context_admission::stored::source_authority_required)
88 }
89
90 fn originless(&self, text: &str) -> bool {
97 !self.protected()
98 && crate::core::context_admission::stored::serve_stored(text, None, None)
99 .is_ok_and(|served| served == text)
100 }
101
102 pub(crate) fn ranking_files(
114 &self,
115 mut files: Vec<String>,
116 edges: &[(String, String)],
117 seeds: &[String],
118 ) -> HashSet<String> {
119 files.sort();
120 files.dedup();
121 if !self.protected() || files.len() <= PROTECTED_FILES {
122 return files.into_iter().filter(|file| self.file(file)).collect();
123 }
124 let known: HashSet<&str> = files.iter().map(String::as_str).collect();
125 let mut asked: HashSet<&str> = HashSet::new();
126 let mut served: HashSet<String> = HashSet::new();
127 let mut frontier: BTreeSet<&str> = seeds
128 .iter()
129 .map(String::as_str)
130 .filter(|seed| known.contains(*seed))
131 .collect();
132 while !frontier.is_empty() && asked.len() < PROTECTED_FILES {
133 let mut next: BTreeSet<&str> = BTreeSet::new();
134 for file in frontier {
135 if asked.len() >= PROTECTED_FILES {
136 break;
137 }
138 if !asked.insert(file) || !self.file(file) {
139 continue;
140 }
141 served.insert(file.to_owned());
142 for (from, to) in edges {
143 if from == file {
144 next.insert(to.as_str());
145 } else if to == file {
146 next.insert(from.as_str());
147 }
148 }
149 }
150 next.retain(|file| known.contains(*file) && !asked.contains(*file));
151 frontier = next;
152 }
153 served
154 }
155
156 fn reference(&self, path: &str, needle: &str) -> bool {
157 let protected = self.protected();
158 if path.contains("://") {
159 return !protected;
160 }
161 self.present(path) && (!protected || self.sources.borrow_mut().serves(path, needle))
162 }
163
164 fn present(&self, path: &str) -> bool {
165 let known = self.present.borrow().get(path).copied();
166 if let Some(known) = known {
167 return known;
168 }
169 if self.present.borrow().len() >= MAX_FILES {
170 return false;
171 }
172 let root = Path::new(&self.project_root);
173 let present = crate::core::pathjail::jail_path(&root.join(path), root)
174 .is_ok_and(|jailed| std::fs::symlink_metadata(jailed).is_ok_and(|meta| meta.is_file()));
175 self.present.borrow_mut().insert(path.to_owned(), present);
176 present
177 }
178}
179
180fn identifier(name: &str) -> &str {
184 name.rsplit(|c: char| !(c.is_alphanumeric() || c == '_'))
185 .find(|part| !part.is_empty())
186 .unwrap_or("")
187}
188
189impl CodeGraph {
190 #[must_use]
193 pub(crate) fn gated_for(mut self, project_root: &str) -> Self {
194 self.gate = Some(GraphGate::for_root(project_root));
195 self
196 }
197
198 pub fn serves_file(&self, path: &str) -> bool {
202 self.gate.as_ref().is_none_or(|gate| gate.file(path))
203 }
204
205 pub fn serves_symbol(&self, path: &str, name: &str) -> bool {
207 self.gate
208 .as_ref()
209 .is_none_or(|gate| gate.symbol(path, name))
210 }
211
212 pub fn serves_node(&self, node: &Node) -> bool {
215 self.serves_node_parts(
216 node.kind.as_str(),
217 &node.file_path,
218 &node.name,
219 node.metadata.as_deref(),
220 )
221 }
222
223 pub fn serves_node_parts(
235 &self,
236 kind: &str,
237 file_path: &str,
238 name: &str,
239 metadata: Option<&str>,
240 ) -> bool {
241 let Some(gate) = self.gate.as_ref() else {
242 return true;
243 };
244 if file_path.is_empty() {
245 return gate.originless(name) && metadata.is_none_or(|text| gate.originless(text));
246 }
247 match kind {
248 "symbol" | "test" if name != file_path => gate.symbol(file_path, name),
249 _ => gate.file(file_path),
250 }
251 }
252
253 pub(super) fn served_neighbours(&self, anchor: &str, neighbours: Vec<String>) -> Vec<String> {
256 if !self.serves_file(anchor) {
257 return Vec::new();
258 }
259 neighbours
260 .into_iter()
261 .filter(|path| self.serves_file(path))
262 .collect()
263 }
264}
265
266pub trait GraphSource {
274 fn store(&self) -> &Connection;
276
277 fn checks_sources(&self) -> bool;
281
282 fn file_served(&self, path: &str) -> bool;
284
285 fn symbol_served(&self, path: &str, name: &str) -> bool;
287
288 fn rankable_files(
291 &self,
292 files: Vec<String>,
293 edges: &[(String, String)],
294 seeds: &[String],
295 ) -> HashSet<String>;
296}
297
298impl<T: GraphSource + ?Sized> GraphSource for &T {
301 fn store(&self) -> &Connection {
302 (**self).store()
303 }
304
305 fn checks_sources(&self) -> bool {
306 (**self).checks_sources()
307 }
308
309 fn file_served(&self, path: &str) -> bool {
310 (**self).file_served(path)
311 }
312
313 fn symbol_served(&self, path: &str, name: &str) -> bool {
314 (**self).symbol_served(path, name)
315 }
316
317 fn rankable_files(
318 &self,
319 files: Vec<String>,
320 edges: &[(String, String)],
321 seeds: &[String],
322 ) -> HashSet<String> {
323 (**self).rankable_files(files, edges, seeds)
324 }
325}
326
327impl GraphSource for Connection {
328 fn store(&self) -> &Connection {
329 self
330 }
331
332 fn checks_sources(&self) -> bool {
333 false
334 }
335
336 fn file_served(&self, _path: &str) -> bool {
337 true
338 }
339
340 fn symbol_served(&self, _path: &str, _name: &str) -> bool {
341 true
342 }
343
344 fn rankable_files(
345 &self,
346 files: Vec<String>,
347 _edges: &[(String, String)],
348 _seeds: &[String],
349 ) -> HashSet<String> {
350 files.into_iter().collect()
351 }
352}
353
354impl GraphSource for CodeGraph {
355 fn store(&self) -> &Connection {
356 self.connection()
357 }
358
359 fn checks_sources(&self) -> bool {
360 self.gate.is_some()
361 }
362
363 fn file_served(&self, path: &str) -> bool {
364 self.serves_file(path)
365 }
366
367 fn symbol_served(&self, path: &str, name: &str) -> bool {
368 self.serves_symbol(path, name)
369 }
370
371 fn rankable_files(
372 &self,
373 files: Vec<String>,
374 edges: &[(String, String)],
375 seeds: &[String],
376 ) -> HashSet<String> {
377 match self.gate.as_ref() {
378 Some(gate) => gate.ranking_files(files, edges, seeds),
379 None => files.into_iter().collect(),
380 }
381 }
382}
383
384#[cfg(test)]
385mod tests {
386 use crate::core::graph_index::{FileEntry, ProjectIndex, SymbolEntry};
387 use crate::core::graph_provider::GraphProvider;
388 use crate::core::policy::runtime::TestPolicyOverride;
389
390 #[test]
393 fn a_graph_reference_is_withheld_once_its_source_file_is_gone() {
394 let _data = crate::core::data_dir::isolated_data_dir();
395 let _community = TestPolicyOverride::set(None);
396 let project = tempfile::tempdir().expect("project");
397 let root = project.path().to_string_lossy().into_owned();
398 std::fs::write(project.path().join("auth.rs"), "pub fn verify() {}\n").expect("source");
399
400 let mut index = ProjectIndex::new(&root);
401 index.files.insert(
402 "auth.rs".to_owned(),
403 FileEntry {
404 path: "auth.rs".to_owned(),
405 hash: String::new(),
406 language: "rust".to_owned(),
407 line_count: 1,
408 token_count: 4,
409 exports: vec!["verify".to_owned()],
410 summary: "pub fn verify() {}".to_owned(),
411 },
412 );
413 index.symbols.insert(
414 "auth.rs::verify".to_owned(),
415 SymbolEntry {
416 file: "auth.rs".to_owned(),
417 name: "verify".to_owned(),
418 kind: "fn".to_owned(),
419 start_line: 1,
420 end_line: 1,
421 is_exported: true,
422 },
423 );
424 let provider = GraphProvider::GraphIndex(index);
425
426 assert_eq!(provider.file_paths(), vec!["auth.rs".to_owned()]);
427 assert_eq!(provider.all_symbols().len(), 1);
428 assert!(provider.get_file_entry("auth.rs").is_some());
429
430 std::fs::remove_file(project.path().join("auth.rs")).expect("remove");
431 assert!(provider.file_paths().is_empty());
432 assert!(provider.all_symbols().is_empty());
433 assert!(provider.get_symbol("auth.rs::verify").is_none());
434 assert!(provider.get_file_entry("auth.rs").is_none());
435 assert!(provider.materialize_project_index(&root).files.is_empty());
436 }
437
438 #[test]
443 fn a_ranking_under_protection_uses_only_files_decided_from_its_seeds() {
444 use super::{GraphGate, PROTECTED_FILES};
445
446 let _data = crate::core::data_dir::isolated_data_dir();
447 let pack = crate::core::policy::parse(
448 "name = \"graph-ranking-test\"\nversion = \"1.0.0\"\ndescription = \"test\"\n\
449 [redaction]\ncustomer = 'K-[0-9]{6}'",
450 )
451 .expect("valid policy");
452 let _protected = TestPolicyOverride::set(Some(
453 crate::core::policy::resolve(&pack).expect("resolved policy"),
454 ));
455 let project = tempfile::tempdir().expect("project");
456 let root = project.path().to_string_lossy().into_owned();
457 let files: Vec<String> = (0..PROTECTED_FILES + 2)
459 .map(|n| format!("f{n:03}.rs"))
460 .collect();
461 for file in &files {
462 std::fs::write(project.path().join(file), "pub fn item() {}\n").expect("source");
463 }
464 let edges: Vec<(String, String)> = files
465 .windows(2)
466 .map(|pair| (pair[0].clone(), pair[1].clone()))
467 .collect();
468
469 let gate = GraphGate::for_root(&root);
470 let local = gate.ranking_files(files.clone(), &edges, &files[..1]);
471 assert!(local.contains(&files[0]), "the seed is decided first");
472 assert!(local.len() <= PROTECTED_FILES);
473 assert!(
474 !local.contains(&files[PROTECTED_FILES + 1]),
475 "beyond the bound"
476 );
477 assert!(local.iter().all(|file| gate.file(file)));
479 assert!(!gate.file(&files[PROTECTED_FILES + 1]));
481
482 let unseeded = GraphGate::for_root(&root).ranking_files(files, &edges, &[]);
483 assert!(unseeded.is_empty());
484 }
485
486 #[test]
491 fn a_traversal_does_not_pass_through_a_withdrawn_file() {
492 use crate::core::property_graph::{CodeGraph, Edge, EdgeKind, Node};
493
494 let _data = crate::core::data_dir::isolated_data_dir();
495 let _community = TestPolicyOverride::set(None);
496 let project = tempfile::tempdir().expect("project");
497 let root = project.path().to_string_lossy().into_owned();
498 for file in ["a.rs", "w.rs", "b.rs", "n.rs"] {
499 std::fs::write(project.path().join(file), "pub fn item() {}\n").expect("source");
500 }
501
502 let graph = CodeGraph::open(&root).expect("graph");
503 let node = |file: &str| graph.upsert_node(&Node::file(file)).expect("node");
504 let (a, w, b, n) = (node("a.rs"), node("w.rs"), node("b.rs"), node("n.rs"));
505 for (from, to) in [(a, w), (w, b), (a, n)] {
506 graph
507 .upsert_edge(&Edge::new(from, to, EdgeKind::Imports))
508 .expect("edge");
509 }
510
511 let chain = graph.dependency_chain("a.rs", "b.rs").expect("chain");
512 assert_eq!(
513 chain.map(|chain| chain.path),
514 Some(vec![
515 "a.rs".to_owned(),
516 "w.rs".to_owned(),
517 "b.rs".to_owned()
518 ])
519 );
520 let impact = graph.impact_analysis("b.rs", 5).expect("impact");
521 assert_eq!(impact.affected_files, ["a.rs", "w.rs"]);
522 assert_eq!(impact.edges_traversed, 2);
523
524 std::fs::remove_file(project.path().join("w.rs")).expect("withdraw");
525 let graph = CodeGraph::open(&root).expect("graph");
527 assert!(
528 graph
529 .dependency_chain("a.rs", "b.rs")
530 .expect("chain")
531 .is_none()
532 );
533 let impact = graph.impact_analysis("b.rs", 5).expect("impact");
534 assert!(impact.affected_files.is_empty(), "{impact:?}");
535 assert_eq!(impact.edges_traversed, 0);
536 assert_eq!(impact.max_depth_reached, 0);
537 let near = graph.dependency_chain("a.rs", "n.rs").expect("chain");
538 assert_eq!(near.map(|chain| chain.depth), Some(1));
539 let related: Vec<String> = graph
540 .related_files("a.rs", 10)
541 .expect("related")
542 .into_iter()
543 .map(|(path, _)| path)
544 .collect();
545 assert_eq!(related, ["n.rs"]);
546 }
547}