use super::tokenizer::codenexus_tokenize;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct FieldWeights {
pub symbol_name: f64,
pub function_name: f64,
pub comment: f64,
pub string_literal: f64,
}
impl Default for FieldWeights {
fn default() -> Self {
Self {
symbol_name: 3.5,
function_name: 2.0,
comment: 0.5,
string_literal: 0.3,
}
}
}
pub(crate) fn relevance_score_with_reason(name: &str, query: &str) -> (f64, &'static str) {
let name_lower = name.to_ascii_lowercase();
let query_lower = query.to_ascii_lowercase();
if name_lower == query_lower {
return (1.0, "exact name match");
}
if name_lower.starts_with(&query_lower) {
return (0.8, "prefix match");
}
let query_tokens = codenexus_tokenize(query);
let name_tokens = codenexus_tokenize(name);
if !query_tokens.is_empty() && !name_tokens.is_empty() {
let all_match = query_tokens
.iter()
.all(|qt| name_tokens.iter().any(|nt| nt == qt));
if all_match {
return (0.7, "token-aligned match");
}
}
if name_lower.contains(&query_lower) {
return (0.5, "substring match");
}
(0.3, "no match")
}
fn bm25f_field_score(field: &str, query: &str) -> f64 {
let (score, _) = relevance_score_with_reason(field, query);
if score <= 0.3 {
0.0
} else {
score
}
}
#[must_use]
pub fn bm25f_score(
query: &str,
symbol_name: &str,
qualified_name: &str,
docstring: &str,
content: &str,
weights: &FieldWeights,
) -> f64 {
if query.trim().is_empty() {
return 0.0;
}
weights.symbol_name * bm25f_field_score(symbol_name, query)
+ weights.function_name * bm25f_field_score(qualified_name, query)
+ weights.comment * bm25f_field_score(docstring, query)
+ weights.string_literal * bm25f_field_score(content, query)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{Language, Node, NodeLabel};
use crate::query::FullTextSearcher;
use crate::storage::Repository;
#[test]
fn test_bm25f_field_weights_default_values() {
let w = FieldWeights::default();
assert_eq!(w.symbol_name, 3.5);
assert_eq!(w.function_name, 2.0);
assert_eq!(w.comment, 0.5);
assert_eq!(w.string_literal, 0.3);
}
#[test]
fn test_bm25f_field_weights_is_copy_and_eq() {
let w1 = FieldWeights::default();
let w2 = w1; assert_eq!(w1, w2);
let w3 = FieldWeights {
symbol_name: 1.0,
..FieldWeights::default()
};
assert_ne!(w1, w3);
}
#[test]
fn test_bm25f_empty_query_returns_zero() {
let w = FieldWeights::default();
let score = bm25f_score("", "parse", "demo.parse", "parse a file", "content", &w);
assert_eq!(score, 0.0);
}
#[test]
fn test_bm25f_whitespace_only_query_returns_zero() {
let w = FieldWeights::default();
let score = bm25f_score(" \t\n", "parse", "demo.parse", "doc", "content", &w);
assert_eq!(score, 0.0);
}
#[test]
fn test_bm25f_no_match_returns_zero() {
let w = FieldWeights::default();
let score = bm25f_score(
"parse",
"read_input",
"demo.read_input",
"read user input",
"stdin content",
&w,
);
assert_eq!(score, 0.0);
}
#[test]
fn test_bm25f_exact_match_in_symbol_name_beats_substring_in_comment() {
let w = FieldWeights::default();
let score_name = bm25f_score("parse", "parse", "demo.parse", "", "", &w);
let score_comment = bm25f_score(
"parse",
"read_input",
"demo.read_input",
"parse a file",
"",
&w,
);
assert!(
score_name > score_comment,
"exact name match ({score_name}) should beat comment substring ({score_comment})"
);
assert!(score_name > 0.0);
assert!(score_comment > 0.0);
}
#[test]
fn test_bm25f_multiple_fields_same_symbol_sums_weights() {
let w = FieldWeights::default();
let score_both = bm25f_score("parse", "parse", "parse", "", "", &w);
let score_name_only = bm25f_score("parse", "parse", "demo.other", "", "", &w);
assert!(
score_both > score_name_only,
"multi-field match ({score_both}) should sum and exceed single-field ({score_name_only})"
);
}
#[test]
fn test_bm25f_custom_weights_change_ranking() {
let w = FieldWeights {
symbol_name: 0.1,
function_name: 0.1,
comment: 10.0,
string_literal: 0.1,
};
let score_name = bm25f_score("parse", "parse", "demo.parse", "", "", &w);
let score_comment = bm25f_score("parse", "read", "demo.read", "parse a file", "", &w);
assert!(
score_comment > score_name,
"with comment weight 10.0, comment match ({score_comment}) should beat name match ({score_name})"
);
}
#[test]
fn test_bm25f_zero_weights_returns_zero() {
let w = FieldWeights {
symbol_name: 0.0,
function_name: 0.0,
comment: 0.0,
string_literal: 0.0,
};
let score = bm25f_score("parse", "parse", "demo.parse", "parse", "parse", &w);
assert_eq!(score, 0.0);
}
#[test]
fn test_bm25f_content_field_contributes_via_string_literal_weight() {
let w = FieldWeights::default();
let score_no_content = bm25f_score("parse", "parse", "demo.parse", "", "", &w);
let score_with_content = bm25f_score("parse", "parse", "demo.parse", "", "parse code", &w);
assert!(
score_with_content > score_no_content,
"content match should add to score: {score_with_content} vs {score_no_content}"
);
}
#[test]
fn relevance_score_with_reason_exact_match_returns_one() {
let (score, reason) = relevance_score_with_reason("parse", "parse");
assert_eq!(score, 1.0);
assert_eq!(reason, "exact name match");
let (score, _) = relevance_score_with_reason("PARSE", "parse");
assert_eq!(score, 1.0);
}
#[test]
fn relevance_score_with_reason_prefix_match_returns_zero_eight() {
let (score, reason) = relevance_score_with_reason("parse_file", "parse");
assert_eq!(score, 0.8);
assert_eq!(reason, "prefix match");
}
#[test]
fn relevance_score_with_reason_token_aligned_match_returns_zero_seven() {
let (score, reason) = relevance_score_with_reason("my_parse_helper", "parse");
assert_eq!(score, 0.7);
assert_eq!(reason, "token-aligned match");
}
#[test]
fn relevance_score_with_reason_substring_match_returns_zero_five() {
let (score, reason) = relevance_score_with_reason("myparsehelper", "parse");
assert_eq!(score, 0.5);
assert_eq!(reason, "substring match");
}
#[test]
fn relevance_score_with_reason_no_match_returns_zero_three() {
let (score, reason) = relevance_score_with_reason("read_input", "parse");
assert_eq!(score, 0.3);
assert_eq!(reason, "no match");
}
#[test]
fn relevance_score_with_reason_camel_case_token_alignment() {
let (score, _) = relevance_score_with_reason("parseFile", "fileparse");
assert_eq!(score, 0.3);
}
#[test]
fn relevance_score_with_reason_camel_case_query_matches_snake_case_name() {
let (score, reason) = relevance_score_with_reason("parse_file", "parseFile");
assert_eq!(
score, 0.7,
"camelCase query must match snake_case name via token alignment"
);
assert_eq!(reason, "token-aligned match");
}
#[test]
fn relevance_score_with_reason_pascal_case_query_matches_snake_case_name() {
let (score, _) = relevance_score_with_reason("parse_file", "ParseFile");
assert_eq!(score, 0.7);
}
#[test]
fn relevance_score_with_reason_snake_case_query_matches_camel_case_name() {
let (score, _) = relevance_score_with_reason("parseFile", "parse_file");
assert_eq!(score, 0.7);
}
#[test]
fn relevance_score_with_reason_mixed_case_query_partial_token_match() {
let (score, _) = relevance_score_with_reason("parse_file", "parseFileBig");
assert_eq!(score, 0.3);
}
#[test]
fn relevance_score_with_reason_name_with_empty_tokens() {
let (score, reason) = relevance_score_with_reason("12345", "parse");
assert_eq!(score, 0.3);
assert_eq!(reason, "no match");
}
#[test]
fn relevance_score_with_reason_query_with_empty_tokens() {
let (score, reason) = relevance_score_with_reason("parse", "12345");
assert_eq!(score, 0.3);
assert_eq!(reason, "no match");
}
fn fresh_repo() -> Repository {
Repository::in_memory().expect("in_memory repository")
}
#[derive(Default)]
struct SampleFunctionSpec {
id: String,
project: String,
name: String,
qn: String,
file: String,
line: u32,
docstring: String,
content: String,
}
fn sample_function_with_docstring(spec: SampleFunctionSpec) -> Node {
let mut builder = Node::builder(NodeLabel::Function, &spec.name, &spec.qn)
.id(&spec.id)
.project(&spec.project)
.file_path(&spec.file)
.start_line(spec.line)
.end_line(spec.line + 10)
.language(Language::Rust)
.signature("fn x()");
if !spec.docstring.is_empty() {
builder = builder.docstring(&spec.docstring);
}
if !spec.content.is_empty() {
builder = builder.properties(serde_json::json!({ "content": spec.content }));
}
builder.build()
}
#[test]
fn test_bm25f_field_weights_boost_symbol_name_matches() {
let repo = fresh_repo();
let mut nodes = Vec::new();
for i in 0..5 {
nodes.push(sample_function_with_docstring(SampleFunctionSpec {
id: format!("name_{i}"),
project: "demo".into(),
name: format!("parse_{i}"),
qn: format!("demo.parse_{i}"),
file: "/a.rs".into(),
line: i as u32 + 1,
..Default::default()
}));
}
for i in 0..5 {
nodes.push(sample_function_with_docstring(SampleFunctionSpec {
id: format!("doc_{i}"),
project: "demo".into(),
name: format!("aaa_{i}"),
qn: format!("demo.aaa_{i}"),
file: "/b.rs".into(),
line: i as u32 + 100,
docstring: format!("parse something {i}"),
..Default::default()
}));
}
repo.save_nodes(&nodes, NodeLabel::Function)
.expect("save_nodes");
let searcher = FullTextSearcher::new(repo.connection());
let results = searcher.search("parse", None, 100).expect("search");
let top5: Vec<&str> = results.iter().take(5).map(|r| r.name.as_str()).collect();
assert!(
!top5.is_empty(),
"expected at least one result, got empty list"
);
for name in &top5 {
assert!(
name.starts_with("parse_"),
"expected symbol_name match in TOP-5, got {name:?}; full top-5: {top5:?}"
);
}
}
#[test]
fn test_bm25f_with_weights_builder_changes_ranking() {
let repo = fresh_repo();
let nodes = vec![
sample_function_with_docstring(SampleFunctionSpec {
id: "n1".into(),
project: "demo".into(),
name: "parse".into(),
qn: "demo.parse".into(),
file: "/a.rs".into(),
line: 1,
..Default::default()
}),
sample_function_with_docstring(SampleFunctionSpec {
id: "n2".into(),
project: "demo".into(),
name: "other".into(),
qn: "demo.other".into(),
file: "/b.rs".into(),
line: 2,
docstring: "parse a file".into(),
..Default::default()
}),
];
repo.save_nodes(&nodes, NodeLabel::Function)
.expect("save_nodes");
let weights = FieldWeights {
symbol_name: 0.1,
function_name: 0.1,
comment: 10.0,
string_literal: 0.1,
};
let searcher = FullTextSearcher::new(repo.connection()).with_weights(weights);
let results = searcher.search("parse", None, 100).expect("search");
assert_eq!(
results.len(),
2,
"expected 2 results, got {}",
results.len()
);
assert_eq!(
results[0].name, "other",
"expected comment-only match to rank first with weights {:?}, got top: {:?}",
weights, results[0]
);
}
#[test]
fn test_bm25f_search_returns_zero_score_for_no_field_match() {
let repo = fresh_repo();
repo.save_nodes(
&[sample_function_with_docstring(SampleFunctionSpec {
id: "f1".into(),
project: "demo".into(),
name: "read_input".into(),
qn: "demo.read_input".into(),
file: "/a.rs".into(),
line: 1,
docstring: "reads from stdin".into(),
content: "let x = read();".into(),
})],
NodeLabel::Function,
)
.expect("save_nodes");
let searcher = FullTextSearcher::new(repo.connection());
let results = searcher.search("parse", None, 100).expect("search");
assert!(results.is_empty(), "expected no results, got {results:?}");
}
#[test]
fn test_bm25f_docstring_match_appears_in_results() {
let repo = fresh_repo();
repo.save_nodes(
&[sample_function_with_docstring(SampleFunctionSpec {
id: "f1".into(),
project: "demo".into(),
name: "compute".into(),
qn: "demo.compute".into(),
file: "/a.rs".into(),
line: 1,
docstring: "parse the input and return result".into(),
..Default::default()
})],
NodeLabel::Function,
)
.expect("save_nodes");
let searcher = FullTextSearcher::new(repo.connection());
let results = searcher.search("parse", None, 100).expect("search");
assert_eq!(
results.len(),
1,
"expected docstring match, got {results:?}"
);
assert_eq!(results[0].name, "compute");
assert!(results[0].score > 0.0, "expected non-zero score");
}
#[test]
fn test_bm25f_content_match_appears_in_results() {
let repo = fresh_repo();
repo.save_nodes(
&[sample_function_with_docstring(SampleFunctionSpec {
id: "f1".into(),
project: "demo".into(),
name: "compute".into(),
qn: "demo.compute".into(),
file: "/a.rs".into(),
line: 1,
content: "let parsed = parse_input();".into(),
..Default::default()
})],
NodeLabel::Function,
)
.expect("save_nodes");
let searcher = FullTextSearcher::new(repo.connection());
let results = searcher.search("parse", None, 100).expect("search");
assert_eq!(results.len(), 1, "expected content match, got {results:?}");
assert_eq!(results[0].name, "compute");
assert!(results[0].score > 0.0, "expected non-zero score");
}
#[test]
fn test_bm25f_qualified_name_match_appears_in_results() {
let repo = fresh_repo();
repo.save_nodes(
&[sample_function_with_docstring(SampleFunctionSpec {
id: "f1".into(),
project: "demo".into(),
name: "compute".into(),
qn: "demo.parse_helper".into(),
file: "/a.rs".into(),
line: 1,
..Default::default()
})],
NodeLabel::Function,
)
.expect("save_nodes");
let searcher = FullTextSearcher::new(repo.connection());
let results = searcher.search("parse", None, 100).expect("search");
assert_eq!(
results.len(),
1,
"expected qualifiedName match, got {results:?}"
);
assert_eq!(results[0].name, "compute");
}
#[test]
fn test_bm25f_label_without_docstring_content_still_works() {
let repo = fresh_repo();
let module = Node::builder(NodeLabel::Module, "parser", "demo.parser")
.id("m1")
.project("demo")
.file_path("/a.rs")
.start_line(1)
.language(Language::Rust)
.build();
repo.save_nodes(&[module], NodeLabel::Module)
.expect("save_nodes");
let searcher = FullTextSearcher::new(repo.connection());
let results = searcher.search("parse", None, 100).expect("search");
assert_eq!(results.len(), 1);
assert_eq!(results[0].name, "parser");
assert_eq!(results[0].label, "Module");
}
}