mod score;
pub(crate) use score::{Boosts, Feature, confidence, match_positions, match_quality, path_stem};
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use crate::store::{Store, SymbolRow};
const CANDIDATE_LIMIT: usize = 8000;
const LIVE_REPO_ID: i64 = -1;
const BRANCH_FILE_BOOST: f64 = 180.0;
const BRANCH_DIR_BOOST: f64 = 60.0;
#[derive(Debug, Default, Clone)]
pub(crate) struct ActiveFiles {
files: HashSet<String>,
dirs: HashSet<String>,
}
impl ActiveFiles {
pub(crate) fn new<I: IntoIterator<Item = String>>(paths: I) -> Self {
let files: HashSet<String> = paths.into_iter().collect();
let dirs = files
.iter()
.filter_map(|f| parent_dir(f))
.map(str::to_string)
.collect();
ActiveFiles { files, dirs }
}
fn is_empty(&self) -> bool {
self.files.is_empty()
}
fn boost(&self, path: &str) -> f64 {
if self.files.contains(path) {
BRANCH_FILE_BOOST
} else if parent_dir(path).is_some_and(|d| self.dirs.contains(d)) {
BRANCH_DIR_BOOST
} else {
0.0
}
}
}
fn parent_dir(path: &str) -> Option<&str> {
path.rfind('/').map(|i| &path[..i])
}
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct Hit {
pub name: String,
pub kind: String,
pub language: String,
pub file: String,
pub line: i64,
#[serde(skip_serializing_if = "Option::is_none")]
pub end_line: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub parent: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub visibility: Option<String>,
#[serde(rename = "repo")]
pub repo_identity: String,
#[serde(skip)]
pub score: f64,
pub confidence: f64,
#[serde(serialize_with = "serialize_feature_names")]
pub features: Vec<Feature>,
#[serde(skip_serializing_if = "Option::is_none")]
pub signature: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub body: Option<String>,
}
fn serialize_feature_names<S: serde::Serializer>(
features: &[Feature],
s: S,
) -> Result<S::Ok, S::Error> {
use serde::Serialize;
let mut sorted: Vec<&Feature> = features.iter().collect();
sorted.sort_by(|a, b| b.value.total_cmp(&a.value));
let names: Vec<&str> = sorted.iter().map(|f| f.name).collect();
names.serialize(s)
}
pub(crate) fn search(
store: &Store,
query: &str,
current_repo_id: Option<i64>,
only_repo: Option<i64>,
active: &ActiveFiles,
limit: usize,
) -> crate::store::Result<Vec<Hit>> {
let (leaf, _) = score::parse_qualified(query);
let stripped;
let recall = if score::has_wildcard(leaf) {
stripped = score::strip_wildcards(leaf);
stripped.as_str()
} else {
leaf
};
let trace_on = crate::trace::enabled();
let t = std::time::Instant::now();
let candidates = store.search_candidates(recall, CANDIDATE_LIMIT, score::has_wildcard(leaf))?;
let n_candidates = candidates.len();
let t_recall = t.elapsed();
let t = std::time::Instant::now();
let now = now_unix();
let learned = learned_boosts(store, query, now)?;
let mut hits: Vec<Hit> = candidates
.into_iter()
.filter_map(|c| {
if only_repo.is_some_and(|r| r != c.repository_id) {
return None;
}
let learned_boost = if learned.is_empty() {
0.0
} else {
let key = (c.repository_id, c.file.clone(), c.name.clone());
learned.get(&key).copied().unwrap_or(0.0)
};
let boosts = Boosts {
learned: learned_boost,
recency: recency_boost(c.git_ts.max(c.mtime.map(|n| n / 1_000_000_000)), now),
branch: if active.is_empty() {
0.0
} else {
active.boost(&c.file)
},
};
rank_one(query, c, current_repo_id, boosts)
})
.collect();
let n_hits = hits.len();
let t_score = t.elapsed();
let t = std::time::Instant::now();
sort_and_truncate(&mut hits, limit);
crate::profile::record("recall", t_recall, || format!("{n_candidates} candidates"));
crate::profile::record("score", t_score, || format!("{n_hits} hits"));
crate::profile::record("sort", t.elapsed(), || format!("top {limit}"));
if trace_on {
crate::trace!(
"search {query:?}: recall {n_candidates} cand in {} ms, score→{n_hits} hits in {} ms, sort {} ms",
t_recall.as_millis(),
t_score.as_millis(),
t.elapsed().as_millis(),
);
}
Ok(hits)
}
fn recency_boost(mtime: Option<i64>, now: i64) -> f64 {
let Some(mtime) = mtime else {
return 0.0;
};
let age_days = (now - mtime).max(0) as f64 / 86_400.0;
let boost = 120.0 * 0.5_f64.powf(age_days / 14.0);
if boost < 1.0 { 0.0 } else { boost }
}
fn learned_boosts(
store: &Store,
query: &str,
now: i64,
) -> crate::store::Result<HashMap<(i64, String, String), f64>> {
let q = query.to_ascii_lowercase();
let mut map: HashMap<(i64, String, String), f64> = HashMap::new();
for s in store.selections_for(&q)? {
let boost = learned_boost(s.selections, s.last_selected_at, now);
let entry = map.entry((s.repository_id, s.file, s.name)).or_insert(0.0);
*entry = entry.max(boost);
}
Ok(map)
}
fn learned_boost(selections: i64, last_selected_at: i64, now: i64) -> f64 {
if selections <= 0 {
return 0.0;
}
let strength = (selections.min(5) as f64) / 5.0;
let age_days = (now - last_selected_at).max(0) as f64 / 86_400.0;
let recency = 0.5_f64.powf(age_days / 30.0);
260.0 * strength * recency
}
fn now_unix() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
pub(crate) fn live_search(
root: &Path,
query: &str,
limit: usize,
skip: &HashSet<String>,
deadline: Option<Instant>,
prefilter: bool,
) -> Vec<Hit> {
let needle = prefilter.then_some(query.as_bytes());
let identity = crate::index::detect_identity(root).to_string();
let mut hits: Vec<Hit> = crate::index::scan(root, skip, deadline, needle)
.into_iter()
.flat_map(|fs| fs.symbols)
.filter_map(|s| {
let row = SymbolRow {
name: s.name,
kind: s.kind.as_str().to_string(),
language: s.language,
file: s.file,
line: s.line as i64,
end_line: Some(s.end_line as i64),
parent: s.parent,
repository_id: LIVE_REPO_ID,
repo_identity: identity.clone(),
mtime: None,
git_ts: None,
visibility: s.visibility.map(str::to_string),
};
rank_one(query, row, Some(LIVE_REPO_ID), Boosts::default())
})
.collect();
sort_and_truncate(&mut hits, limit);
hits
}
pub(crate) fn merge(a: Vec<Hit>, b: Vec<Hit>, limit: usize) -> Vec<Hit> {
use std::collections::HashMap;
let mut by_key: HashMap<(String, i64, String), Hit> = HashMap::new();
for hit in a.into_iter().chain(b) {
let key = (hit.file.clone(), hit.line, hit.name.clone());
match by_key.get(&key) {
Some(existing) if existing.score >= hit.score => {}
_ => {
by_key.insert(key, hit);
}
}
}
let mut hits: Vec<Hit> = by_key.into_values().collect();
sort_and_truncate(&mut hits, limit);
hits
}
pub(crate) fn apply_scope_gate(query: &str, hits: &mut Vec<Hit>) {
if score::parse_qualified(query).1.is_none() {
return; }
let in_scope = |h: &Hit| h.features.iter().any(|f| f.name == "parent");
if hits.iter().any(in_scope) {
hits.retain(in_scope);
}
}
fn sort_and_truncate(hits: &mut Vec<Hit>, limit: usize) {
hits.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.name.len().cmp(&b.name.len()))
.then_with(|| a.name.cmp(&b.name))
.then_with(|| (&a.file, a.line).cmp(&(&b.file, b.line)))
});
hits.truncate(limit);
}
fn rank_one(
query: &str,
c: SymbolRow,
current_repo_id: Option<i64>,
boosts: Boosts,
) -> Option<Hit> {
let scored = score::score(query, &c, current_repo_id, boosts)?;
Some(Hit {
name: c.name,
kind: c.kind,
language: c.language,
file: c.file,
line: c.line,
end_line: c.end_line,
parent: c.parent,
visibility: c.visibility,
repo_identity: c.repo_identity,
score: scored.total,
confidence: 0.0, features: scored.features,
signature: None,
body: None,
})
}
#[cfg(test)]
mod tests {
#[test]
fn identical_names_rank_in_a_stable_order() {
let hit = |file: &str, line: i64| Hit {
name: "Transaction".into(),
kind: "class".into(),
language: "ruby".into(),
file: file.into(),
line,
end_line: None,
parent: None,
visibility: None,
score: 1.0,
confidence: 0.5,
signature: None,
repo_identity: "local:/tmp/x".into(),
features: Vec::new(),
body: None,
};
let ordered = |mut hits: Vec<Hit>| {
sort_and_truncate(&mut hits, 10);
hits.into_iter()
.map(|h| (h.file, h.line))
.collect::<Vec<_>>()
};
let a = ordered(vec![
hit("app/models/b.rb", 1),
hit("app/models/a.rb", 9),
hit("app/models/a.rb", 2),
]);
let b = ordered(vec![
hit("app/models/a.rb", 2),
hit("app/models/b.rb", 1),
hit("app/models/a.rb", 9),
]);
assert_eq!(a, b, "ranking must not depend on row order");
assert_eq!(
a,
vec![
("app/models/a.rb".to_string(), 2),
("app/models/a.rb".to_string(), 9),
("app/models/b.rb".to_string(), 1),
]
);
}
use super::*;
use crate::core::{Kind, Symbol};
fn sym(name: &str, kind: Kind) -> Symbol {
Symbol {
name: name.into(),
kind,
language: "ruby".into(),
file: "app/x.rb".into(),
line: 1,
end_line: 1,
parent: None,
visibility: None,
}
}
fn store_with(symbols: &[Symbol]) -> Store {
let mut store = Store::open_in_memory().unwrap();
let repo = store
.upsert_repository(&crate::core::RepoIdentity::local("/tmp/x"), None)
.unwrap();
store
.replace_file_symbols(repo, "app/x.rb", "ruby", None, "h", symbols)
.unwrap();
store
}
fn names(hits: &[Hit]) -> Vec<&str> {
hits.iter().map(|h| h.name.as_str()).collect()
}
fn store_two_repos() -> (Store, i64, i64) {
let mut store = Store::open_in_memory().unwrap();
let a = store
.upsert_repository(&crate::core::RepoIdentity::local("/tmp/a"), None)
.unwrap();
let b = store
.upsert_repository(&crate::core::RepoIdentity::local("/tmp/b"), None)
.unwrap();
store
.replace_file_symbols(a, "a.rb", "ruby", None, "h", &[sym("Widget", Kind::Class)])
.unwrap();
store
.replace_file_symbols(b, "b.rb", "ruby", None, "h", &[sym("Widget", Kind::Class)])
.unwrap();
(store, a, b)
}
#[test]
fn only_repo_scopes_results_to_that_repo() {
let (store, a, b) = store_two_repos();
let hits = search(
&store,
"Widget",
Some(a),
Some(a),
&ActiveFiles::default(),
10,
)
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].repo_identity, "local:/tmp/a");
let all = search(&store, "Widget", Some(a), None, &ActiveFiles::default(), 10).unwrap();
assert_eq!(all.len(), 2);
let _ = b;
}
#[test]
fn scoped_search_reports_no_match_rather_than_leaking_another_repo() {
let (store, a, _b) = store_two_repos();
let hits = search(
&store,
"Gadget",
Some(a),
Some(a),
&ActiveFiles::default(),
10,
)
.unwrap();
assert!(hits.is_empty());
}
#[test]
fn ranks_exact_match_first() {
let store = store_with(&[
sym("Users", Kind::Class),
sym("User", Kind::Class),
sym("UserMailer", Kind::Class),
]);
let hits = search(&store, "user", None, None, &ActiveFiles::default(), 10).unwrap();
assert_eq!(hits[0].name, "User");
}
#[test]
fn abbreviation_finds_the_intended_symbol() {
let store = store_with(&[
sym("RefundProcessor", Kind::Class),
sym("Refund", Kind::Class),
sym("Payment", Kind::Class),
]);
let hits = search(
&store,
"refundproc",
None,
None,
&ActiveFiles::default(),
10,
)
.unwrap();
assert_eq!(hits[0].name, "RefundProcessor");
assert!(!names(&hits).contains(&"Payment"));
}
#[test]
fn short_fuzzy_query_still_resolves() {
let store = store_with(&[sym("User", Kind::Class), sym("Account", Kind::Class)]);
let hits = search(&store, "usr", None, None, &ActiveFiles::default(), 10).unwrap();
assert_eq!(hits[0].name, "User");
}
#[test]
fn no_match_returns_empty() {
let store = store_with(&[sym("User", Kind::Class)]);
let hits = search(&store, "zzzzz", None, None, &ActiveFiles::default(), 10).unwrap();
assert!(hits.is_empty());
}
#[test]
fn merge_dedups_by_location_keeping_higher_score() {
let mk = |name: &str, score: f64| Hit {
name: name.into(),
kind: "class".into(),
language: "ruby".into(),
file: "a.rb".into(),
line: 1,
end_line: Some(1),
parent: None,
visibility: None,
repo_identity: "r".into(),
score,
confidence: 0.0,
features: vec![],
signature: None,
body: None,
};
let from_index = vec![mk("User", 100.0)];
let from_live = vec![mk("User", 500.0), mk("Account", 200.0)];
let merged = merge(from_index, from_live, 10);
assert_eq!(merged.len(), 2, "the duplicate User is collapsed");
assert_eq!(merged[0].name, "User");
assert_eq!(merged[0].score, 500.0, "the higher-scored duplicate wins");
}
#[test]
fn active_files_boosts_the_file_and_its_neighbors() {
let active = ActiveFiles::new(["app/services/refund.rb".to_string()]);
assert_eq!(active.boost("app/services/refund.rb"), BRANCH_FILE_BOOST);
assert_eq!(active.boost("app/services/charge.rb"), BRANCH_DIR_BOOST);
assert_eq!(active.boost("app/models/user.rb"), 0.0);
}
fn nested(name: &str, kind: Kind, parent: &str) -> Symbol {
Symbol {
parent: Some(parent.into()),
..sym(name, kind)
}
}
#[test]
fn qualified_query_ranks_the_definition_in_the_named_scope() {
let store = store_with(&[
nested("Config", Kind::Class, "Baz"),
nested("Config", Kind::Class, "Foo"),
nested("Config", Kind::Class, "Qux"),
]);
let hits = search(
&store,
"Foo::Config",
None,
None,
&ActiveFiles::default(),
10,
)
.unwrap();
assert_eq!(hits[0].parent.as_deref(), Some("Foo"));
assert!(hits[0].features.iter().any(|f| f.name == "parent"));
}
#[test]
fn qualifier_resolves_modules_and_methods_too() {
let store = store_with(&[
nested("perform", Kind::Method, "Bar::Worker"),
nested("perform", Kind::Method, "Other::Worker"),
nested("Worker", Kind::Module, "Bar"),
]);
let m = search(
&store,
"Bar::Worker#perform",
None,
None,
&ActiveFiles::default(),
10,
)
.unwrap();
assert_eq!(m[0].kind, "method");
assert_eq!(m[0].parent.as_deref(), Some("Bar::Worker"));
let w = search(
&store,
"Bar::Worker",
None,
None,
&ActiveFiles::default(),
10,
)
.unwrap();
assert_eq!(w[0].name, "Worker");
assert_eq!(w[0].parent.as_deref(), Some("Bar"));
}
fn hit(name: &str, in_scope: bool) -> Hit {
Hit {
name: name.into(),
kind: "method".into(),
language: "ruby".into(),
file: "a.rb".into(),
line: 1,
end_line: Some(1),
parent: None,
visibility: None,
repo_identity: "r".into(),
score: 1.0,
confidence: 0.0,
features: if in_scope {
vec![Feature {
name: "parent",
value: 180.0,
}]
} else {
vec![]
},
signature: None,
body: None,
}
}
#[test]
fn scope_gate_keeps_only_in_scope_results_when_some_match() {
let mut hits = vec![hit("baz", true), hit("baz", false), hit("baz", false)];
apply_scope_gate("Foo::Bar#baz", &mut hits);
assert_eq!(hits.len(), 1, "out-of-scope baz methods are dropped");
assert!(hits[0].features.iter().any(|f| f.name == "parent"));
}
#[test]
fn scope_gate_falls_back_when_nothing_matches_the_scope() {
let mut hits = vec![hit("baz", false), hit("baz", false)];
apply_scope_gate("Foo::Bar#baz", &mut hits);
assert_eq!(hits.len(), 2, "fall back rather than return empty");
}
#[test]
fn scope_gate_is_a_noop_for_an_unqualified_query() {
let mut hits = vec![hit("baz", true), hit("baz", false)];
apply_scope_gate("baz", &mut hits);
assert_eq!(hits.len(), 2, "no qualifier — nothing to gate on");
}
#[test]
fn branch_boost_lifts_an_active_file() {
let store = store_with(&[sym("User", Kind::Class)]); let active = ActiveFiles::new(["app/x.rb".to_string()]);
let hits = search(&store, "user", None, None, &active, 10).unwrap();
assert!(hits[0].features.iter().any(|f| f.name == "branch"));
}
}