use std::cell::{OnceCell, RefCell};
use std::collections::{BTreeSet, HashMap, HashSet};
use std::path::Path;
use rusqlite::Connection;
use super::{CodeGraph, Node};
use crate::core::bm25_index::source_gate::SourceGate;
const GRAPH_TOOL: &str = "ctx_graph";
const PROTECTED_FILES: usize = crate::core::bm25_index::source_gate::MAX_SOURCES;
const MAX_FILES: usize = 65_536;
pub(crate) struct GraphGate {
project_root: String,
protected: OnceCell<bool>,
present: RefCell<HashMap<String, bool>>,
sources: RefCell<SourceGate>,
}
impl GraphGate {
pub(crate) fn for_root(project_root: &str) -> Self {
Self {
project_root: project_root.to_owned(),
protected: OnceCell::new(),
present: RefCell::new(HashMap::new()),
sources: RefCell::new(SourceGate::new(Path::new(project_root), GRAPH_TOOL)),
}
}
pub(crate) fn file(&self, path: &str) -> bool {
self.reference(path, "")
}
pub(crate) fn symbol(&self, path: &str, name: &str) -> bool {
self.reference(path, identifier(name))
}
fn protected(&self) -> bool {
*self
.protected
.get_or_init(crate::core::context_admission::stored::source_authority_required)
}
fn originless(&self, text: &str) -> bool {
!self.protected()
&& crate::core::context_admission::stored::serve_stored(text, None, None)
.is_ok_and(|served| served == text)
}
pub(crate) fn ranking_files(
&self,
mut files: Vec<String>,
edges: &[(String, String)],
seeds: &[String],
) -> HashSet<String> {
files.sort();
files.dedup();
if !self.protected() || files.len() <= PROTECTED_FILES {
return files.into_iter().filter(|file| self.file(file)).collect();
}
let known: HashSet<&str> = files.iter().map(String::as_str).collect();
let mut asked: HashSet<&str> = HashSet::new();
let mut served: HashSet<String> = HashSet::new();
let mut frontier: BTreeSet<&str> = seeds
.iter()
.map(String::as_str)
.filter(|seed| known.contains(*seed))
.collect();
while !frontier.is_empty() && asked.len() < PROTECTED_FILES {
let mut next: BTreeSet<&str> = BTreeSet::new();
for file in frontier {
if asked.len() >= PROTECTED_FILES {
break;
}
if !asked.insert(file) || !self.file(file) {
continue;
}
served.insert(file.to_owned());
for (from, to) in edges {
if from == file {
next.insert(to.as_str());
} else if to == file {
next.insert(from.as_str());
}
}
}
next.retain(|file| known.contains(*file) && !asked.contains(*file));
frontier = next;
}
served
}
fn reference(&self, path: &str, needle: &str) -> bool {
let protected = self.protected();
if path.contains("://") {
return !protected;
}
self.present(path) && (!protected || self.sources.borrow_mut().serves(path, needle))
}
fn present(&self, path: &str) -> bool {
let known = self.present.borrow().get(path).copied();
if let Some(known) = known {
return known;
}
if self.present.borrow().len() >= MAX_FILES {
return false;
}
let root = Path::new(&self.project_root);
let present = crate::core::pathjail::jail_path(&root.join(path), root)
.is_ok_and(|jailed| std::fs::symlink_metadata(jailed).is_ok_and(|meta| meta.is_file()));
self.present.borrow_mut().insert(path.to_owned(), present);
present
}
}
fn identifier(name: &str) -> &str {
name.rsplit(|c: char| !(c.is_alphanumeric() || c == '_'))
.find(|part| !part.is_empty())
.unwrap_or("")
}
impl CodeGraph {
#[must_use]
pub(crate) fn gated_for(mut self, project_root: &str) -> Self {
self.gate = Some(GraphGate::for_root(project_root));
self
}
pub fn serves_file(&self, path: &str) -> bool {
self.gate.as_ref().is_none_or(|gate| gate.file(path))
}
pub fn serves_symbol(&self, path: &str, name: &str) -> bool {
self.gate
.as_ref()
.is_none_or(|gate| gate.symbol(path, name))
}
pub fn serves_node(&self, node: &Node) -> bool {
self.serves_node_parts(
node.kind.as_str(),
&node.file_path,
&node.name,
node.metadata.as_deref(),
)
}
pub fn serves_node_parts(
&self,
kind: &str,
file_path: &str,
name: &str,
metadata: Option<&str>,
) -> bool {
let Some(gate) = self.gate.as_ref() else {
return true;
};
if file_path.is_empty() {
return gate.originless(name) && metadata.is_none_or(|text| gate.originless(text));
}
match kind {
"symbol" | "test" if name != file_path => gate.symbol(file_path, name),
_ => gate.file(file_path),
}
}
pub(super) fn served_neighbours(&self, anchor: &str, neighbours: Vec<String>) -> Vec<String> {
if !self.serves_file(anchor) {
return Vec::new();
}
neighbours
.into_iter()
.filter(|path| self.serves_file(path))
.collect()
}
}
pub trait GraphSource {
fn store(&self) -> &Connection;
fn checks_sources(&self) -> bool;
fn file_served(&self, path: &str) -> bool;
fn symbol_served(&self, path: &str, name: &str) -> bool;
fn rankable_files(
&self,
files: Vec<String>,
edges: &[(String, String)],
seeds: &[String],
) -> HashSet<String>;
}
impl<T: GraphSource + ?Sized> GraphSource for &T {
fn store(&self) -> &Connection {
(**self).store()
}
fn checks_sources(&self) -> bool {
(**self).checks_sources()
}
fn file_served(&self, path: &str) -> bool {
(**self).file_served(path)
}
fn symbol_served(&self, path: &str, name: &str) -> bool {
(**self).symbol_served(path, name)
}
fn rankable_files(
&self,
files: Vec<String>,
edges: &[(String, String)],
seeds: &[String],
) -> HashSet<String> {
(**self).rankable_files(files, edges, seeds)
}
}
impl GraphSource for Connection {
fn store(&self) -> &Connection {
self
}
fn checks_sources(&self) -> bool {
false
}
fn file_served(&self, _path: &str) -> bool {
true
}
fn symbol_served(&self, _path: &str, _name: &str) -> bool {
true
}
fn rankable_files(
&self,
files: Vec<String>,
_edges: &[(String, String)],
_seeds: &[String],
) -> HashSet<String> {
files.into_iter().collect()
}
}
impl GraphSource for CodeGraph {
fn store(&self) -> &Connection {
self.connection()
}
fn checks_sources(&self) -> bool {
self.gate.is_some()
}
fn file_served(&self, path: &str) -> bool {
self.serves_file(path)
}
fn symbol_served(&self, path: &str, name: &str) -> bool {
self.serves_symbol(path, name)
}
fn rankable_files(
&self,
files: Vec<String>,
edges: &[(String, String)],
seeds: &[String],
) -> HashSet<String> {
match self.gate.as_ref() {
Some(gate) => gate.ranking_files(files, edges, seeds),
None => files.into_iter().collect(),
}
}
}
#[cfg(test)]
mod tests {
use crate::core::graph_index::{FileEntry, ProjectIndex, SymbolEntry};
use crate::core::graph_provider::GraphProvider;
use crate::core::policy::runtime::TestPolicyOverride;
#[test]
fn a_graph_reference_is_withheld_once_its_source_file_is_gone() {
let _data = crate::core::data_dir::isolated_data_dir();
let _community = TestPolicyOverride::set(None);
let project = tempfile::tempdir().expect("project");
let root = project.path().to_string_lossy().into_owned();
std::fs::write(project.path().join("auth.rs"), "pub fn verify() {}\n").expect("source");
let mut index = ProjectIndex::new(&root);
index.files.insert(
"auth.rs".to_owned(),
FileEntry {
path: "auth.rs".to_owned(),
hash: String::new(),
language: "rust".to_owned(),
line_count: 1,
token_count: 4,
exports: vec!["verify".to_owned()],
summary: "pub fn verify() {}".to_owned(),
},
);
index.symbols.insert(
"auth.rs::verify".to_owned(),
SymbolEntry {
file: "auth.rs".to_owned(),
name: "verify".to_owned(),
kind: "fn".to_owned(),
start_line: 1,
end_line: 1,
is_exported: true,
},
);
let provider = GraphProvider::GraphIndex(index);
assert_eq!(provider.file_paths(), vec!["auth.rs".to_owned()]);
assert_eq!(provider.all_symbols().len(), 1);
assert!(provider.get_file_entry("auth.rs").is_some());
std::fs::remove_file(project.path().join("auth.rs")).expect("remove");
assert!(provider.file_paths().is_empty());
assert!(provider.all_symbols().is_empty());
assert!(provider.get_symbol("auth.rs::verify").is_none());
assert!(provider.get_file_entry("auth.rs").is_none());
assert!(provider.materialize_project_index(&root).files.is_empty());
}
#[test]
fn a_ranking_under_protection_uses_only_files_decided_from_its_seeds() {
use super::{GraphGate, PROTECTED_FILES};
let _data = crate::core::data_dir::isolated_data_dir();
let pack = crate::core::policy::parse(
"name = \"graph-ranking-test\"\nversion = \"1.0.0\"\ndescription = \"test\"\n\
[redaction]\ncustomer = 'K-[0-9]{6}'",
)
.expect("valid policy");
let _protected = TestPolicyOverride::set(Some(
crate::core::policy::resolve(&pack).expect("resolved policy"),
));
let project = tempfile::tempdir().expect("project");
let root = project.path().to_string_lossy().into_owned();
let files: Vec<String> = (0..PROTECTED_FILES + 2)
.map(|n| format!("f{n:03}.rs"))
.collect();
for file in &files {
std::fs::write(project.path().join(file), "pub fn item() {}\n").expect("source");
}
let edges: Vec<(String, String)> = files
.windows(2)
.map(|pair| (pair[0].clone(), pair[1].clone()))
.collect();
let gate = GraphGate::for_root(&root);
let local = gate.ranking_files(files.clone(), &edges, &files[..1]);
assert!(local.contains(&files[0]), "the seed is decided first");
assert!(local.len() <= PROTECTED_FILES);
assert!(
!local.contains(&files[PROTECTED_FILES + 1]),
"beyond the bound"
);
assert!(local.iter().all(|file| gate.file(file)));
assert!(!gate.file(&files[PROTECTED_FILES + 1]));
let unseeded = GraphGate::for_root(&root).ranking_files(files, &edges, &[]);
assert!(unseeded.is_empty());
}
#[test]
fn a_traversal_does_not_pass_through_a_withdrawn_file() {
use crate::core::property_graph::{CodeGraph, Edge, EdgeKind, Node};
let _data = crate::core::data_dir::isolated_data_dir();
let _community = TestPolicyOverride::set(None);
let project = tempfile::tempdir().expect("project");
let root = project.path().to_string_lossy().into_owned();
for file in ["a.rs", "w.rs", "b.rs", "n.rs"] {
std::fs::write(project.path().join(file), "pub fn item() {}\n").expect("source");
}
let graph = CodeGraph::open(&root).expect("graph");
let node = |file: &str| graph.upsert_node(&Node::file(file)).expect("node");
let (a, w, b, n) = (node("a.rs"), node("w.rs"), node("b.rs"), node("n.rs"));
for (from, to) in [(a, w), (w, b), (a, n)] {
graph
.upsert_edge(&Edge::new(from, to, EdgeKind::Imports))
.expect("edge");
}
let chain = graph.dependency_chain("a.rs", "b.rs").expect("chain");
assert_eq!(
chain.map(|chain| chain.path),
Some(vec![
"a.rs".to_owned(),
"w.rs".to_owned(),
"b.rs".to_owned()
])
);
let impact = graph.impact_analysis("b.rs", 5).expect("impact");
assert_eq!(impact.affected_files, ["a.rs", "w.rs"]);
assert_eq!(impact.edges_traversed, 2);
std::fs::remove_file(project.path().join("w.rs")).expect("withdraw");
let graph = CodeGraph::open(&root).expect("graph");
assert!(
graph
.dependency_chain("a.rs", "b.rs")
.expect("chain")
.is_none()
);
let impact = graph.impact_analysis("b.rs", 5).expect("impact");
assert!(impact.affected_files.is_empty(), "{impact:?}");
assert_eq!(impact.edges_traversed, 0);
assert_eq!(impact.max_depth_reached, 0);
let near = graph.dependency_chain("a.rs", "n.rs").expect("chain");
assert_eq!(near.map(|chain| chain.depth), Some(1));
let related: Vec<String> = graph
.related_files("a.rs", 10)
.expect("related")
.into_iter()
.map(|(path, _)| path)
.collect();
assert_eq!(related, ["n.rs"]);
}
}