use crate::security::{sanitize_label, MAX_GRAPH_FILE_BYTES};
use std::collections::{HashMap, HashSet};
#[derive(Clone, Debug)]
pub struct NodeData {
pub label: String,
pub source_file: String,
pub source_location: String,
pub community: Option<i64>,
}
#[derive(Clone, Debug)]
pub struct EdgeData {
pub relation: String,
pub confidence: String,
pub context: Option<String>,
}
#[derive(Debug)]
pub struct QueryGraph {
pub nodes: HashMap<String, NodeData>,
adj: HashMap<String, Vec<(String, EdgeData)>>,
pub idf_cache: HashMap<String, f64>,
}
impl Default for QueryGraph {
fn default() -> Self {
Self::new()
}
}
impl QueryGraph {
pub fn new() -> Self {
QueryGraph {
nodes: HashMap::new(),
adj: HashMap::new(),
idf_cache: HashMap::new(),
}
}
pub fn add_node(&mut self, id: impl Into<String>, data: NodeData) {
let id = id.into();
self.adj.entry(id.clone()).or_default();
self.nodes.insert(id, data);
}
pub fn add_edge(&mut self, u: impl Into<String>, v: impl Into<String>, data: EdgeData) {
let u: String = u.into();
let v: String = v.into();
let reverse = EdgeData {
relation: data.relation.clone(),
confidence: data.confidence.clone(),
context: data.context.clone(),
};
self.adj
.entry(u.clone())
.or_default()
.push((v.clone(), data));
self.adj.entry(v.clone()).or_default().push((u, reverse));
self.nodes.entry(v).or_insert_with(|| NodeData {
label: String::new(),
source_file: String::new(),
source_location: String::new(),
community: None,
});
}
pub fn neighbors(&self, id: &str) -> Vec<(&str, &EdgeData)> {
self.adj
.get(id)
.map(|v| v.iter().map(|(n, e)| (n.as_str(), e)).collect())
.unwrap_or_default()
}
pub fn degree(&self, id: &str) -> usize {
self.adj.get(id).map(|v| v.len()).unwrap_or(0)
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
fn all_degrees(&self) -> Vec<usize> {
self.nodes.keys().map(|id| self.degree(id)).collect()
}
fn hub_threshold(&self) -> usize {
let mut degrees = self.all_degrees();
if degrees.is_empty() {
return 50;
}
degrees.sort_unstable();
let p99_idx = (degrees.len() as f64 * 0.99) as usize;
let p99_idx = p99_idx.min(degrees.len() - 1);
degrees[p99_idx].max(50)
}
}
pub fn has_chinese(text: &str) -> bool {
text.chars().any(|c| ('\u{4e00}'..='\u{9fff}').contains(&c))
}
fn search_tokens(text: &str) -> Vec<String> {
let lower = text.to_lowercase();
let mut tokens = Vec::new();
let mut current = String::new();
for c in lower.chars() {
if c.is_alphanumeric() || c == '_' {
current.push(c);
} else if !current.is_empty() {
tokens.push(current.clone());
current.clear();
}
}
if !current.is_empty() {
tokens.push(current);
}
tokens
}
fn is_searchable(term: &str) -> bool {
let all_ascii_alpha = term.chars().all(|c| c.is_ascii_lowercase());
if all_ascii_alpha {
term.len() > 2
} else {
true
}
}
fn segment_chinese_bigrams(text: &str) -> Vec<String> {
let chars: Vec<char> = text.chars().collect();
let mut segments: Vec<String> = chars
.windows(2)
.map(|w| w.iter().collect::<String>())
.collect();
if segments.is_empty() && !chars.is_empty() {
segments.push(text.to_string());
}
if text.len() > 1 && !segments.contains(&text.to_string()) {
segments.push(text.to_string());
}
segments
}
pub fn query_terms(question: &str) -> Vec<String> {
let mut terms: Vec<String> = Vec::new();
for raw in question.split_whitespace() {
if has_chinese(raw) {
let lower = raw.to_lowercase();
let lower = lower.trim();
for seg in segment_chinese_bigrams(lower) {
let seg = seg.trim().to_string();
if !seg.is_empty() && is_searchable(&seg) {
terms.push(seg);
}
}
} else {
for tok in search_tokens(raw) {
if is_searchable(&tok) {
terms.push(tok);
}
}
}
}
terms
}
const EXACT_MATCH_BONUS: f64 = 1000.0;
const PREFIX_MATCH_BONUS: f64 = 100.0;
const SUBSTRING_MATCH_BONUS: f64 = 1.0;
const SOURCE_MATCH_BONUS: f64 = 0.5;
pub fn compute_idf(g: &mut QueryGraph, terms: &[String]) -> HashMap<String, f64> {
let n = g.node_count().max(1) as f64;
let uncached: Vec<String> = terms
.iter()
.filter(|t| !g.idf_cache.contains_key(*t))
.cloned()
.collect();
if !uncached.is_empty() {
let mut df: HashMap<String, usize> = uncached.iter().map(|t| (t.clone(), 0)).collect();
for data in g.nodes.values() {
let norm_label = data.label.to_lowercase();
for t in &uncached {
if norm_label.contains(t.as_str()) {
*df.get_mut(t).unwrap() += 1;
}
}
}
for t in &uncached {
let d = *df.get(t).unwrap_or(&0) as f64;
g.idf_cache.insert(t.clone(), (1.0 + n / (1.0 + d)).ln());
}
}
terms
.iter()
.map(|t| {
let w = g
.idf_cache
.get(t)
.copied()
.unwrap_or_else(|| (1.0 + n).ln());
(t.clone(), w)
})
.collect()
}
pub fn score_nodes(g: &mut QueryGraph, terms: &[String]) -> Vec<(f64, String)> {
let norm_terms: Vec<String> = terms.iter().flat_map(|t| search_tokens(t)).collect();
if norm_terms.is_empty() {
return vec![];
}
let idf = compute_idf(g, &norm_terms);
let node_ids: Vec<String> = g.nodes.keys().cloned().collect();
let mut scored: Vec<(f64, String)> = Vec::new();
for nid in node_ids {
let data = g.nodes.get(&nid).unwrap();
let norm_label = data.label.to_lowercase();
let bare_label = norm_label.trim_end_matches(['(', ')']).to_string();
let source = data.source_file.to_lowercase();
let mut score = 0.0f64;
for t in &norm_terms {
let w = idf.get(t).copied().unwrap_or(1.0);
if *t == norm_label || *t == bare_label {
score += EXACT_MATCH_BONUS * w;
} else if norm_label.starts_with(t.as_str()) || bare_label.starts_with(t.as_str()) {
score += PREFIX_MATCH_BONUS * w;
} else if norm_label.contains(t.as_str()) {
score += SUBSTRING_MATCH_BONUS * w;
}
if source.contains(t.as_str()) {
score += SOURCE_MATCH_BONUS * w;
}
}
if score > 0.0 {
scored.push((score, nid));
}
}
scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
scored
}
pub fn pick_seeds(scored: &[(f64, String)], max_k: usize, gap_ratio: f64) -> Vec<String> {
if scored.is_empty() {
return vec![];
}
let top_score = scored[0].0;
let mut seeds = Vec::new();
for (score, nid) in scored.iter().take(max_k) {
if !seeds.is_empty() && *score < top_score * gap_ratio {
break;
}
seeds.push(nid.clone());
}
seeds
}
static CONTEXT_HINTS: &[(&str, &[&str])] = &[
(
"call",
&["call", "calls", "called", "invoke", "invokes", "invoked"],
),
(
"import",
&["import", "imports", "imported", "module", "modules"],
),
(
"field",
&[
"field",
"fields",
"member",
"members",
"property",
"properties",
],
),
(
"parameter_type",
&[
"parameter",
"parameters",
"param",
"params",
"argument",
"arguments",
],
),
("return_type", &["return", "returns", "returned"]),
(
"generic_arg",
&["generic", "generics", "template", "templates"],
),
];
fn context_alias(key: &str) -> &str {
match key {
"param" | "params" | "parameter" | "parameters" | "argument" | "arguments" | "arg"
| "args" => "parameter_type",
"return" | "returns" | "returned" => "return_type",
"generic" | "generics" | "template" | "templates" => "generic_arg",
"annotation" | "annotations" | "decorator" | "decorators" => "attribute",
"calls" | "called" | "invoke" | "invocation" => "call",
"fields" | "property" | "properties" | "member" | "members" => "field",
"imports" | "imported" | "module" | "modules" => "import",
"exports" | "exported" => "export",
other => other,
}
}
pub fn normalize_context_filters(filters: &[String]) -> Vec<String> {
let mut seen = HashSet::new();
let mut result = Vec::new();
for v in filters {
let key = v.trim().to_lowercase();
if key.is_empty() {
continue;
}
let canonical = context_alias(&key).to_string();
if seen.insert(canonical.clone()) {
result.push(canonical);
}
}
result
}
pub fn infer_context_filters(question: &str) -> Vec<String> {
let lowered: HashSet<String> = question
.replace(['?', ','], " ")
.split_whitespace()
.map(|tok| tok.to_lowercase())
.collect();
let mut inferred = Vec::new();
for (context, hints) in CONTEXT_HINTS {
if hints.iter().any(|h| lowered.contains(*h)) {
inferred.push(context.to_string());
}
}
inferred
}
pub fn resolve_context_filters(
question: &str,
explicit_filters: &[String],
) -> (Vec<String>, Option<String>) {
let normalized = normalize_context_filters(explicit_filters);
if !normalized.is_empty() {
return (normalized, Some("explicit".to_string()));
}
let inferred = infer_context_filters(question);
if !inferred.is_empty() {
return (inferred, Some("heuristic".to_string()));
}
(vec![], None)
}
pub fn filter_graph_by_context(g: &QueryGraph, filters: &[String]) -> QueryGraph {
let filter_set: HashSet<String> = normalize_context_filters(filters).into_iter().collect();
if filter_set.is_empty() {
let mut h = QueryGraph::new();
for (id, data) in &g.nodes {
h.add_node(id.clone(), data.clone());
}
for (u, neighbors) in &g.adj {
for (v, edge) in neighbors {
if u < v {
h.adj
.entry(u.clone())
.or_default()
.push((v.clone(), edge.clone()));
h.adj.entry(v.clone()).or_default().push((
u.clone(),
EdgeData {
relation: edge.relation.clone(),
confidence: edge.confidence.clone(),
context: edge.context.clone(),
},
));
}
}
}
return h;
}
let mut h = QueryGraph::new();
for (id, data) in &g.nodes {
h.add_node(id.clone(), data.clone());
}
let mut added: HashSet<(String, String)> = HashSet::new();
for (u, neighbors) in &g.adj {
for (v, edge) in neighbors {
let matches = edge
.context
.as_deref()
.map(|c| filter_set.contains(c))
.unwrap_or(false);
if matches {
let key = if u <= v {
(u.clone(), v.clone())
} else {
(v.clone(), u.clone())
};
if added.insert(key) {
h.adj
.entry(u.clone())
.or_default()
.push((v.clone(), edge.clone()));
h.adj.entry(v.clone()).or_default().push((
u.clone(),
EdgeData {
relation: edge.relation.clone(),
confidence: edge.confidence.clone(),
context: edge.context.clone(),
},
));
}
}
}
}
h
}
pub fn bfs(
g: &QueryGraph,
start_nodes: &[String],
depth: usize,
) -> (HashSet<String>, Vec<(String, String)>) {
let hub_threshold = g.hub_threshold();
let seed_set: HashSet<&str> = start_nodes.iter().map(|s| s.as_str()).collect();
let mut visited: HashSet<String> = start_nodes.iter().cloned().collect();
let mut frontier: HashSet<String> = start_nodes.iter().cloned().collect();
let mut edges_seen: Vec<(String, String)> = Vec::new();
for _ in 0..depth {
let mut next_frontier = HashSet::new();
for n in &frontier {
if !seed_set.contains(n.as_str()) && g.degree(n) >= hub_threshold {
continue;
}
for (neighbor, _) in g.neighbors(n) {
if !visited.contains(neighbor) {
next_frontier.insert(neighbor.to_string());
edges_seen.push((n.clone(), neighbor.to_string()));
}
}
}
visited.extend(next_frontier.iter().cloned());
frontier = next_frontier;
if frontier.is_empty() {
break;
}
}
(visited, edges_seen)
}
pub fn dfs(
g: &QueryGraph,
start_nodes: &[String],
depth: usize,
) -> (HashSet<String>, Vec<(String, String)>) {
let hub_threshold = g.hub_threshold();
let seed_set: HashSet<&str> = start_nodes.iter().map(|s| s.as_str()).collect();
let mut visited: HashSet<String> = HashSet::new();
let mut edges_seen: Vec<(String, String)> = Vec::new();
let mut stack: Vec<(String, usize)> =
start_nodes.iter().rev().map(|n| (n.clone(), 0)).collect();
while let Some((node, d)) = stack.pop() {
if visited.contains(&node) || d > depth {
continue;
}
visited.insert(node.clone());
if !seed_set.contains(node.as_str()) && g.degree(&node) >= hub_threshold {
continue;
}
for (neighbor, _) in g.neighbors(&node) {
if !visited.contains(neighbor) {
stack.push((neighbor.to_string(), d + 1));
edges_seen.push((node.clone(), neighbor.to_string()));
}
}
}
(visited, edges_seen)
}
pub fn subgraph_to_text(
g: &QueryGraph,
nodes: &HashSet<String>,
edges: &[(String, String)],
token_budget: usize,
seeds: Option<&[String]>,
) -> String {
let char_budget = token_budget * 3;
let mut lines: Vec<String> = Vec::new();
let _seed_set: HashSet<&str> = seeds.unwrap_or(&[]).iter().map(|s| s.as_str()).collect();
let mut ordered: Vec<String> = seeds
.unwrap_or(&[])
.iter()
.filter(|n| nodes.contains(*n))
.cloned()
.collect();
let rest_set: HashSet<&str> = ordered.iter().map(|s| s.as_str()).collect();
let mut rest: Vec<String> = nodes
.iter()
.filter(|n| !rest_set.contains(n.as_str()))
.cloned()
.collect();
rest.sort_by_key(|n| std::cmp::Reverse(g.degree(n)));
ordered.extend(rest);
for nid in &ordered {
let d = g.nodes.get(nid.as_str());
let (label, src, loc, community) = match d {
Some(nd) => (
nd.label.as_str(),
nd.source_file.as_str(),
nd.source_location.as_str(),
nd.community.map(|c| c.to_string()),
),
None => (nid.as_str(), "", "", None),
};
let line = format!(
"NODE {} [src={} loc={} community={}]",
sanitize_label(Some(label)),
sanitize_label(Some(src)),
sanitize_label(Some(loc)),
sanitize_label(community.as_deref()),
);
lines.push(line);
}
for (u, v) in edges {
if nodes.contains(u) && nodes.contains(v) {
let edge = g
.adj
.get(u)
.and_then(|adj| adj.iter().find(|(n, _)| n == v))
.map(|(_, e)| e);
let (relation, confidence, context) = match edge {
Some(e) => (
e.relation.as_str(),
e.confidence.as_str(),
e.context.as_deref(),
),
None => ("", "", None),
};
let context_suffix = if let Some(ctx) = context {
format!(" context={}", sanitize_label(Some(ctx)))
} else {
String::new()
};
let u_label = g
.nodes
.get(u)
.map(|d| d.label.as_str())
.unwrap_or(u.as_str());
let v_label = g
.nodes
.get(v)
.map(|d| d.label.as_str())
.unwrap_or(v.as_str());
let line = format!(
"EDGE {} --{} [{}{}]--> {}",
sanitize_label(Some(u_label)),
sanitize_label(Some(relation)),
sanitize_label(Some(confidence)),
context_suffix,
sanitize_label(Some(v_label)),
);
lines.push(line);
}
}
let output = lines.join("\n");
if output.len() > char_budget {
let cut_at = output[..char_budget]
.rfind('\n')
.filter(|&p| p > 0)
.unwrap_or(char_budget);
let total_nodes = lines.iter().filter(|l| l.starts_with("NODE ")).count();
let shown_nodes = output[..cut_at]
.split('\n')
.filter(|l| l.starts_with("NODE "))
.count();
let cut_count = total_nodes.saturating_sub(shown_nodes);
return format!(
"{}\n... (truncated — {} more nodes cut by ~{}-token budget. Narrow with context_filter=['call'] or use get_node for a specific symbol)",
&output[..cut_at],
cut_count,
token_budget,
);
}
output
}
pub fn find_node(g: &QueryGraph, label: &str) -> Vec<String> {
let tokens = search_tokens(label);
let term = tokens.join(" ");
if term.is_empty() {
return vec![];
}
let mut exact = Vec::new();
let mut prefix = Vec::new();
let mut substring = Vec::new();
for (nid, data) in &g.nodes {
let norm_label = data.label.to_lowercase();
let bare_label = norm_label.trim_end_matches(['(', ')']).to_string();
let nid_lower = nid.to_lowercase();
if term == norm_label || term == bare_label || term == nid_lower {
exact.push(nid.clone());
} else if norm_label.starts_with(&term)
|| bare_label.starts_with(&term)
|| nid_lower.starts_with(&term)
{
prefix.push(nid.clone());
} else if norm_label.contains(&term) {
substring.push(nid.clone());
}
}
let mut result = exact;
result.extend(prefix);
result.extend(substring);
result
}
pub fn communities_from_graph(g: &QueryGraph) -> HashMap<i64, Vec<String>> {
let mut communities: HashMap<i64, Vec<String>> = HashMap::new();
for (nid, data) in &g.nodes {
if let Some(cid) = data.community {
communities.entry(cid).or_default().push(nid.clone());
}
}
communities
}
pub fn query_graph_text(
g: &mut QueryGraph,
question: &str,
mode: &str,
depth: usize,
token_budget: usize,
context_filters: &[String],
) -> String {
let terms = query_terms(question);
let scored = score_nodes(g, &terms);
let start_nodes = pick_seeds(&scored, 3, 0.2);
if start_nodes.is_empty() {
return "No matching nodes found.".to_string();
}
let (resolved_filters, filter_source) = resolve_context_filters(question, context_filters);
let traversal_graph = filter_graph_by_context(g, &resolved_filters);
let (nodes, edges) = if mode == "dfs" {
dfs(&traversal_graph, &start_nodes, depth)
} else {
bfs(&traversal_graph, &start_nodes, depth)
};
let seed_labels: Vec<String> = start_nodes
.iter()
.map(|n| {
g.nodes
.get(n)
.map(|d| d.label.clone())
.unwrap_or_else(|| n.clone())
})
.collect();
let mut header_parts = vec![
format!("Traversal: {} depth={}", mode.to_uppercase(), depth),
format!("Start: {:?}", seed_labels),
];
if !resolved_filters.is_empty() {
let src = filter_source.as_deref().unwrap_or("unknown");
header_parts.push(format!(
"Context: {} ({})",
resolved_filters.join(", "),
src
));
}
header_parts.push(format!("{} nodes found", nodes.len()));
let header = header_parts.join(" | ") + "\n\n";
header
+ &subgraph_to_text(
&traversal_graph,
&nodes,
&edges,
token_budget,
Some(&start_nodes),
)
}
pub fn load_graph_with_cap(path: &str, max_bytes: u64) -> Result<QueryGraph, String> {
use std::path::Path;
let resolved = Path::new(path);
if resolved.extension().and_then(|e| e.to_str()) != Some("json") {
return Err(format!("Graph path must be a .json file, got: {:?}", path));
}
if !resolved.exists() {
return Err(format!("Graph file not found: {}", path));
}
let meta = std::fs::metadata(resolved).map_err(|e| e.to_string())?;
if meta.len() > max_bytes {
let msg = format!(
"error: graph.json ({} bytes) exceeds {} byte cap",
meta.len(),
max_bytes
);
eprintln!("{}", msg);
return Err(msg);
}
let content = std::fs::read_to_string(resolved).map_err(|e| e.to_string())?;
let data: serde_json::Value = serde_json::from_str(&content)
.map_err(|e| format!("graph.json is corrupted ({}). Re-run to rebuild.", e))?;
let mut g = QueryGraph::new();
if let Some(nodes) = data.get("nodes").and_then(|v| v.as_array()) {
for node in nodes {
let id = node
.get("id")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
if id.is_empty() {
continue;
}
let label = node
.get("label")
.and_then(|v| v.as_str())
.unwrap_or(&id)
.to_string();
let source_file = node
.get("source_file")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let source_location = node
.get("source_location")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let community = node.get("community").and_then(|v| v.as_i64());
g.add_node(
id,
NodeData {
label,
source_file,
source_location,
community,
},
);
}
}
let links_key = if data.get("links").is_some() {
"links"
} else {
"edges"
};
if let Some(links) = data.get(links_key).and_then(|v| v.as_array()) {
for link in links {
let source = link
.get("source")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let target = link
.get("target")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
if source.is_empty() || target.is_empty() {
continue;
}
let relation = link
.get("relation")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let confidence = link
.get("confidence")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let context = link
.get("context")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
g.add_edge(
source,
target,
EdgeData {
relation,
confidence,
context,
},
);
}
}
Ok(g)
}
pub fn load_graph(path: &str) -> Result<QueryGraph, String> {
load_graph_with_cap(path, MAX_GRAPH_FILE_BYTES)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
fn make_graph() -> QueryGraph {
let mut g = QueryGraph::new();
g.add_node(
"n1",
NodeData {
label: "extract".into(),
source_file: "extract.py".into(),
source_location: "L10".into(),
community: Some(0),
},
);
g.add_node(
"n2",
NodeData {
label: "cluster".into(),
source_file: "cluster.py".into(),
source_location: "L5".into(),
community: Some(0),
},
);
g.add_node(
"n3",
NodeData {
label: "build".into(),
source_file: "build.py".into(),
source_location: "L1".into(),
community: Some(1),
},
);
g.add_node(
"n4",
NodeData {
label: "report".into(),
source_file: "report.py".into(),
source_location: "L1".into(),
community: Some(1),
},
);
g.add_node(
"n5",
NodeData {
label: "isolated".into(),
source_file: "other.py".into(),
source_location: "L1".into(),
community: Some(2),
},
);
g.add_edge(
"n1",
"n2",
EdgeData {
relation: "calls".into(),
confidence: "INFERRED".into(),
context: Some("call".into()),
},
);
g.add_edge(
"n2",
"n3",
EdgeData {
relation: "imports".into(),
confidence: "EXTRACTED".into(),
context: Some("import".into()),
},
);
g.add_edge(
"n3",
"n4",
EdgeData {
relation: "uses".into(),
confidence: "EXTRACTED".into(),
context: None,
},
);
g
}
fn write_graph_json(nodes: &[&str]) -> NamedTempFile {
let nodes_json: Vec<serde_json::Value> = nodes
.iter()
.map(|n| {
serde_json::json!({
"id": n,
"label": n,
"community": 0
})
})
.collect();
let data = serde_json::json!({
"directed": false,
"nodes": nodes_json,
"links": []
});
let mut f = NamedTempFile::with_suffix(".json").unwrap();
f.write_all(data.to_string().as_bytes()).unwrap();
f
}
#[test]
fn test_communities_from_graph_basic() {
let g = make_graph();
let c = communities_from_graph(&g);
assert!(c.contains_key(&0));
assert!(c.contains_key(&1));
assert!(c[&0].contains(&"n1".to_string()));
assert!(c[&0].contains(&"n2".to_string()));
assert!(c[&1].contains(&"n3".to_string()));
}
#[test]
fn test_communities_from_graph_no_community_attr() {
let mut g = QueryGraph::new();
g.add_node(
"a",
NodeData {
label: "foo".into(),
source_file: "".into(),
source_location: "".into(),
community: None,
},
);
let c = communities_from_graph(&g);
assert!(c.is_empty());
}
#[test]
fn test_communities_from_graph_isolated() {
let g = make_graph();
let c = communities_from_graph(&g);
assert!(c.contains_key(&2));
assert!(c[&2].contains(&"n5".to_string()));
}
#[test]
fn test_score_nodes_exact_label_match() {
let mut g = make_graph();
let scored = score_nodes(&mut g, &["extract".to_string()]);
let nids: Vec<&str> = scored.iter().map(|(_, n)| n.as_str()).collect();
assert!(nids.contains(&"n1"));
assert_eq!(scored[0].1, "n1");
}
#[test]
fn test_score_nodes_no_match() {
let mut g = make_graph();
let scored = score_nodes(&mut g, &["xyzzy".to_string()]);
assert!(scored.is_empty());
}
#[test]
fn test_score_nodes_source_file_partial() {
let mut g = make_graph();
let scored = score_nodes(&mut g, &["cluster".to_string()]);
let nids: Vec<&str> = scored.iter().map(|(_, n)| n.as_str()).collect();
assert!(nids.contains(&"n2"));
}
#[test]
fn test_score_nodes_ignores_trailing_punctuation() {
let mut g = make_graph();
let scored = score_nodes(&mut g, &["extract?".to_string()]);
assert!(!scored.is_empty());
assert_eq!(scored[0].1, "n1");
}
#[test]
fn test_find_node_ignores_trailing_punctuation() {
let g = make_graph();
let result = find_node(&g, "extract?");
assert!(result.contains(&"n1".to_string()));
}
#[test]
fn test_query_terms_strips_search_punctuation() {
let terms = query_terms("what calls extract?");
assert!(terms.contains(&"what".to_string()));
assert!(terms.contains(&"calls".to_string()));
assert!(terms.contains(&"extract".to_string()));
assert!(!terms.iter().any(|t| t.contains('?')));
}
#[test]
fn test_query_graph_text_keeps_short_non_english_terms() {
let mut g = QueryGraph::new();
g.add_node(
"frontend",
NodeData {
label: "前端".into(),
source_file: "docs/前端.md".into(),
source_location: "L1".into(),
community: Some(0),
},
);
let text = query_graph_text(&mut g, "前端", "bfs", 1, 2000, &[]);
assert!(!text.contains("No matching nodes found."));
assert!(text.contains("NODE 前端"));
}
#[test]
fn test_infer_context_filters_for_calls_question() {
let filters = infer_context_filters("who calls extract");
assert!(filters.contains(&"call".to_string()));
}
#[test]
fn test_resolve_context_filters_explicit_overrides_heuristic() {
let (filters, source) =
resolve_context_filters("who calls extract", &["field".to_string()]);
assert_eq!(filters, vec!["field"]);
assert_eq!(source.as_deref(), Some("explicit"));
}
#[test]
fn test_bfs_depth_1() {
let g = make_graph();
let (visited, _) = bfs(&g, &["n1".to_string()], 1);
assert!(visited.contains("n1"));
assert!(visited.contains("n2"));
assert!(!visited.contains("n3"));
}
#[test]
fn test_bfs_depth_2() {
let g = make_graph();
let (visited, _) = bfs(&g, &["n1".to_string()], 2);
assert!(visited.contains("n3"));
}
#[test]
fn test_bfs_disconnected() {
let g = make_graph();
let (visited, _) = bfs(&g, &["n5".to_string()], 3);
assert_eq!(visited, HashSet::from(["n5".to_string()]));
}
#[test]
fn test_bfs_returns_edges() {
let g = make_graph();
let (_, edges) = bfs(&g, &["n1".to_string()], 1);
assert!(!edges.is_empty());
assert!(edges.iter().any(|(u, v)| u == "n1" || v == "n1"));
}
#[test]
fn test_filter_graph_by_context_limits_traversal() {
let g = make_graph();
let filtered = filter_graph_by_context(&g, &["call".to_string()]);
let (visited, edges) = bfs(&filtered, &["n1".to_string()], 2);
assert!(visited.contains("n2"));
assert!(!visited.contains("n3"));
assert_eq!(edges.len(), 1);
}
#[test]
fn test_dfs_depth_1() {
let g = make_graph();
let (visited, _) = dfs(&g, &["n1".to_string()], 1);
assert!(visited.contains("n1"));
assert!(visited.contains("n2"));
assert!(!visited.contains("n3"));
}
#[test]
fn test_dfs_full_chain() {
let g = make_graph();
let (visited, _) = dfs(&g, &["n1".to_string()], 5);
for n in &["n1", "n2", "n3", "n4"] {
assert!(visited.contains(*n), "missing {}", n);
}
}
#[test]
fn test_subgraph_to_text_contains_labels() {
let g = make_graph();
let nodes: HashSet<String> = ["n1", "n2"].iter().map(|s| s.to_string()).collect();
let edges = vec![("n1".to_string(), "n2".to_string())];
let text = subgraph_to_text(&g, &nodes, &edges, 2000, None);
assert!(text.contains("extract"));
assert!(text.contains("cluster"));
}
#[test]
fn test_subgraph_to_text_truncates() {
let g = make_graph();
let nodes: HashSet<String> = ["n1", "n2", "n3", "n4"]
.iter()
.map(|s| s.to_string())
.collect();
let edges = vec![("n1".to_string(), "n2".to_string())];
let text = subgraph_to_text(&g, &nodes, &edges, 1, None);
assert!(text.contains("truncated"));
}
#[test]
fn test_subgraph_to_text_edge_included() {
let g = make_graph();
let nodes: HashSet<String> = ["n1", "n2"].iter().map(|s| s.to_string()).collect();
let edges = vec![("n1".to_string(), "n2".to_string())];
let text = subgraph_to_text(&g, &nodes, &edges, 2000, None);
assert!(text.contains("EDGE"));
assert!(text.contains("calls"));
}
#[test]
fn test_subgraph_to_text_includes_edge_context() {
let g = make_graph();
let nodes: HashSet<String> = ["n1", "n2"].iter().map(|s| s.to_string()).collect();
let edges = vec![("n1".to_string(), "n2".to_string())];
let text = subgraph_to_text(&g, &nodes, &edges, 2000, None);
assert!(text.contains("context=call"));
}
#[test]
fn test_query_graph_text_explicit_context_filter_changes_traversal() {
let mut g = make_graph();
let text = query_graph_text(&mut g, "extract", "bfs", 2, 2000, &["call".to_string()]);
assert!(text.contains("Context: call (explicit)"));
assert!(text.contains("cluster"));
assert!(!text.contains("build"));
}
#[test]
fn test_query_graph_text_heuristic_context_filter_changes_traversal() {
let mut g = make_graph();
let text = query_graph_text(&mut g, "who calls extract", "bfs", 2, 2000, &[]);
assert!(text.contains("Context: call (heuristic)"));
assert!(text.contains("cluster"));
assert!(!text.contains("build"));
}
#[test]
fn test_load_graph_roundtrip() {
let f = write_graph_json(&["alpha", "beta"]);
let g = load_graph(f.path().to_str().unwrap()).unwrap();
assert_eq!(g.node_count(), 2);
}
#[test]
fn test_load_graph_missing_file() {
let result = load_graph("/tmp/does_not_exist_codesynapse.json");
assert!(result.is_err());
}
#[test]
fn test_load_graph_rejects_oversized_file() {
let f = write_graph_json(&["a"]);
let result = load_graph_with_cap(f.path().to_str().unwrap(), 16);
assert!(result.is_err());
let msg = result.unwrap_err();
assert!(msg.contains("exceeds") && msg.contains("byte cap"));
}
#[test]
fn test_load_graph_accepts_under_cap() {
let f = write_graph_json(&["a"]);
let result = load_graph_with_cap(f.path().to_str().unwrap(), 10 * 1024 * 1024);
assert!(result.is_ok());
assert_eq!(result.unwrap().node_count(), 1);
}
#[test]
fn test_maybe_reload_detects_graph_change() {
let f1 = write_graph_json(&["alpha", "beta"]);
let path = f1.path().to_str().unwrap();
let g1 = load_graph(path).unwrap();
assert_eq!(g1.node_count(), 2);
let nodes_json: Vec<serde_json::Value> = ["alpha", "beta", "gamma"]
.iter()
.map(|n| serde_json::json!({"id": n, "label": n, "community": 0}))
.collect();
let data = serde_json::json!({"directed": false, "nodes": nodes_json, "links": []});
std::fs::write(path, data.to_string()).unwrap();
let g2 = load_graph(path).unwrap();
assert!(g2.nodes.contains_key("gamma"));
}
#[test]
fn test_load_graph_cache_key_changes_with_content() {
let f = write_graph_json(&["a"]);
let path = f.path();
let s1 = std::fs::metadata(path).unwrap();
let key1 = (s1.modified().unwrap(), s1.len());
std::thread::sleep(std::time::Duration::from_millis(10));
let nodes_json = vec![
serde_json::json!({"id": "a", "label": "a", "community": 0}),
serde_json::json!({"id": "b", "label": "b", "community": 0}),
];
let data = serde_json::json!({"directed": false, "nodes": nodes_json, "links": []});
std::fs::write(path, data.to_string()).unwrap();
let s2 = std::fs::metadata(path).unwrap();
let key2 = (s2.modified().unwrap(), s2.len());
assert_ne!(key1, key2, "stat key must change when file content changes");
}
fn make_noisy_graph() -> QueryGraph {
let mut g = QueryGraph::new();
for i in 0..20usize {
let id = format!("err{}", i);
g.add_node(
id.clone(),
NodeData {
label: format!("error_handler_{}", i),
source_file: format!("err{}.py", i),
source_location: "L1".into(),
community: Some(0),
},
);
if i > 0 {
let prev = format!("err{}", i - 1);
g.add_edge(
prev,
id,
EdgeData {
relation: "calls".into(),
confidence: "EXTRACTED".into(),
context: Some("call".into()),
},
);
}
}
g.add_node(
"fbs",
NodeData {
label: "FooBarService".into(),
source_file: "service.py".into(),
source_location: "L1".into(),
community: Some(1),
},
);
g.add_node(
"fbs_dep",
NodeData {
label: "ServiceClient".into(),
source_file: "client.py".into(),
source_location: "L1".into(),
community: Some(1),
},
);
g.add_edge(
"fbs",
"fbs_dep",
EdgeData {
relation: "uses".into(),
confidence: "EXTRACTED".into(),
context: None,
},
);
g
}
#[test]
fn test_idf_downweights_common_terms() {
let mut g = make_noisy_graph();
let scored = score_nodes(&mut g, &["foobarservice".to_string(), "error".to_string()]);
assert!(!scored.is_empty());
assert_eq!(scored[0].1, "fbs", "FooBarService should rank first");
}
#[test]
fn test_idf_cached_on_graph() {
let mut g = make_graph();
score_nodes(&mut g, &["extract".to_string()]);
assert!(g.idf_cache.contains_key("extract"));
}
#[test]
fn test_idf_new_graph_starts_fresh() {
let mut g1 = make_graph();
let g2 = make_graph();
score_nodes(&mut g1, &["extract".to_string()]);
assert!(g1.idf_cache.contains_key("extract"));
assert!(!g2.idf_cache.contains_key("extract"));
}
#[test]
fn test_idf_rare_term_gets_high_weight() {
let mut g = make_graph();
let idf = compute_idf(&mut g, &["extract".to_string()]);
assert!(idf["extract"] > 1.0, "rare term should get IDF > 1");
}
#[test]
fn test_idf_common_term_gets_low_weight() {
let mut g = QueryGraph::new();
for i in 0..20usize {
g.add_node(
format!("n{}", i),
NodeData {
label: format!("handle_{}", i),
source_file: format!("f{}.py", i),
source_location: "L1".into(),
community: None,
},
);
}
let mut g_ref = g;
let idf = compute_idf(&mut g_ref, &["handle".to_string()]);
assert!(idf["handle"] < 1.0, "common term should get IDF < 1");
}
#[test]
fn test_pick_seeds_dominant_identifier_gives_one_seed() {
let scored = vec![
(1000.0, "fbs".to_string()),
(1.0, "err1".to_string()),
(0.9, "err2".to_string()),
];
let seeds = pick_seeds(&scored, 3, 0.2);
assert_eq!(seeds, vec!["fbs"]);
}
#[test]
fn test_pick_seeds_close_scores_keeps_multiple() {
let scored = vec![
(10.0, "a".to_string()),
(9.0, "b".to_string()),
(8.5, "c".to_string()),
];
let seeds = pick_seeds(&scored, 3, 0.2);
assert_eq!(seeds.len(), 3);
}
#[test]
fn test_pick_seeds_empty() {
assert!(pick_seeds(&[], 3, 0.2).is_empty());
}
#[test]
fn test_pick_seeds_single() {
let scored = vec![(5.0, "x".to_string())];
assert_eq!(pick_seeds(&scored, 3, 0.2), vec!["x"]);
}
#[test]
fn test_pick_seeds_respects_max_k() {
let scored: Vec<(f64, String)> = (0..10).map(|i| (10.0, format!("n{}", i))).collect();
let seeds = pick_seeds(&scored, 3, 0.2);
assert_eq!(seeds.len(), 3);
}
#[test]
fn test_subgraph_to_text_truncation_hint_is_actionable() {
let g = make_graph();
let nodes: HashSet<String> = ["n1", "n2", "n3", "n4"]
.iter()
.map(|s| s.to_string())
.collect();
let edges = vec![("n1".to_string(), "n2".to_string())];
let text = subgraph_to_text(&g, &nodes, &edges, 1, None);
assert!(text.contains("truncated"));
assert!(text.contains("get_node") || text.contains("context_filter"));
}
#[test]
fn test_query_seeds_from_identifier_not_noise() {
let mut g = make_noisy_graph();
let text = query_graph_text(&mut g, "FooBarService error handling", "bfs", 2, 2000, &[]);
assert!(text.contains("FooBarService"));
assert!(text.contains("ServiceClient"));
}
#[test]
fn test_query_graph_text_parameter_type_context_filter() {
let mut g = QueryGraph::new();
g.add_node(
"process",
NodeData {
label: "process".into(),
source_file: "sample.cs".into(),
source_location: "L20".into(),
community: None,
},
);
g.add_node(
"payload",
NodeData {
label: "Payload".into(),
source_file: "sample.cs".into(),
source_location: "L5".into(),
community: None,
},
);
g.add_node(
"other",
NodeData {
label: "PayloadFactory".into(),
source_file: "sample.cs".into(),
source_location: "L40".into(),
community: None,
},
);
g.add_edge(
"process",
"payload",
EdgeData {
relation: "references".into(),
confidence: "EXTRACTED".into(),
context: Some("parameter_type".into()),
},
);
g.add_edge(
"process",
"other",
EdgeData {
relation: "calls".into(),
confidence: "EXTRACTED".into(),
context: Some("call".into()),
},
);
let text = query_graph_text(
&mut g,
"who accepts Payload",
"bfs",
2,
2000,
&["parameter_type".to_string()],
);
assert!(text.contains("parameter_type"));
assert!(text.contains("Payload"));
assert!(!text.contains("PayloadFactory"));
}
#[test]
fn test_query_graph_text_context_filter_aliases_resolve() {
assert_eq!(
normalize_context_filters(&["param".to_string()]),
vec!["parameter_type"]
);
assert_eq!(
normalize_context_filters(&["parameter".to_string()]),
vec!["parameter_type"]
);
assert_eq!(
normalize_context_filters(&["return".to_string()]),
vec!["return_type"]
);
assert_eq!(
normalize_context_filters(&["returns".to_string()]),
vec!["return_type"]
);
assert_eq!(
normalize_context_filters(&["generic".to_string()]),
vec!["generic_arg"]
);
assert_eq!(
normalize_context_filters(&["generics".to_string()]),
vec!["generic_arg"]
);
assert_eq!(
normalize_context_filters(&["annotation".to_string()]),
vec!["attribute"]
);
assert_eq!(
normalize_context_filters(&["decorator".to_string()]),
vec!["attribute"]
);
assert_eq!(
normalize_context_filters(&["parameter_type".to_string()]),
vec!["parameter_type"]
);
assert_eq!(
normalize_context_filters(&["field".to_string()]),
vec!["field"]
);
}
#[test]
fn test_query_terms_chinese_mixed() {
let terms = query_terms("前端 router 路由配置");
assert!(terms.iter().any(|t| t.contains("前端") || t == "前端"));
assert!(terms.contains(&"router".to_string()));
assert!(terms.iter().any(|t| t.contains("路由")));
assert!(terms.iter().any(|t| t.contains("配置")));
}
#[test]
fn test_query_terms_non_chinese_scripts_are_not_segmented() {
assert!(!has_chinese("かなカナ한글"));
let terms = query_terms("かなカナ한글");
assert!(terms.contains(&"かなカナ한글".to_string()));
}
#[test]
fn test_query_terms_chinese_no_jieba_fallback() {
let terms = query_terms("页面路由");
assert!(terms.iter().any(|t| t.contains("页面")));
assert!(terms.iter().any(|t| t.contains("路由")));
assert!(terms.contains(&"页面路由".to_string()));
assert_eq!(terms.len(), 4); }
#[test]
fn test_score_nodes_chinese_substring_match() {
let mut g = QueryGraph::new();
g.add_node(
"n1",
NodeData {
label: "路由桥接核对表".into(),
source_file: "doc.md".into(),
source_location: "L1".into(),
community: Some(0),
},
);
g.add_node(
"n2",
NodeData {
label: "其他内容".into(),
source_file: "doc.md".into(),
source_location: "L1".into(),
community: Some(0),
},
);
let scored = score_nodes(&mut g, &["路由".to_string()]);
let nids: Vec<&str> = scored.iter().map(|(_, n)| n.as_str()).collect();
assert!(nids.contains(&"n1"));
assert!(!nids.contains(&"n2"));
}
#[test]
fn test_query_text_chinese_finds_routing_nodes() {
let mut g = QueryGraph::new();
g.add_node(
"parent",
NodeData {
label: "页面路由规范".into(),
source_file: "doc.md".into(),
source_location: "L1".into(),
community: Some(0),
},
);
g.add_node(
"child",
NodeData {
label: "路由桥接核对表".into(),
source_file: "doc.md".into(),
source_location: "L10".into(),
community: Some(0),
},
);
g.add_edge(
"parent",
"child",
EdgeData {
relation: "contains".into(),
confidence: "EXTRACTED".into(),
context: None,
},
);
let text = query_graph_text(&mut g, "页面路由", "bfs", 2, 2000, &[]);
assert!(!text.contains("No matching nodes found."));
assert!(text.contains("路由"));
}
}