use std::collections::{HashMap, HashSet, VecDeque};
use std::io::Write;
use std::path::Path;
use serde_json::Value;
use crate::error::Result;
use crate::security::{check_file_size, MAX_GRAPH_FILE_BYTES};
const CHARS_PER_TOKEN: usize = 4;
pub const SAMPLE_QUESTIONS: &[&str] = &[
"how does authentication work",
"what is the main entry point",
"how are errors handled",
"what connects the data layer to the api",
"what are the core abstractions",
];
pub fn safe_char<'a>(unicode: &'a str, ascii: &'a str, supports_unicode: bool) -> &'a str {
if supports_unicode {
unicode
} else {
ascii
}
}
pub fn hr(width: usize, supports_unicode: bool) -> String {
safe_char("─", "-", supports_unicode).repeat(width)
}
#[derive(Debug, Default)]
struct BenchGraph {
nodes: Vec<BenchNode>,
node_index: HashMap<String, usize>,
adj: HashMap<String, Vec<String>>,
edge_count: usize,
}
#[derive(Debug, Clone)]
struct BenchNode {
id: String,
label: String,
source_file: String,
source_location: String,
}
fn value_to_id(v: &Value) -> Option<String> {
match v {
Value::String(s) => Some(s.clone()),
Value::Number(n) => Some(n.to_string()),
_ => None,
}
}
impl BenchGraph {
fn from_json(data: &Value) -> Option<Self> {
let mut g = BenchGraph::default();
let nodes = data.get("nodes")?.as_array()?;
for n in nodes {
let id = n.get("id").and_then(value_to_id).unwrap_or_default();
if id.is_empty() {
continue;
}
let label = n
.get("label")
.and_then(Value::as_str)
.unwrap_or("")
.to_string();
let source_file = n
.get("source_file")
.and_then(Value::as_str)
.unwrap_or("")
.to_string();
let source_location = n
.get("source_location")
.and_then(Value::as_str)
.unwrap_or("")
.to_string();
let idx = g.nodes.len();
g.node_index.insert(id.clone(), idx);
g.adj.entry(id.clone()).or_default();
g.nodes.push(BenchNode {
id,
label,
source_file,
source_location,
});
}
let links_key = if data.get("links").is_some() {
"links"
} else {
"edges"
};
if let Some(links) = data.get(links_key).and_then(Value::as_array) {
for e in links {
let src = e.get("source").and_then(value_to_id).unwrap_or_default();
let tgt = e.get("target").and_then(value_to_id).unwrap_or_default();
if src.is_empty() || tgt.is_empty() {
continue;
}
g.adj.entry(src.clone()).or_default().push(tgt.clone());
g.adj.entry(tgt.clone()).or_default().push(src.clone());
g.edge_count += 1;
}
}
Some(g)
}
fn node_count(&self) -> usize {
self.nodes.len()
}
fn edge_count(&self) -> usize {
self.edge_count
}
}
fn estimate_tokens(text: &str) -> u64 {
(text.len() / CHARS_PER_TOKEN).max(1) as u64
}
fn has_non_ascii(s: &str) -> bool {
!s.is_ascii()
}
fn is_searchable(term: &str) -> bool {
if term
.chars()
.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit())
{
term.len() > 2
} else {
true
}
}
fn bench_query_terms(question: &str) -> Vec<String> {
let mut terms = Vec::new();
for word in question.split_whitespace() {
if has_non_ascii(word) {
let s = word.to_string();
if is_searchable(&s) {
terms.push(s);
}
} else {
let lower = word.to_lowercase();
let mut start: Option<usize> = None;
let chars: Vec<(usize, char)> = lower.char_indices().collect();
for (i, c) in &chars {
if c.is_alphanumeric() || *c == '_' {
if start.is_none() {
start = Some(*i);
}
} else if let Some(s) = start {
let tok = &lower[s..*i];
if is_searchable(tok) {
terms.push(tok.to_string());
}
start = None;
}
}
if let Some(s) = start {
let tok = &lower[s..];
if is_searchable(tok) {
terms.push(tok.to_string());
}
}
}
}
terms
}
pub fn query_subgraph_tokens(g_data: &Value, question: &str, depth: u32) -> u64 {
let g = match BenchGraph::from_json(g_data) {
Some(g) => g,
None => return 0,
};
query_subgraph_tokens_inner(&g, question, depth)
}
fn query_subgraph_tokens_inner(g: &BenchGraph, question: &str, depth: u32) -> u64 {
let terms = bench_query_terms(question);
if terms.is_empty() {
return 0;
}
let mut scored: Vec<(usize, &str)> = g
.nodes
.iter()
.map(|n| {
let label = n.label.to_lowercase();
let score = terms.iter().filter(|t| label.contains(t.as_str())).count();
(score, n.id.as_str())
})
.filter(|(s, _)| *s > 0)
.collect();
scored.sort_by_key(|k| std::cmp::Reverse(k.0));
let start_ids: Vec<&str> = scored.iter().take(3).map(|(_, id)| *id).collect();
if start_ids.is_empty() {
return 0;
}
let mut visited: HashSet<&str> = start_ids.iter().copied().collect();
let mut frontier: VecDeque<&str> = start_ids.iter().copied().collect();
let mut edges_seen: Vec<(&str, &str)> = Vec::new();
for _ in 0..depth {
let current: Vec<&str> = frontier.drain(..).collect();
for n in ¤t {
if let Some(neighbors) = g.adj.get(*n) {
for nb in neighbors {
let nb_s = nb.as_str();
if !visited.contains(nb_s) {
visited.insert(nb_s);
frontier.push_back(nb_s);
edges_seen.push((n, nb_s));
}
}
}
}
}
let mut lines = Vec::new();
for nid in &visited {
if let Some(&idx) = g.node_index.get(*nid) {
let n = &g.nodes[idx];
lines.push(format!(
"NODE {} src={} loc={}",
n.label, n.source_file, n.source_location
));
}
}
for (u, v) in &edges_seen {
if visited.contains(u) && visited.contains(v) {
lines.push(format!("EDGE {} --> {}", u, v));
}
}
if lines.is_empty() {
return 0;
}
estimate_tokens(&lines.join("\n"))
}
#[derive(Debug, Clone)]
pub struct QuestionResult {
pub question: String,
pub query_tokens: u64,
pub reduction: f64,
}
#[derive(Debug, Clone, Default)]
pub struct BenchmarkResult {
pub corpus_tokens: u64,
pub corpus_words: u64,
pub nodes: usize,
pub edges: usize,
pub avg_query_tokens: u64,
pub reduction_ratio: f64,
pub per_question: Vec<QuestionResult>,
pub error: Option<String>,
}
pub fn run_benchmark(
graph_path: &Path,
corpus_words: Option<u64>,
questions: Option<&[&str]>,
max_bytes: Option<u64>,
) -> Result<BenchmarkResult> {
let cap = max_bytes.unwrap_or(MAX_GRAPH_FILE_BYTES);
check_file_size(graph_path, cap)?;
let text = std::fs::read_to_string(graph_path)?;
let data: Value = serde_json::from_str(&text)?;
let g = match BenchGraph::from_json(&data) {
Some(g) => g,
None => {
return Ok(BenchmarkResult {
error: Some("Failed to parse graph JSON".to_string()),
..Default::default()
})
}
};
let node_count = g.node_count();
let edge_count = g.edge_count();
let corpus_words = corpus_words.unwrap_or(node_count as u64 * 50);
let corpus_tokens = corpus_words * 100 / 75;
let qs = questions.unwrap_or(SAMPLE_QUESTIONS);
let mut per_question: Vec<QuestionResult> = Vec::new();
for q in qs {
let qt = query_subgraph_tokens_inner(&g, q, 3);
if qt > 0 {
let reduction = if qt > 0 {
(corpus_tokens as f64 / qt as f64 * 10.0).round() / 10.0
} else {
0.0
};
per_question.push(QuestionResult {
question: q.to_string(),
query_tokens: qt,
reduction,
});
}
}
if per_question.is_empty() {
return Ok(BenchmarkResult {
error: Some(
"No matching nodes found for sample questions. Build the graph first.".to_string(),
),
nodes: node_count,
edges: edge_count,
corpus_tokens,
corpus_words,
..Default::default()
});
}
let avg_query_tokens =
per_question.iter().map(|p| p.query_tokens).sum::<u64>() / per_question.len() as u64;
let reduction_ratio = if avg_query_tokens > 0 {
(corpus_tokens as f64 / avg_query_tokens as f64 * 10.0).round() / 10.0
} else {
0.0
};
Ok(BenchmarkResult {
corpus_tokens,
corpus_words,
nodes: node_count,
edges: edge_count,
avg_query_tokens,
reduction_ratio,
per_question,
error: None,
})
}
pub fn print_benchmark(result: &BenchmarkResult) {
print_benchmark_to(result, &mut std::io::stdout(), true);
}
pub fn print_benchmark_to<W: Write>(result: &BenchmarkResult, out: &mut W, supports_unicode: bool) {
if let Some(ref err) = result.error {
let _ = writeln!(out, "Benchmark error: {err}");
return;
}
let arrow = safe_char("\u{2192}", "->", supports_unicode);
let _ = writeln!(out, "\ncodesynapse token reduction benchmark");
let _ = writeln!(out, "{}", hr(50, supports_unicode));
let _ = writeln!(
out,
" Corpus: {} words {} ~{} tokens (naive)",
result.corpus_words, arrow, result.corpus_tokens
);
let _ = writeln!(
out,
" Graph: {} nodes, {} edges",
result.nodes, result.edges
);
let _ = writeln!(
out,
" Avg query cost: ~{} tokens",
result.avg_query_tokens
);
let _ = writeln!(
out,
" Reduction: {}x fewer tokens per query",
result.reduction_ratio
);
let _ = writeln!(out, "\n Per question:");
for p in &result.per_question {
let q = if p.question.len() > 55 {
&p.question[..55]
} else {
&p.question
};
let _ = writeln!(out, " [{}x] {}", p.reduction, q);
}
let _ = writeln!(out);
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::io::Write;
use tempfile::NamedTempFile;
fn make_graph_value() -> Value {
json!({
"nodes": [
{"id": "n1", "label": "authentication", "source_file": "auth.py", "source_location": "L1", "community": 0},
{"id": "n2", "label": "api_handler", "source_file": "api.py", "source_location": "L5", "community": 0},
{"id": "n3", "label": "main_entry", "source_file": "main.py", "source_location": "L1", "community": 1},
{"id": "n4", "label": "error_handler", "source_file": "errors.py","source_location": "L1", "community": 1},
{"id": "n5", "label": "database_layer", "source_file": "db.py", "source_location": "L1", "community": 2}
],
"links": [
{"source": "n1", "target": "n2", "relation": "calls", "confidence": "INFERRED"},
{"source": "n2", "target": "n3", "relation": "imports", "confidence": "EXTRACTED"},
{"source": "n3", "target": "n4", "relation": "uses", "confidence": "EXTRACTED"},
{"source": "n5", "target": "n2", "relation": "provides","confidence": "EXTRACTED"}
]
})
}
fn write_graph_file(data: &Value) -> NamedTempFile {
let mut f = NamedTempFile::new().unwrap();
write!(f, "{}", serde_json::to_string(data).unwrap()).unwrap();
f
}
#[test]
fn test_query_returns_positive_for_matching_question() {
let g = make_graph_value();
let tokens = query_subgraph_tokens(&g, "how does authentication work", 3);
assert!(tokens > 0, "expected >0 tokens, got {tokens}");
}
#[test]
fn test_query_returns_zero_for_no_match() {
let g = make_graph_value();
let tokens = query_subgraph_tokens(&g, "xyzzy plugh zorkmid", 3);
assert_eq!(tokens, 0);
}
#[test]
fn test_query_bfs_expands_neighbors() {
let g = make_graph_value();
let tokens_deep = query_subgraph_tokens(&g, "authentication", 3);
let tokens_shallow = query_subgraph_tokens(&g, "authentication", 1);
assert!(
tokens_deep >= tokens_shallow,
"deep={tokens_deep} shallow={tokens_shallow}"
);
}
#[test]
fn test_query_keeps_short_non_english_terms() {
let g = json!({
"nodes": [
{"id": "frontend", "label": "\u{524d}\u{7aef}", "source_file": "docs/frontend.md", "source_location": "L1", "community": 0}
],
"links": []
});
let tokens = query_subgraph_tokens(&g, "\u{524d}\u{7aef}", 1);
assert!(tokens > 0, "expected >0 for Chinese query term");
}
#[test]
fn test_run_benchmark_returns_reduction() {
let data = make_graph_value();
let f = write_graph_file(&data);
let result = run_benchmark(f.path(), Some(10_000), None, None).unwrap();
assert!(result.error.is_none());
assert!(
result.reduction_ratio > 1.0,
"reduction_ratio={}",
result.reduction_ratio
);
}
#[test]
fn test_run_benchmark_corpus_tokens_proportional() {
let data = make_graph_value();
let f = write_graph_file(&data);
let r1 = run_benchmark(f.path(), Some(1_000), None, None).unwrap();
let r2 = run_benchmark(f.path(), Some(10_000), None, None).unwrap();
let diff = (r2.corpus_tokens as i64 - r1.corpus_tokens as i64 * 10).unsigned_abs();
assert!(
diff <= r1.corpus_tokens,
"corpus_tokens not proportional: r1={} r2={} diff={}",
r1.corpus_tokens,
r2.corpus_tokens,
diff
);
}
#[test]
fn test_run_benchmark_per_question_list() {
let data = make_graph_value();
let f = write_graph_file(&data);
let qs = &["how does authentication work", "what is the main entry"];
let result = run_benchmark(f.path(), Some(5_000), Some(qs), None).unwrap();
assert!(!result.per_question.is_empty());
for p in &result.per_question {
assert!(!p.question.is_empty());
assert!(p.query_tokens > 0);
assert!(p.reduction > 0.0);
}
}
#[test]
fn test_run_benchmark_estimates_corpus_if_no_words() {
let data = make_graph_value();
let f = write_graph_file(&data);
let result = run_benchmark(f.path(), None, None, None).unwrap();
assert!(result.corpus_words > 0, "corpus_words should be estimated");
}
#[test]
fn test_run_benchmark_error_on_empty_graph() {
let data = json!({"nodes": [], "links": []});
let f = write_graph_file(&data);
let result = run_benchmark(f.path(), Some(1_000), None, None).unwrap();
assert!(result.error.is_some(), "expected error for empty graph");
}
#[test]
fn test_run_benchmark_includes_node_edge_counts() {
let data = make_graph_value();
let f = write_graph_file(&data);
let result = run_benchmark(f.path(), Some(5_000), None, None).unwrap();
assert_eq!(result.nodes, 5, "nodes={}", result.nodes);
assert_eq!(result.edges, 4, "edges={}", result.edges);
}
#[test]
fn test_print_benchmark_no_crash() {
let data = make_graph_value();
let f = write_graph_file(&data);
let result = run_benchmark(f.path(), Some(5_000), None, None).unwrap();
let mut buf = Vec::new();
print_benchmark_to(&result, &mut buf, true);
let out = String::from_utf8(buf).unwrap();
assert!(out.to_lowercase().contains("reduction"), "out={out}");
assert!(out.contains('x'), "out={out}");
}
#[test]
fn test_print_benchmark_error_message() {
let result = BenchmarkResult {
error: Some("test error message".to_string()),
..Default::default()
};
let mut buf = Vec::new();
print_benchmark_to(&result, &mut buf, true);
let out = String::from_utf8(buf).unwrap();
assert!(out.contains("test error message"), "out={out}");
}
#[test]
fn test_safe_char_unicode_mode() {
assert_eq!(safe_char("\u{2192}", "->", true), "\u{2192}");
assert_eq!(hr(5, true), "\u{2500}".repeat(5));
}
#[test]
fn test_safe_char_ascii_mode() {
assert_eq!(safe_char("\u{2192}", "->", false), "->");
assert_eq!(hr(5, false), "-----");
}
#[test]
fn test_print_benchmark_ascii_only_output() {
let data = make_graph_value();
let f = write_graph_file(&data);
let result = run_benchmark(f.path(), Some(5_000), None, None).unwrap();
let mut buf = Vec::new();
print_benchmark_to(&result, &mut buf, false);
let out = String::from_utf8(buf).unwrap();
assert!(out.to_lowercase().contains("reduction"), "out={out}");
assert!(!out.contains('\u{2500}'), "should not contain ─");
assert!(!out.contains('\u{2192}'), "should not contain →");
}
#[test]
fn test_run_benchmark_rejects_oversized_graph() {
let data = make_graph_value();
let f = write_graph_file(&data);
let result = run_benchmark(f.path(), Some(5_000), None, Some(8));
assert!(result.is_err(), "expected Err for oversized graph");
let msg = result.unwrap_err().to_string();
assert!(msg.contains("exceeds"), "msg={msg}");
}
}