use std::collections::{HashMap, HashSet};
use std::time::Instant;
use crate::search::summarize::{self, ReadChunk};
use crate::types::{Chunk, ChunkType, SearchResult};
const CONTEXT_BUDGET_BYTES: usize = 32_768;
const CONFIDENCE_NOISE_FLOOR: f32 = 0.08;
const CONFIDENCE_LOW: f32 = 0.20;
#[derive(Debug, Clone)]
pub struct DeepResult {
pub answer: String,
pub sources: Vec<DeepSource>,
pub metrics: DeepMetrics,
pub confidence: f32,
}
#[derive(Debug, Clone)]
pub struct DeepSource {
pub file: String,
pub start_line: u32,
pub end_line: u32,
pub name: Option<String>,
pub kind: Option<String>,
pub content: String,
pub also_matched_via: Vec<String>,
}
#[derive(Debug, Clone, Default)]
pub struct DeepMetrics {
pub search_ms: u64,
pub triage_ms: u64,
pub graph_ms: u64,
pub read_ms: u64,
pub summarize_ms: u64,
pub total_ms: u64,
pub chunks_searched: usize,
pub chunks_read: usize,
pub confidence_zone: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PreferredKind {
FnMethod,
StructClassType,
Neutral,
}
fn classify_query_preference(query: &str) -> PreferredKind {
let q = query.to_lowercase();
if q.contains("how does")
|| q.contains("how is")
|| q.contains("how do")
|| q.contains("how are")
{
PreferredKind::FnMethod
} else if q.contains("what is") || q.contains("what are") {
PreferredKind::StructClassType
} else {
PreferredKind::Neutral
}
}
fn compute_docstring_boost(query: &str, chunk: &Chunk) -> f32 {
let ChunkType::AstNode {
structured_meta: Some(meta),
..
} = &chunk.chunk_type
else {
return 0.0;
};
let terms = summarize::extract_query_terms(query);
if terms.is_empty() {
return 0.0;
}
let mut text = String::new();
if let Some(ref doc) = meta.docstring {
text.push_str(doc);
text.push(' ');
}
if !meta.nl_summary.is_empty() {
text.push_str(&meta.nl_summary);
}
if text.is_empty() {
return 0.0;
}
let lower = text.to_lowercase();
let match_count = terms.iter().filter(|t| lower.contains(t.as_str())).count();
(0.15 * match_count as f32).min(0.6)
}
fn kind_matches_preference(kind: Option<&str>, pref: PreferredKind) -> bool {
match pref {
PreferredKind::Neutral => false,
PreferredKind::FnMethod => {
matches!(kind, Some("fn" | "method" | "function"))
}
PreferredKind::StructClassType => {
matches!(kind, Some("struct" | "class" | "interface" | "enum"))
}
}
}
fn score_triage_candidates(
query: &str,
pref: PreferredKind,
results: &[SearchResult],
max_possible: f32,
apply_noise_floor: bool,
) -> Vec<(usize, f32, String, u32, u32)> {
let mut scored: Vec<(usize, f32, String, u32, u32)> = Vec::new();
for (idx, result) in results.iter().enumerate() {
let file = result.chunk.file_path.display().to_string();
let start = result.chunk.start_line;
let end = result.chunk.end_line;
if matches!(result.chunk.chunk_type, ChunkType::TextWindow { .. }) && result.score < 0.1 {
continue;
}
if apply_noise_floor && max_possible > 0.0 {
let normalized = result.score / max_possible;
if normalized < CONFIDENCE_NOISE_FLOOR {
continue;
}
}
let kind = match &result.chunk.chunk_type {
ChunkType::AstNode { kind, .. } => Some(kind.to_string()),
_ => None,
};
let kind_boost = if kind_matches_preference(kind.as_deref(), pref) {
0.1
} else {
0.0
};
let docstring_boost = compute_docstring_boost(query, &result.chunk);
let channel_count = [
result.score_dense > 0.0,
result.score_sparse > 0.0,
result.score_exact > 0.0,
]
.iter()
.filter(|&&b| b)
.count();
let consensus_bonus = match channel_count {
3 => 0.15,
2 => 0.05,
_ => 0.0,
};
let adjusted_score = result.score + kind_boost + docstring_boost + consensus_bonus;
scored.push((idx, adjusted_score, file, start, end));
}
scored
}
fn triage_results(
query: &str,
results: &[SearchResult],
max_chunks: usize,
max_possible: f32,
) -> Vec<usize> {
let pref = classify_query_preference(query);
let mut scored = score_triage_candidates(query, pref, results, max_possible, true);
if scored.is_empty() && !results.is_empty() {
scored = score_triage_candidates(query, pref, results, max_possible, false);
}
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let mut selected_ranges: Vec<(String, u32, u32)> = Vec::new();
let mut file_counts: HashMap<String, usize> = HashMap::new();
let mut indices: Vec<usize> = Vec::new();
for (idx, _score, file, start, end) in &scored {
if file_counts.get(file).copied().unwrap_or(0) >= 3 {
continue;
}
let mut overlaps = false;
for (sel_file, sel_start, sel_end) in &selected_ranges {
if sel_file != file {
continue;
}
let overlap_start = start.max(sel_start);
let overlap_end = end.min(sel_end);
if overlap_end >= overlap_start {
let overlap_lines = (overlap_end.saturating_sub(*overlap_start) + 1) as f32;
let my_lines = (end.saturating_sub(*start) + 1).max(1) as f32;
if overlap_lines / my_lines > 0.5 {
overlaps = true;
break;
}
}
}
if overlaps {
continue;
}
selected_ranges.push((file.clone(), *start, *end));
*file_counts.entry(file.clone()).or_insert(0) += 1;
indices.push(*idx);
if indices.len() >= max_chunks {
break;
}
}
indices.sort_unstable();
indices
}
#[derive(Debug, Clone)]
struct DedupedSelection<'a> {
result: &'a SearchResult,
also_matched_via: Vec<String>,
}
fn channels_label(result: &SearchResult) -> String {
let mut parts: Vec<&'static str> = Vec::new();
if result.score_dense > 0.0 {
parts.push("dense");
}
if result.score_sparse > 0.0 {
parts.push("sparse");
}
if result.score_exact > 0.0 {
parts.push("exact");
}
if parts.is_empty() {
format!("{:?}", result.source).to_lowercase()
} else {
parts.join("+")
}
}
fn dedup_by_range<'a>(selected: &[&'a SearchResult]) -> Vec<DedupedSelection<'a>> {
use std::collections::HashMap;
const MAX_LABELS: usize = 4;
let mut order: Vec<(String, u32, u32)> = Vec::new();
let mut by_key: HashMap<(String, u32, u32), (usize, Vec<String>)> = HashMap::new();
for (idx, result) in selected.iter().enumerate() {
let key = (
result.chunk.file_path.display().to_string(),
result.chunk.start_line,
result.chunk.end_line,
);
if let Some(entry) = by_key.get_mut(&key) {
let (best_idx, ref mut also) = *entry;
let best_score = selected[best_idx].score;
let this_score = result.score;
if this_score > best_score {
also.push(channels_label(selected[best_idx]));
*entry = (idx, std::mem::take(also));
} else {
also.push(channels_label(result));
}
} else {
order.push(key.clone());
by_key.insert(key, (idx, Vec::new()));
}
}
order
.into_iter()
.map(|key| {
let (best_idx, mut also) = by_key.remove(&key).expect("key inserted above");
also.sort();
also.dedup();
also.truncate(MAX_LABELS);
DedupedSelection {
result: selected[best_idx],
also_matched_via: also,
}
})
.collect()
}
fn append_dedup_footer(answer: String, sources: &[DeepSource]) -> String {
use std::fmt::Write as _;
let lines: Vec<String> = sources
.iter()
.filter(|s| !s.also_matched_via.is_empty())
.map(|s| {
let mut line = format!("- {}:{}-{}", s.file, s.start_line, s.end_line);
if let Some(name) = &s.name {
line.push(' ');
line.push_str(name);
}
let _ = write!(
line,
" (also matched via: {})",
s.also_matched_via.join(", ")
);
line
})
.collect();
if lines.is_empty() {
return answer;
}
let mut out = answer;
if !out.is_empty() {
out.push_str("\n\n");
}
out.push_str("Cross-source dedup:");
for line in &lines {
out.push('\n');
out.push_str(line);
}
out
}
fn expand_with_graph(
store: &crate::index::storage::ChunkStore,
selected_ids: &[u64],
_query: &str,
budget: usize,
) -> anyhow::Result<Vec<u64>> {
if selected_ids.is_empty() || budget == 0 {
return Ok(Vec::new());
}
let selected_set: HashSet<u64> = selected_ids.iter().copied().collect();
let mut candidate_counts: HashMap<u64, usize> = HashMap::new();
let outgoing_calls = store.get_call_edges_from(selected_ids)?;
for (_, callee_id) in outgoing_calls {
if !selected_set.contains(&callee_id) {
*candidate_counts.entry(callee_id).or_insert(0) += 1;
}
}
let incoming_calls = store.get_call_edges_to(selected_ids)?;
for (_, caller_id) in incoming_calls {
if !selected_set.contains(&caller_id) {
*candidate_counts.entry(caller_id).or_insert(0) += 1;
}
}
let type_def_edges = store.get_type_ref_edges_to_defs(selected_ids)?;
for (_, usage_id) in type_def_edges {
if !selected_set.contains(&usage_id) {
*candidate_counts.entry(usage_id).or_insert(0) += 1;
}
}
let type_usage_edges = store.get_type_ref_edges_from_usages(selected_ids)?;
for (_, def_id) in type_usage_edges {
if !selected_set.contains(&def_id) {
*candidate_counts.entry(def_id).or_insert(0) += 1;
}
}
let hierarchy_edges = store.get_hierarchy_edges_for(selected_ids)?;
for (_, related_id) in hierarchy_edges {
if !selected_set.contains(&related_id) {
*candidate_counts.entry(related_id).or_insert(0) += 1;
}
}
if candidate_counts.is_empty() {
return Ok(Vec::new());
}
let mut scored: Vec<(u64, f32)> = candidate_counts
.into_iter()
.map(|(id, count)| {
let edge_score = 0.5 + (count.saturating_sub(1)) as f32 * 0.2;
(id, edge_score)
})
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let max_additional = (budget.saturating_sub(selected_ids.len())).min(7);
let additional: Vec<u64> = scored
.into_iter()
.take(max_additional)
.map(|(id, _)| id)
.collect();
Ok(additional)
}
fn parse_name_list(text: &str) -> Vec<String> {
text.split(',')
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_string)
.collect()
}
struct GraphContext {
callers: Vec<String>,
callees: Vec<String>,
type_refs: Vec<String>,
}
#[allow(clippy::similar_names)]
fn parse_graph_context(nl_summary: &str) -> GraphContext {
let callees = nl_summary
.find("; calls ")
.map(|pos| {
let after = &nl_summary[pos + "; calls ".len()..];
let end = after.find(';').unwrap_or(after.len());
parse_name_list(&after[..end])
})
.unwrap_or_default();
let callers = nl_summary
.find("; called by ")
.map(|pos| {
let after = &nl_summary[pos + "; called by ".len()..];
let end = after.find(';').unwrap_or(after.len());
parse_name_list(&after[..end])
})
.unwrap_or_default();
let type_refs = nl_summary
.find("; uses types ")
.map(|pos| {
let after = &nl_summary[pos + "; uses types ".len()..];
let end = after.find(';').unwrap_or(after.len());
parse_name_list(&after[..end])
})
.unwrap_or_default();
GraphContext {
callers,
callees,
type_refs,
}
}
fn chunk_to_read_chunk(chunk: &crate::types::Chunk) -> ReadChunk {
let full_path = chunk.file_path.display().to_string();
match &chunk.chunk_type {
ChunkType::AstNode {
name,
kind,
structured_meta,
..
} => {
let summary = structured_meta.as_ref().and_then(|meta| {
if meta.nl_summary.is_empty() {
None
} else {
Some(meta.nl_summary.clone())
}
});
let docstring = structured_meta
.as_ref()
.and_then(|meta| meta.docstring.clone());
let graph_ctx = structured_meta.as_ref().map_or(
GraphContext {
callers: Vec::new(),
callees: Vec::new(),
type_refs: Vec::new(),
},
|meta| parse_graph_context(&meta.nl_summary),
);
ReadChunk {
file: full_path.clone(),
start_line: chunk.start_line,
end_line: chunk.end_line,
name: Some(name.clone()),
kind: Some(kind.to_string()),
content: chunk.content.clone(),
summary,
docstring,
callers: graph_ctx.callers,
callees: graph_ctx.callees,
type_refs: graph_ctx.type_refs,
full_path,
}
}
ChunkType::TextWindow { .. } | ChunkType::PdfPage { .. } => ReadChunk {
file: full_path.clone(),
start_line: chunk.start_line,
end_line: chunk.end_line,
name: None,
kind: None,
content: chunk.content.clone(),
summary: None,
docstring: None,
callers: Vec::new(),
callees: Vec::new(),
type_refs: Vec::new(),
full_path,
},
}
}
fn try_symbol_shortcut(
searcher: &crate::search::hybrid::HybridSearcher,
query: &str,
start: Instant,
) -> anyhow::Result<Option<DeepResult>> {
let symbol_hits = searcher.with_store(|store| store.lookup_symbol_exact(query))?;
if symbol_hits.is_empty() {
return Ok(None);
}
let chunk_ids: Vec<u64> = symbol_hits.iter().map(|(id, _, _)| *id as u64).collect();
let chunks = searcher.with_store(|store| store.get_chunks(&chunk_ids))?;
if chunks.is_empty() {
return Ok(None);
}
let read_chunks: Vec<ReadChunk> = chunks.iter().map(chunk_to_read_chunk).collect();
let answer = summarize::extractive_summarize(query, &read_chunks, true);
let sources: Vec<DeepSource> = chunks
.iter()
.map(|chunk| {
let (name, kind) = match &chunk.chunk_type {
ChunkType::AstNode { name, kind, .. } => {
(Some(name.clone()), Some(kind.to_string()))
}
_ => (None, None),
};
DeepSource {
file: chunk.file_path.display().to_string(),
start_line: chunk.start_line,
end_line: chunk.end_line,
name,
kind,
content: chunk.content.clone(),
also_matched_via: Vec::new(),
}
})
.collect();
let total_ms = start.elapsed().as_millis() as u64;
Ok(Some(DeepResult {
answer,
sources,
metrics: DeepMetrics {
search_ms: 0,
triage_ms: 0,
graph_ms: 0,
read_ms: total_ms,
summarize_ms: 0,
total_ms,
chunks_searched: 0,
chunks_read: chunks.len(),
confidence_zone: "high".to_string(), },
confidence: 1.0, }))
}
fn decompose_query(query: &str) -> Option<Vec<String>> {
let q = query.to_lowercase();
if let Some(rest) = q.strip_prefix("how does ")
&& let Some(pos) = rest.find(" use ").or_else(|| rest.find(" interact with "))
{
let x = rest[..pos].trim().to_string();
let sep_len = if rest[pos..].starts_with(" use ") {
5
} else {
14
};
let y = rest[pos + sep_len..].trim().to_string();
if !x.is_empty() && !y.is_empty() && x.len() > 2 && y.len() > 2 {
return Some(vec![x, y]);
}
}
for sep in [" across ", " from ", " between "] {
if let Some(main_pos) = q.find(sep) {
let subject = q[..main_pos].trim();
let rest = &q[main_pos + sep.len()..];
if let Some(and_pos) = rest.find(" and ").or_else(|| rest.find(" to ")) {
let part_a = rest[..and_pos].trim();
let and_len = if rest[and_pos..].starts_with(" and ") {
5
} else {
4
};
let part_b = rest[and_pos + and_len..].trim();
if !part_a.is_empty() && !part_b.is_empty() {
return Some(vec![
format!("{subject} {part_a}"),
format!("{subject} {part_b}"),
]);
}
}
}
}
for prefix in ["add ", "implement ", "where would i add ", "where to add "] {
if let Some(stripped) = q.strip_prefix(prefix) {
let feature = stripped.trim().trim_end_matches(" to this project").trim();
if feature.len() > 3 {
return Some(vec![
format!("existing {feature} implementation patterns"),
"main entry points and initialization".to_string(),
]);
}
}
}
None
}
pub type ProgressFn<'a> = &'a dyn Fn(u32, u32, &str);
pub fn deep_search(
searcher: &crate::search::hybrid::HybridSearcher,
query: &str,
max_results: usize,
use_graph: bool,
) -> anyhow::Result<DeepResult> {
deep_search_inner(searcher, query, max_results, use_graph, None)
}
pub fn deep_search_with_progress(
searcher: &crate::search::hybrid::HybridSearcher,
query: &str,
max_results: usize,
use_graph: bool,
on_progress: ProgressFn<'_>,
) -> anyhow::Result<DeepResult> {
deep_search_inner(searcher, query, max_results, use_graph, Some(on_progress))
}
fn deep_search_inner(
searcher: &crate::search::hybrid::HybridSearcher,
query: &str,
max_results: usize,
use_graph: bool,
on_progress: Option<ProgressFn<'_>>,
) -> anyhow::Result<DeepResult> {
let total_start = Instant::now();
let mut metrics = DeepMetrics::default();
let report = |step, total, msg| {
if let Some(cb) = &on_progress {
cb(step, total, msg);
}
};
let query_type_for_shortcut = crate::search::query_classifier::classify(query);
if query_type_for_shortcut == crate::search::query_classifier::QueryType::Identifier
&& let Some(shortcut_result) = try_symbol_shortcut(searcher, query, total_start)?
{
return Ok(shortcut_result);
}
report(0, 5, "Searching code index...");
let search_start = Instant::now();
let effective_max = max_results.max(20);
let query_type_pre = crate::search::query_classifier::classify(query);
let should_decompose = matches!(
query_type_pre,
crate::search::query_classifier::QueryType::Semantic
| crate::search::query_classifier::QueryType::Mixed
);
let output = if should_decompose && let Some(sub_queries) = decompose_query(query) {
let mut merged_results: Vec<crate::types::SearchResult> = Vec::new();
let mut seen_chunk_ids: HashSet<u64> = HashSet::new();
for sub_q in &sub_queries {
let sub_search_query = crate::search::SearchQuery::new(sub_q)
.max_results(effective_max)
.code_only();
if let Ok(sub_output) = searcher.search(&sub_search_query) {
for result in sub_output.results {
if seen_chunk_ids.insert(result.chunk.id) {
merged_results.push(result);
}
}
}
}
merged_results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
crate::search::SearchOutput {
results: merged_results,
metrics: crate::search::SearchMetrics {
total_ms: 0,
dense_ms: None,
sparse_ms: None,
exact_ms: None,
fusion_ms: None,
rerank_ms: None,
dense_count: 0,
sparse_count: 0,
exact_count: 0,
fused_count: 0,
result_count: 0,
query_type: "decomposed".to_string(),
response_bytes: None,
},
}
} else {
let search_query = crate::search::SearchQuery::new(query)
.max_results(effective_max)
.code_only();
searcher.search(&search_query)?
};
metrics.search_ms = search_start.elapsed().as_millis() as u64;
metrics.chunks_searched = output.results.len();
let query_type = crate::search::query_classifier::classify(query);
let max_possible = query_type.triple_fusion_weights().max_possible();
let confidence = if output.results.is_empty() {
0.0
} else {
let top_scores: Vec<f32> = output
.results
.iter()
.take(3)
.map(|r| (r.score / max_possible).min(1.0))
.collect();
top_scores.iter().sum::<f32>() / top_scores.len() as f32
};
let confidence_zone = if output.results.is_empty() {
"no_results"
} else if confidence >= CONFIDENCE_LOW {
"high"
} else if confidence >= CONFIDENCE_NOISE_FLOOR {
"medium"
} else {
"low"
};
report(1, 5, "Triaging results...");
let triage_start = Instant::now();
let triage_cap = if confidence > 0.7 {
5 } else if confidence > 0.3 {
8 } else {
14 };
let selected_indices = triage_results(query, &output.results, triage_cap, max_possible);
let selected_results: Vec<&SearchResult> = selected_indices
.iter()
.map(|&i| &output.results[i])
.collect();
metrics.triage_ms = triage_start.elapsed().as_millis() as u64;
let deduped_selected: Vec<DedupedSelection<'_>> = dedup_by_range(&selected_results);
let dedup_lookup: HashMap<(String, u32, u32), Vec<String>> = deduped_selected
.iter()
.filter(|d| !d.also_matched_via.is_empty())
.map(|d| {
(
(
d.result.chunk.file_path.display().to_string(),
d.result.chunk.start_line,
d.result.chunk.end_line,
),
d.also_matched_via.clone(),
)
})
.collect();
let selected_results: Vec<&SearchResult> = deduped_selected.iter().map(|d| d.result).collect();
let selected_ids: Vec<u64> = selected_results.iter().map(|r| r.chunk.id).collect();
report(2, 5, "Expanding via call graph...");
let graph_start = Instant::now();
let additional_ids: Vec<u64> = if use_graph && !selected_ids.is_empty() {
searcher.with_store(|store| expand_with_graph(store, &selected_ids, query, 12))?
} else {
Vec::new()
};
metrics.graph_ms = graph_start.elapsed().as_millis() as u64;
report(3, 5, "Reading full chunk content...");
let read_start = Instant::now();
let extra_chunks: Vec<crate::types::Chunk> = if additional_ids.is_empty() {
Vec::new()
} else {
searcher.with_store(|store| store.get_chunks(&additional_ids))?
};
let mut combined_chunks: Vec<crate::types::Chunk> =
selected_results.iter().map(|r| r.chunk.clone()).collect();
let existing_ids: HashSet<u64> = combined_chunks.iter().map(|c| c.id).collect();
for chunk in extra_chunks {
if !existing_ids.contains(&chunk.id) {
combined_chunks.push(chunk);
}
}
metrics.read_ms = read_start.elapsed().as_millis() as u64;
let mut total_bytes: usize = 0;
let mut budget_chunks: Vec<crate::types::Chunk> = Vec::new();
for chunk in combined_chunks {
let chunk_bytes = chunk.content.len();
if total_bytes + chunk_bytes > CONTEXT_BUDGET_BYTES && !budget_chunks.is_empty() {
break;
}
total_bytes += chunk_bytes;
budget_chunks.push(chunk);
}
let combined_chunks = budget_chunks;
metrics.chunks_read = combined_chunks.len();
report(4, 5, "Synthesizing answer...");
let summarize_start = Instant::now();
let read_chunks: Vec<ReadChunk> = combined_chunks.iter().map(chunk_to_read_chunk).collect();
let answer = summarize::extractive_summarize(query, &read_chunks, true);
metrics.summarize_ms = summarize_start.elapsed().as_millis() as u64;
metrics.total_ms = total_start.elapsed().as_millis() as u64;
report(5, 5, "Complete");
let sources: Vec<DeepSource> = combined_chunks
.iter()
.map(|chunk| {
let (name, kind) = match &chunk.chunk_type {
ChunkType::AstNode { name, kind, .. } => {
(Some(name.clone()), Some(kind.to_string()))
}
_ => (None, None),
};
let file = chunk.file_path.display().to_string();
let also_matched_via = if dedup_lookup.is_empty() {
Vec::new()
} else {
dedup_lookup
.get(&(file.clone(), chunk.start_line, chunk.end_line))
.cloned()
.unwrap_or_default()
};
DeepSource {
file,
start_line: chunk.start_line,
end_line: chunk.end_line,
name,
kind,
content: chunk.content.clone(),
also_matched_via,
}
})
.collect();
let answer = append_dedup_footer(answer, &sources);
let answer = if confidence_zone == "low" || confidence_zone == "no_results" {
if answer.is_empty() || answer == "No relevant code found for this query." {
"No relevant code found for this query.".to_string()
} else {
format!(
"{answer}\n\n[Low confidence: results may not be relevant to the query. Consider rephrasing or checking if the codebase covers this topic.]"
)
}
} else {
answer
};
metrics.confidence_zone = confidence_zone.to_string();
Ok(DeepResult {
answer,
sources,
metrics,
confidence,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{AstNodeKind, Chunk, ChunkType, SearchResult, SearchSource};
use std::path::PathBuf;
fn make_ast_result(file: &str, name: &str, start: u32, end: u32, score: f32) -> SearchResult {
SearchResult {
chunk: Chunk {
id: u64::from(start) * 1000 + u64::from(end),
file_path: PathBuf::from(file),
start_line: start,
end_line: end,
content: format!("pub fn {name}(query: &str) -> Vec<Result> {{ todo!() }}"),
chunk_type: ChunkType::AstNode {
name: name.to_string(),
kind: AstNodeKind::Function,
language: "rust".to_string(),
structured_meta: None,
},
},
score,
source: SearchSource::Hybrid,
score_dense: 0.0,
score_sparse: 0.0,
score_exact: 0.0,
confidence: crate::types::Confidence::Inferred,
confidence_score: 0.0,
}
}
#[test]
fn test_triage_basic() {
let results = vec![
make_ast_result("a.rs", "search_index", 1, 20, 0.9),
make_ast_result("b.rs", "index_file", 1, 15, 0.7),
make_ast_result("c.rs", "other_fn", 1, 10, 0.5),
];
let indices = triage_results("search index", &results, 8, 5.0);
assert!(!indices.is_empty(), "Should select at least one result");
assert!(indices.len() <= 8);
}
#[test]
fn test_triage_per_file_cap() {
let results: Vec<SearchResult> = (0..6u32)
.map(|i| {
make_ast_result(
"same_file.rs",
&format!("fn_{i}"),
i * 20 + 1,
i * 20 + 19,
0.8,
)
})
.collect();
let indices = triage_results("search query", &results, 8, 5.0);
let from_same_file = indices
.iter()
.filter(|&&i| results[i].chunk.file_path.display().to_string() == "same_file.rs")
.count();
assert!(
from_same_file <= 3,
"Per-file cap of 3 should be enforced, got {from_same_file}"
);
}
#[test]
fn test_triage_low_score_text_window_skipped() {
let tw_result = SearchResult {
chunk: Chunk {
id: 999,
file_path: PathBuf::from("doc.md"),
start_line: 1,
end_line: 10,
content: "some text".to_string(),
chunk_type: ChunkType::TextWindow { window_index: 0 },
},
score: 0.05, source: SearchSource::Sparse,
score_dense: 0.0,
score_sparse: 0.0,
score_exact: 0.0,
confidence: crate::types::Confidence::Inferred,
confidence_score: 0.0,
};
let results = vec![tw_result];
let indices = triage_results("some text", &results, 8, 100.0);
assert!(
indices.is_empty(),
"Low-score TextWindow should be filtered out"
);
}
#[test]
fn test_chunk_to_read_chunk_ast_node() {
let chunk = Chunk {
id: 1,
file_path: PathBuf::from("src/lib.rs"),
start_line: 10,
end_line: 30,
content: "pub fn authenticate(token: &str) -> bool { true }".to_string(),
chunk_type: ChunkType::AstNode {
name: "authenticate".to_string(),
kind: AstNodeKind::Function,
language: "rust".to_string(),
structured_meta: None,
},
};
let rc = chunk_to_read_chunk(&chunk);
assert_eq!(rc.name.as_deref(), Some("authenticate"));
assert_eq!(rc.kind.as_deref(), Some("fn"));
assert_eq!(rc.start_line, 10);
assert_eq!(rc.end_line, 30);
assert!(rc.summary.is_none());
}
#[test]
fn test_chunk_to_read_chunk_text_window() {
let chunk = Chunk {
id: 2,
file_path: PathBuf::from("README.md"),
start_line: 1,
end_line: 5,
content: "Overview of the project".to_string(),
chunk_type: ChunkType::TextWindow { window_index: 0 },
};
let rc = chunk_to_read_chunk(&chunk);
assert!(rc.name.is_none());
assert!(rc.kind.is_none());
assert!(rc.summary.is_none());
}
#[test]
fn test_expand_with_graph_empty_selected() {
let result: Vec<u64> = Vec::new(); assert!(result.is_empty());
}
#[test]
fn test_classify_query_preference_how_does() {
assert_eq!(
classify_query_preference("how does authentication work"),
PreferredKind::FnMethod
);
}
#[test]
fn test_classify_query_preference_what_is() {
assert_eq!(
classify_query_preference("what is the ChunkStore struct"),
PreferredKind::StructClassType
);
}
#[test]
fn test_classify_query_preference_neutral() {
assert_eq!(
classify_query_preference("search index BM25"),
PreferredKind::Neutral
);
}
fn make_chunk_with_docstring(docstring: &str, nl_summary: &str) -> Chunk {
use crate::chunking::structured_meta::StructuredChunkMeta;
Chunk {
id: 1,
file_path: PathBuf::from("src/search.rs"),
start_line: 1,
end_line: 20,
content: "fn placeholder() {}".to_string(),
chunk_type: ChunkType::AstNode {
name: "placeholder".to_string(),
kind: AstNodeKind::Function,
language: "rust".to_string(),
structured_meta: Some(Box::new(StructuredChunkMeta {
docstring: if docstring.is_empty() {
None
} else {
Some(docstring.to_string())
},
nl_summary: nl_summary.to_string(),
..Default::default()
})),
},
}
}
#[test]
fn test_docstring_boost_matching_terms() {
let chunk = make_chunk_with_docstring("Run the full deep search pipeline for a query.", "");
let boost = compute_docstring_boost("how does the deep search pipeline work", &chunk);
assert!(
boost > 0.3,
"Expected boost > 0.3 for 3 matching terms, got {boost}"
);
assert!(boost <= 0.6, "Boost should be capped at 0.6, got {boost}");
}
#[test]
fn test_docstring_boost_no_meta() {
let chunk = Chunk {
id: 2,
file_path: PathBuf::from("lib.rs"),
start_line: 1,
end_line: 5,
content: "fn foo() {}".to_string(),
chunk_type: ChunkType::TextWindow { window_index: 0 },
};
let boost = compute_docstring_boost("deep search", &chunk);
assert_eq!(boost, 0.0, "TextWindow chunks should get no boost");
}
#[test]
fn test_docstring_boost_no_matching_terms() {
let chunk =
make_chunk_with_docstring("Authenticate user credentials against the database.", "");
let boost = compute_docstring_boost("deep search pipeline", &chunk);
assert_eq!(boost, 0.0, "No matching terms should produce zero boost");
}
#[test]
fn test_docstring_boost_capped_at_0_6() {
let chunk = make_chunk_with_docstring(
"deep search pipeline triage graph expand read summarize query",
"deep search pipeline triage graph expand read summarize query",
);
let boost = compute_docstring_boost(
"deep search pipeline triage graph expand read summarize query",
&chunk,
);
assert!(
(boost - 0.6).abs() < f32::EPSILON,
"Boost should be capped at 0.6, got {boost}"
);
}
#[test]
fn test_docstring_boost_from_nl_summary() {
let chunk = make_chunk_with_docstring("", "Run the full deep search pipeline");
let boost = compute_docstring_boost("deep search pipeline", &chunk);
assert!(
boost > 0.3,
"NL summary should also contribute to boost, got {boost}"
);
}
#[test]
fn test_triage_docstring_boost_reranks() {
let mut results = vec![
make_ast_result("a.rs", "classify_query_preference", 1, 20, 0.8),
make_ast_result("b.rs", "deep_search_inner", 1, 20, 0.8),
];
if let ChunkType::AstNode {
structured_meta, ..
} = &mut results[1].chunk.chunk_type
{
use crate::chunking::structured_meta::StructuredChunkMeta;
*structured_meta = Some(Box::new(StructuredChunkMeta {
docstring: Some("Run the full deep search pipeline for a query.".to_string()),
..Default::default()
}));
}
let indices = triage_results("how does the deep search pipeline work", &results, 8, 5.0);
assert_eq!(indices, vec![0, 1]); let boost_a =
compute_docstring_boost("how does the deep search pipeline work", &results[0].chunk);
let boost_b =
compute_docstring_boost("how does the deep search pipeline work", &results[1].chunk);
assert!(
boost_b > boost_a,
"deep_search_inner should have higher docstring boost ({boost_b}) than classify ({boost_a})"
);
}
#[test]
fn test_confidence_high_score_results() {
let max_possible = 1.7_f32;
let scores = [1.2_f32, 1.0, 0.9];
let top_scores: Vec<f32> = scores
.iter()
.take(3)
.map(|&s| (s / max_possible).min(1.0))
.collect();
let confidence = top_scores.iter().sum::<f32>() / top_scores.len() as f32;
let zone = if confidence >= CONFIDENCE_LOW {
"high"
} else if confidence >= CONFIDENCE_NOISE_FLOOR {
"medium"
} else {
"low"
};
assert_eq!(zone, "high");
assert!(confidence > CONFIDENCE_LOW);
}
#[test]
fn test_confidence_low_score_results() {
let max_possible = 5.8_f32;
let scores = [0.2_f32, 0.15, 0.1];
let top_scores: Vec<f32> = scores
.iter()
.take(3)
.map(|&s| (s / max_possible).min(1.0))
.collect();
let confidence = top_scores.iter().sum::<f32>() / top_scores.len() as f32;
let zone = if confidence >= CONFIDENCE_LOW {
"high"
} else if confidence >= CONFIDENCE_NOISE_FLOOR {
"medium"
} else {
"low"
};
assert_eq!(zone, "low");
assert!(confidence < CONFIDENCE_NOISE_FLOOR);
}
#[test]
fn test_confidence_empty_results() {
let confidence = 0.0_f32;
let zone = "no_results";
assert_eq!(confidence, 0.0);
assert_eq!(zone, "no_results");
}
#[test]
fn test_triage_noise_floor_filter() {
let results = vec![
make_ast_result("a.rs", "good_fn", 1, 20, 0.5), make_ast_result("b.rs", "noise_fn", 1, 20, 0.05), make_ast_result("c.rs", "barely_noise", 1, 20, 0.07), ];
let indices = triage_results("test query", &results, 8, 1.0);
assert_eq!(indices.len(), 1);
assert_eq!(indices[0], 0);
}
#[test]
fn test_triage_noise_floor_rescue_on_total_filter() {
let results = vec![
make_ast_result("a.rs", "fn_a", 1, 20, 0.04), make_ast_result("b.rs", "fn_b", 1, 20, 0.05), make_ast_result("c.rs", "fn_c", 1, 20, 0.06), ];
let indices = triage_results("test query", &results, 8, 1.0);
assert_eq!(
indices.len(),
3,
"Rescue pass should restore all candidates when noise floor would have \
dropped everything; got {indices:?}"
);
assert_eq!(indices, vec![0, 1, 2]);
}
#[test]
fn test_triage_rescue_respects_max_chunks() {
let results: Vec<SearchResult> = (0..5u32)
.map(|i| make_ast_result(&format!("f{i}.rs"), &format!("fn_{i}"), 1, 20, 0.05))
.collect();
let indices = triage_results("test query", &results, 2, 1.0);
assert_eq!(
indices.len(),
2,
"Rescue pass must still honor max_chunks; got {indices:?}"
);
}
#[test]
fn test_triage_rescue_skips_when_input_empty() {
let indices = triage_results("test query", &[], 8, 1.0);
assert!(indices.is_empty());
}
#[test]
fn test_symbol_shortcut_identifier_query_classified() {
use crate::search::query_classifier;
let qt = query_classifier::classify("getUserById");
assert_eq!(qt, query_classifier::QueryType::Identifier);
let result = DeepResult {
answer: "Found definition of getUserById".to_string(),
sources: vec![DeepSource {
file: "src/users.rs".to_string(),
start_line: 10,
end_line: 30,
name: Some("getUserById".to_string()),
kind: Some("fn".to_string()),
content: "fn getUserById() {}".to_string(),
also_matched_via: Vec::new(),
}],
metrics: DeepMetrics {
confidence_zone: "high".to_string(),
..DeepMetrics::default()
},
confidence: 1.0,
};
assert_eq!(result.confidence, 1.0);
assert_eq!(result.metrics.confidence_zone, "high");
}
#[test]
fn test_symbol_shortcut_semantic_query_skipped() {
use crate::search::query_classifier;
let qt = query_classifier::classify("how does the search engine work");
assert_eq!(qt, query_classifier::QueryType::Semantic);
}
#[test]
fn test_symbol_shortcut_miss_falls_through() {
let empty_hits: Vec<(i64, String, String)> = vec![];
assert!(empty_hits.is_empty()); }
#[test]
fn test_decompose_query_how_does_use() {
let result = decompose_query("how does authentication use sessions");
assert!(result.is_some());
let parts = result.unwrap();
assert_eq!(parts.len(), 2);
assert_eq!(parts[0], "authentication");
assert_eq!(parts[1], "sessions");
}
#[test]
fn test_decompose_query_simple_returns_none() {
let result = decompose_query("search index BM25");
assert!(result.is_none());
}
#[test]
fn test_decompose_query_across_pattern() {
let result = decompose_query("error handling across frontend and backend");
assert!(result.is_some());
let parts = result.unwrap();
assert_eq!(parts.len(), 2);
}
#[test]
fn test_decompose_interact_with() {
let result = decompose_query("how does authentication interact with session management");
assert!(result.is_some());
let parts = result.unwrap();
assert_eq!(parts.len(), 2);
assert!(parts[0].contains("authentication"));
assert!(parts[1].contains("session management"));
}
#[test]
fn test_decompose_from_pattern() {
let result = decompose_query("data flow from request parsing to database write");
assert!(result.is_some());
let parts = result.unwrap();
assert_eq!(parts.len(), 2);
assert!(parts[0].contains("request parsing"));
assert!(parts[1].contains("database write"));
}
#[test]
fn test_decompose_between_pattern() {
let result = decompose_query("error propagation between service layer and handler");
assert!(result.is_some());
let parts = result.unwrap();
assert_eq!(parts.len(), 2);
}
#[test]
fn test_decompose_implement_pattern() {
let result = decompose_query("implement rate limiting");
assert!(result.is_some());
let parts = result.unwrap();
assert_eq!(parts.len(), 2);
assert!(parts[0].contains("rate limiting"));
assert!(parts[1].contains("entry points"));
}
#[test]
fn test_decompose_where_to_add_pattern() {
let result = decompose_query("where to add request validation");
assert!(result.is_some());
let parts = result.unwrap();
assert_eq!(parts.len(), 2);
assert!(parts[0].contains("request validation"));
}
#[test]
fn test_decompose_no_match_keyword_query() {
assert!(decompose_query("authentication").is_none());
assert!(decompose_query("search").is_none());
}
#[test]
fn test_decompose_feature_too_short() {
assert!(decompose_query("add foo").is_none());
}
#[test]
fn test_decompose_feature_second_query_generic() {
let result = decompose_query("add logging");
assert!(result.is_some());
let parts = result.unwrap();
assert!(parts[1].contains("entry points"));
assert!(!parts[1].contains("middleware")); assert!(!parts[1].contains("request handling"));
let result = decompose_query("add structured logging");
assert!(result.is_some());
let parts = result.unwrap();
assert!(parts[1].contains("entry points"));
assert!(!parts[1].contains("middleware")); assert!(!parts[1].contains("request handling")); }
#[allow(clippy::too_many_arguments)] fn make_result_with_channels(
chunk_id: u64,
file: &str,
start: u32,
end: u32,
score: f32,
score_dense: f32,
score_sparse: f32,
score_exact: f32,
) -> SearchResult {
SearchResult {
chunk: Chunk {
id: chunk_id,
file_path: PathBuf::from(file),
start_line: start,
end_line: end,
content: format!("// chunk {chunk_id}"),
chunk_type: ChunkType::AstNode {
name: format!("fn_{chunk_id}"),
kind: AstNodeKind::Function,
language: "rust".to_string(),
structured_meta: None,
},
},
score,
source: SearchSource::Hybrid,
score_dense,
score_sparse,
score_exact,
confidence: crate::types::Confidence::Inferred,
confidence_score: 0.0,
}
}
#[test]
fn test_channels_label_combinations() {
let r = make_result_with_channels(1, "a.rs", 1, 10, 0.5, 0.3, 0.2, 0.0);
assert_eq!(channels_label(&r), "dense+sparse");
let r = make_result_with_channels(2, "a.rs", 1, 10, 0.5, 0.0, 0.0, 0.5);
assert_eq!(channels_label(&r), "exact");
let r = make_result_with_channels(3, "a.rs", 1, 10, 1.0, 0.3, 0.3, 0.4);
assert_eq!(channels_label(&r), "dense+sparse+exact");
let r = make_result_with_channels(4, "a.rs", 1, 10, 0.5, 0.0, 0.0, 0.0);
assert_eq!(channels_label(&r), "hybrid");
}
#[test]
fn test_dedup_by_range_collapses_three_duplicates() {
let r1 = make_result_with_channels(1, "a.rs", 10, 20, 0.9, 0.5, 0.4, 0.0); let r2 = make_result_with_channels(2, "a.rs", 10, 20, 0.5, 0.0, 0.5, 0.0); let r3 = make_result_with_channels(3, "a.rs", 10, 20, 0.3, 0.0, 0.0, 0.3); let r4 = make_result_with_channels(4, "b.rs", 30, 40, 0.8, 0.4, 0.4, 0.0); let r5 = make_result_with_channels(5, "c.rs", 50, 60, 0.7, 0.0, 0.7, 0.0);
let selected = vec![&r1, &r2, &r3, &r4, &r5];
let deduped = dedup_by_range(&selected);
assert_eq!(
deduped.len(),
3,
"5 entries with 3 collapsing should yield 3"
);
assert_eq!(deduped[0].result.chunk.id, 1);
assert_eq!(deduped[0].result.chunk.file_path, PathBuf::from("a.rs"));
assert_eq!(deduped[0].also_matched_via, vec!["exact", "sparse"]);
assert_eq!(deduped[1].result.chunk.id, 4);
assert!(deduped[1].also_matched_via.is_empty());
assert_eq!(deduped[2].result.chunk.id, 5);
assert!(deduped[2].also_matched_via.is_empty());
}
#[test]
fn test_dedup_keeps_highest_score_on_out_of_order_input() {
let r1 = make_result_with_channels(1, "x.rs", 5, 15, 0.2, 0.2, 0.0, 0.0);
let r2 = make_result_with_channels(2, "x.rs", 5, 15, 0.5, 0.0, 0.5, 0.0);
let r3 = make_result_with_channels(3, "x.rs", 5, 15, 0.9, 0.0, 0.0, 0.9);
let selected = vec![&r1, &r2, &r3];
let deduped = dedup_by_range(&selected);
assert_eq!(deduped.len(), 1);
assert_eq!(deduped[0].result.chunk.id, 3);
assert_eq!(deduped[0].also_matched_via, vec!["dense", "sparse"]);
}
#[test]
fn test_dedup_preserves_order_for_uniques() {
let r1 = make_result_with_channels(1, "a.rs", 1, 10, 0.9, 0.5, 0.4, 0.0);
let r2 = make_result_with_channels(2, "b.rs", 1, 10, 0.8, 0.4, 0.4, 0.0);
let r3 = make_result_with_channels(3, "c.rs", 1, 10, 0.7, 0.3, 0.4, 0.0);
let selected = vec![&r1, &r2, &r3];
let deduped = dedup_by_range(&selected);
assert_eq!(deduped.len(), 3);
assert_eq!(deduped[0].result.chunk.id, 1);
assert_eq!(deduped[1].result.chunk.id, 2);
assert_eq!(deduped[2].result.chunk.id, 3);
for d in &deduped {
assert!(d.also_matched_via.is_empty());
}
}
#[test]
fn test_dedup_empty_input() {
let selected: Vec<&SearchResult> = vec![];
let deduped = dedup_by_range(&selected);
assert!(deduped.is_empty());
}
#[test]
fn test_dedup_caps_also_matched_via_at_four() {
let r1 = make_result_with_channels(1, "x.rs", 1, 10, 0.9, 0.5, 0.0, 0.0);
let r2 = make_result_with_channels(2, "x.rs", 1, 10, 0.1, 0.1, 0.0, 0.0); let r3 = make_result_with_channels(3, "x.rs", 1, 10, 0.1, 0.0, 0.1, 0.0); let r4 = make_result_with_channels(4, "x.rs", 1, 10, 0.1, 0.0, 0.0, 0.1); let r5 = make_result_with_channels(5, "x.rs", 1, 10, 0.1, 0.1, 0.1, 0.0); let r6 = make_result_with_channels(6, "x.rs", 1, 10, 0.1, 0.1, 0.1, 0.1);
let selected = vec![&r1, &r2, &r3, &r4, &r5, &r6];
let deduped = dedup_by_range(&selected);
assert_eq!(deduped.len(), 1);
assert_eq!(deduped[0].result.chunk.id, 1);
assert!(
deduped[0].also_matched_via.len() <= 4,
"got: {:?}",
deduped[0].also_matched_via
);
let mut sorted = deduped[0].also_matched_via.clone();
sorted.sort();
assert_eq!(deduped[0].also_matched_via, sorted);
}
#[test]
fn test_append_dedup_footer_with_matches() {
let sources = vec![DeepSource {
file: "src/auth.rs".to_string(),
start_line: 10,
end_line: 30,
name: Some("authenticate".to_string()),
kind: Some("fn".to_string()),
content: "fn authenticate() {}".to_string(),
also_matched_via: vec!["sparse".to_string(), "exact".to_string()],
}];
let answer = "Auth uses JWT.".to_string();
let out = append_dedup_footer(answer, &sources);
assert!(out.starts_with("Auth uses JWT."));
assert!(out.contains("Cross-source dedup:"));
assert!(out.contains("src/auth.rs:10-30 authenticate"));
assert!(out.contains("(also matched via: sparse, exact)"));
}
#[test]
fn test_append_dedup_footer_no_op_when_no_duplicates() {
let sources = vec![DeepSource {
file: "src/auth.rs".to_string(),
start_line: 10,
end_line: 30,
name: Some("authenticate".to_string()),
kind: Some("fn".to_string()),
content: "fn authenticate() {}".to_string(),
also_matched_via: Vec::new(),
}];
let answer = "Auth uses JWT.".to_string();
let out = append_dedup_footer(answer.clone(), &sources);
assert_eq!(
out, answer,
"footer must not be added when no duplicates exist"
);
}
#[test]
fn test_append_dedup_footer_handles_empty_answer() {
let sources = vec![DeepSource {
file: "a.rs".to_string(),
start_line: 1,
end_line: 5,
name: None,
kind: None,
content: "fn a() {}".to_string(),
also_matched_via: vec!["dense".to_string()],
}];
let out = append_dedup_footer(String::new(), &sources);
assert!(out.starts_with("Cross-source dedup:"));
assert!(out.contains("a.rs:1-5"));
assert!(out.contains("(also matched via: dense)"));
}
}