1mod score;
9
10pub use score::{Boosts, Feature, Scored, confidence, match_positions, match_quality, path_stem};
11
12use std::collections::{HashMap, HashSet};
13use std::path::Path;
14use std::time::{Instant, SystemTime, UNIX_EPOCH};
15
16use crate::store::{Store, SymbolRow};
17
18const CANDIDATE_LIMIT: usize = 8000;
23
24const LIVE_REPO_ID: i64 = -1;
27
28const BRANCH_FILE_BOOST: f64 = 180.0;
30const BRANCH_DIR_BOOST: f64 = 60.0;
32
33#[derive(Debug, Default, Clone)]
37pub struct ActiveFiles {
38 files: HashSet<String>,
39 dirs: HashSet<String>,
40}
41
42impl ActiveFiles {
43 pub fn new<I: IntoIterator<Item = String>>(paths: I) -> Self {
45 let files: HashSet<String> = paths.into_iter().collect();
46 let dirs = files
47 .iter()
48 .filter_map(|f| parent_dir(f))
49 .map(str::to_string)
50 .collect();
51 ActiveFiles { files, dirs }
52 }
53
54 fn is_empty(&self) -> bool {
55 self.files.is_empty()
56 }
57
58 fn boost(&self, path: &str) -> f64 {
61 if self.files.contains(path) {
62 BRANCH_FILE_BOOST
63 } else if parent_dir(path).is_some_and(|d| self.dirs.contains(d)) {
64 BRANCH_DIR_BOOST
65 } else {
66 0.0
67 }
68 }
69}
70
71fn parent_dir(path: &str) -> Option<&str> {
74 path.rfind('/').map(|i| &path[..i])
75}
76
77#[derive(Debug, Clone, PartialEq, serde::Serialize)]
79pub struct Hit {
80 pub name: String,
81 pub kind: String,
82 pub language: String,
83 pub file: String,
84 pub line: i64,
85 #[serde(skip_serializing_if = "Option::is_none")]
88 pub end_line: Option<i64>,
89 #[serde(skip_serializing_if = "Option::is_none")]
90 pub parent: Option<String>,
91 #[serde(skip_serializing_if = "Option::is_none")]
94 pub visibility: Option<String>,
95 #[serde(rename = "repo")]
96 pub repo_identity: String,
97 #[serde(skip)]
100 pub score: f64,
101 pub confidence: f64,
104 #[serde(serialize_with = "serialize_feature_names")]
107 pub features: Vec<Feature>,
108 #[serde(skip_serializing_if = "Option::is_none")]
111 pub signature: Option<String>,
112 #[serde(skip_serializing_if = "Option::is_none")]
114 pub body: Option<String>,
115}
116
117fn serialize_feature_names<S: serde::Serializer>(
120 features: &[Feature],
121 s: S,
122) -> Result<S::Ok, S::Error> {
123 use serde::Serialize;
124 let mut sorted: Vec<&Feature> = features.iter().collect();
125 sorted.sort_by(|a, b| b.value.total_cmp(&a.value));
126 let names: Vec<&str> = sorted.iter().map(|f| f.name).collect();
127 names.serialize(s)
128}
129
130pub fn search(
136 store: &Store,
137 query: &str,
138 current_repo_id: Option<i64>,
139 only_repo: Option<i64>,
140 active: &ActiveFiles,
141 limit: usize,
142) -> crate::store::Result<Vec<Hit>> {
143 let (leaf, _) = score::parse_qualified(query);
148 let stripped;
149 let recall = if score::has_wildcard(leaf) {
150 stripped = score::strip_wildcards(leaf);
151 stripped.as_str()
152 } else {
153 leaf
154 };
155 let trace_on = crate::trace::enabled();
156 let t = std::time::Instant::now();
157 let candidates = store.search_candidates(recall, CANDIDATE_LIMIT, score::has_wildcard(leaf))?;
158 let n_candidates = candidates.len();
159 let t_recall = t.elapsed();
160 let t = std::time::Instant::now();
161 let now = now_unix();
162 let learned = learned_boosts(store, query, now)?;
163
164 let mut hits: Vec<Hit> = candidates
165 .into_iter()
166 .filter_map(|c| {
167 if only_repo.is_some_and(|r| r != c.repository_id) {
170 return None;
171 }
172 let learned_boost = if learned.is_empty() {
175 0.0
176 } else {
177 let key = (c.repository_id, c.file.clone(), c.name.clone());
178 learned.get(&key).copied().unwrap_or(0.0)
179 };
180 let boosts = Boosts {
181 learned: learned_boost,
182 recency: recency_boost(c.git_ts.max(c.mtime.map(|n| n / 1_000_000_000)), now),
186 branch: if active.is_empty() {
187 0.0
188 } else {
189 active.boost(&c.file)
190 },
191 };
192 rank_one(query, c, current_repo_id, boosts)
193 })
194 .collect();
195 let n_hits = hits.len();
196 let t_score = t.elapsed();
197
198 let t = std::time::Instant::now();
199 sort_and_truncate(&mut hits, limit);
200 crate::profile::record("recall", t_recall, || format!("{n_candidates} candidates"));
203 crate::profile::record("score", t_score, || format!("{n_hits} hits"));
204 crate::profile::record("sort", t.elapsed(), || format!("top {limit}"));
205 if trace_on {
206 crate::trace!(
207 "search {query:?}: recall {n_candidates} cand in {} ms, score→{n_hits} hits in {} ms, sort {} ms",
208 t_recall.as_millis(),
209 t_score.as_millis(),
210 t.elapsed().as_millis(),
211 );
212 }
213 Ok(hits)
214}
215
216fn recency_boost(mtime: Option<i64>, now: i64) -> f64 {
219 let Some(mtime) = mtime else {
220 return 0.0;
221 };
222 let age_days = (now - mtime).max(0) as f64 / 86_400.0;
223 let boost = 120.0 * 0.5_f64.powf(age_days / 14.0);
224 if boost < 1.0 { 0.0 } else { boost }
225}
226
227fn learned_boosts(
229 store: &Store,
230 query: &str,
231 now: i64,
232) -> crate::store::Result<HashMap<(i64, String, String), f64>> {
233 let q = query.to_ascii_lowercase();
234 let mut map: HashMap<(i64, String, String), f64> = HashMap::new();
235 for s in store.selections_for(&q)? {
236 let boost = learned_boost(s.selections, s.last_selected_at, now);
239 let entry = map.entry((s.repository_id, s.file, s.name)).or_insert(0.0);
240 *entry = entry.max(boost);
241 }
242 Ok(map)
243}
244
245fn learned_boost(selections: i64, last_selected_at: i64, now: i64) -> f64 {
256 if selections <= 0 {
257 return 0.0;
258 }
259 let strength = (selections.min(5) as f64) / 5.0;
260 let age_days = (now - last_selected_at).max(0) as f64 / 86_400.0;
261 let recency = 0.5_f64.powf(age_days / 30.0);
262 260.0 * strength * recency
263}
264
265fn now_unix() -> i64 {
266 SystemTime::now()
267 .duration_since(UNIX_EPOCH)
268 .map(|d| d.as_secs() as i64)
269 .unwrap_or(0)
270}
271
272pub fn live_search(
280 root: &Path,
281 query: &str,
282 limit: usize,
283 skip: &HashSet<String>,
284 deadline: Option<Instant>,
285 prefilter: bool,
286) -> Vec<Hit> {
287 let needle = prefilter.then_some(query.as_bytes());
288 let identity = crate::index::detect_identity(root).to_string();
289 let mut hits: Vec<Hit> = crate::index::scan(root, skip, deadline, needle)
290 .into_iter()
291 .flat_map(|fs| fs.symbols)
292 .filter_map(|s| {
293 let row = SymbolRow {
294 name: s.name,
295 kind: s.kind.as_str().to_string(),
296 language: s.language,
297 file: s.file,
298 line: s.line as i64,
299 end_line: Some(s.end_line as i64),
300 parent: s.parent,
301 repository_id: LIVE_REPO_ID,
302 repo_identity: identity.clone(),
303 mtime: None,
304 git_ts: None,
305 visibility: s.visibility.map(str::to_string),
306 };
307 rank_one(query, row, Some(LIVE_REPO_ID), Boosts::default())
308 })
309 .collect();
310 sort_and_truncate(&mut hits, limit);
311 hits
312}
313
314pub fn merge(a: Vec<Hit>, b: Vec<Hit>, limit: usize) -> Vec<Hit> {
318 use std::collections::HashMap;
319 let mut by_key: HashMap<(String, i64, String), Hit> = HashMap::new();
320 for hit in a.into_iter().chain(b) {
321 let key = (hit.file.clone(), hit.line, hit.name.clone());
322 match by_key.get(&key) {
323 Some(existing) if existing.score >= hit.score => {}
324 _ => {
325 by_key.insert(key, hit);
326 }
327 }
328 }
329 let mut hits: Vec<Hit> = by_key.into_values().collect();
330 sort_and_truncate(&mut hits, limit);
331 hits
332}
333
334pub fn apply_scope_gate(query: &str, hits: &mut Vec<Hit>) {
345 if score::parse_qualified(query).1.is_none() {
346 return; }
348 let in_scope = |h: &Hit| h.features.iter().any(|f| f.name == "parent");
349 if hits.iter().any(in_scope) {
350 hits.retain(in_scope);
351 }
352}
353
354fn sort_and_truncate(hits: &mut Vec<Hit>, limit: usize) {
365 hits.sort_by(|a, b| {
366 b.score
367 .partial_cmp(&a.score)
368 .unwrap_or(std::cmp::Ordering::Equal)
369 .then_with(|| a.name.len().cmp(&b.name.len()))
370 .then_with(|| a.name.cmp(&b.name))
371 .then_with(|| (&a.file, a.line).cmp(&(&b.file, b.line)))
372 });
373 hits.truncate(limit);
374}
375
376fn rank_one(
377 query: &str,
378 c: SymbolRow,
379 current_repo_id: Option<i64>,
380 boosts: Boosts,
381) -> Option<Hit> {
382 let scored = score::score(query, &c, current_repo_id, boosts)?;
383 Some(Hit {
384 name: c.name,
385 kind: c.kind,
386 language: c.language,
387 file: c.file,
388 line: c.line,
389 end_line: c.end_line,
390 parent: c.parent,
391 visibility: c.visibility,
392 repo_identity: c.repo_identity,
393 score: scored.total,
394 confidence: 0.0, features: scored.features,
396 signature: None,
397 body: None,
398 })
399}
400
401#[cfg(test)]
402mod tests {
403 #[test]
404 fn identical_names_rank_in_a_stable_order() {
405 let hit = |file: &str, line: i64| Hit {
410 name: "Transaction".into(),
411 kind: "class".into(),
412 language: "ruby".into(),
413 file: file.into(),
414 line,
415 end_line: None,
416 parent: None,
417 visibility: None,
418 score: 1.0,
419 confidence: 0.5,
420 signature: None,
421 repo_identity: "local:/tmp/x".into(),
422 features: Vec::new(),
423 body: None,
424 };
425 let ordered = |mut hits: Vec<Hit>| {
426 sort_and_truncate(&mut hits, 10);
427 hits.into_iter()
428 .map(|h| (h.file, h.line))
429 .collect::<Vec<_>>()
430 };
431
432 let a = ordered(vec![
433 hit("app/models/b.rb", 1),
434 hit("app/models/a.rb", 9),
435 hit("app/models/a.rb", 2),
436 ]);
437 let b = ordered(vec![
439 hit("app/models/a.rb", 2),
440 hit("app/models/b.rb", 1),
441 hit("app/models/a.rb", 9),
442 ]);
443 assert_eq!(a, b, "ranking must not depend on row order");
444 assert_eq!(
445 a,
446 vec![
447 ("app/models/a.rb".to_string(), 2),
448 ("app/models/a.rb".to_string(), 9),
449 ("app/models/b.rb".to_string(), 1),
450 ]
451 );
452 }
453
454 use super::*;
455 use crate::core::{Kind, Symbol};
456
457 fn sym(name: &str, kind: Kind) -> Symbol {
458 Symbol {
459 name: name.into(),
460 kind,
461 language: "ruby".into(),
462 file: "app/x.rb".into(),
463 line: 1,
464 end_line: 1,
465 parent: None,
466 visibility: None,
467 }
468 }
469
470 fn store_with(symbols: &[Symbol]) -> Store {
471 let mut store = Store::open_in_memory().unwrap();
472 let repo = store
473 .upsert_repository(&crate::core::RepoIdentity::local("/tmp/x"), None)
474 .unwrap();
475 store
476 .replace_file_symbols(repo, "app/x.rb", "ruby", None, "h", symbols)
477 .unwrap();
478 store
479 }
480
481 fn names(hits: &[Hit]) -> Vec<&str> {
482 hits.iter().map(|h| h.name.as_str()).collect()
483 }
484
485 fn store_two_repos() -> (Store, i64, i64) {
487 let mut store = Store::open_in_memory().unwrap();
488 let a = store
489 .upsert_repository(&crate::core::RepoIdentity::local("/tmp/a"), None)
490 .unwrap();
491 let b = store
492 .upsert_repository(&crate::core::RepoIdentity::local("/tmp/b"), None)
493 .unwrap();
494 store
495 .replace_file_symbols(a, "a.rb", "ruby", None, "h", &[sym("Widget", Kind::Class)])
496 .unwrap();
497 store
498 .replace_file_symbols(b, "b.rb", "ruby", None, "h", &[sym("Widget", Kind::Class)])
499 .unwrap();
500 (store, a, b)
501 }
502
503 #[test]
504 fn only_repo_scopes_results_to_that_repo() {
505 let (store, a, b) = store_two_repos();
506 let hits = search(
508 &store,
509 "Widget",
510 Some(a),
511 Some(a),
512 &ActiveFiles::default(),
513 10,
514 )
515 .unwrap();
516 assert_eq!(hits.len(), 1);
517 assert_eq!(hits[0].repo_identity, "local:/tmp/a");
518 let all = search(&store, "Widget", Some(a), None, &ActiveFiles::default(), 10).unwrap();
520 assert_eq!(all.len(), 2);
521 let _ = b;
522 }
523
524 #[test]
525 fn scoped_search_reports_no_match_rather_than_leaking_another_repo() {
526 let (store, a, _b) = store_two_repos();
527 let hits = search(
529 &store,
530 "Gadget",
531 Some(a),
532 Some(a),
533 &ActiveFiles::default(),
534 10,
535 )
536 .unwrap();
537 assert!(hits.is_empty());
538 }
539
540 #[test]
541 fn ranks_exact_match_first() {
542 let store = store_with(&[
543 sym("Users", Kind::Class),
544 sym("User", Kind::Class),
545 sym("UserMailer", Kind::Class),
546 ]);
547 let hits = search(&store, "user", None, None, &ActiveFiles::default(), 10).unwrap();
548 assert_eq!(hits[0].name, "User");
549 }
550
551 #[test]
552 fn abbreviation_finds_the_intended_symbol() {
553 let store = store_with(&[
554 sym("RefundProcessor", Kind::Class),
555 sym("Refund", Kind::Class),
556 sym("Payment", Kind::Class),
557 ]);
558 let hits = search(
559 &store,
560 "refundproc",
561 None,
562 None,
563 &ActiveFiles::default(),
564 10,
565 )
566 .unwrap();
567 assert_eq!(hits[0].name, "RefundProcessor");
568 assert!(!names(&hits).contains(&"Payment"));
569 }
570
571 #[test]
572 fn short_fuzzy_query_still_resolves() {
573 let store = store_with(&[sym("User", Kind::Class), sym("Account", Kind::Class)]);
574 let hits = search(&store, "usr", None, None, &ActiveFiles::default(), 10).unwrap();
575 assert_eq!(hits[0].name, "User");
576 }
577
578 #[test]
579 fn no_match_returns_empty() {
580 let store = store_with(&[sym("User", Kind::Class)]);
581 let hits = search(&store, "zzzzz", None, None, &ActiveFiles::default(), 10).unwrap();
582 assert!(hits.is_empty());
583 }
584
585 #[test]
586 fn merge_dedups_by_location_keeping_higher_score() {
587 let mk = |name: &str, score: f64| Hit {
588 name: name.into(),
589 kind: "class".into(),
590 language: "ruby".into(),
591 file: "a.rb".into(),
592 line: 1,
593 end_line: Some(1),
594 parent: None,
595 visibility: None,
596 repo_identity: "r".into(),
597 score,
598 confidence: 0.0,
599 features: vec![],
600 signature: None,
601 body: None,
602 };
603 let from_index = vec![mk("User", 100.0)];
604 let from_live = vec![mk("User", 500.0), mk("Account", 200.0)];
605 let merged = merge(from_index, from_live, 10);
606 assert_eq!(merged.len(), 2, "the duplicate User is collapsed");
607 assert_eq!(merged[0].name, "User");
608 assert_eq!(merged[0].score, 500.0, "the higher-scored duplicate wins");
609 }
610
611 #[test]
612 fn active_files_boosts_the_file_and_its_neighbors() {
613 let active = ActiveFiles::new(["app/services/refund.rb".to_string()]);
614 assert_eq!(active.boost("app/services/refund.rb"), BRANCH_FILE_BOOST);
616 assert_eq!(active.boost("app/services/charge.rb"), BRANCH_DIR_BOOST);
618 assert_eq!(active.boost("app/models/user.rb"), 0.0);
620 }
621
622 fn nested(name: &str, kind: Kind, parent: &str) -> Symbol {
623 Symbol {
624 parent: Some(parent.into()),
625 ..sym(name, kind)
626 }
627 }
628
629 #[test]
630 fn qualified_query_ranks_the_definition_in_the_named_scope() {
631 let store = store_with(&[
632 nested("Config", Kind::Class, "Baz"),
633 nested("Config", Kind::Class, "Foo"),
634 nested("Config", Kind::Class, "Qux"),
635 ]);
636 let hits = search(
638 &store,
639 "Foo::Config",
640 None,
641 None,
642 &ActiveFiles::default(),
643 10,
644 )
645 .unwrap();
646 assert_eq!(hits[0].parent.as_deref(), Some("Foo"));
647 assert!(hits[0].features.iter().any(|f| f.name == "parent"));
648 }
649
650 #[test]
651 fn qualifier_resolves_modules_and_methods_too() {
652 let store = store_with(&[
653 nested("perform", Kind::Method, "Bar::Worker"),
654 nested("perform", Kind::Method, "Other::Worker"),
655 nested("Worker", Kind::Module, "Bar"),
656 ]);
657 let m = search(
659 &store,
660 "Bar::Worker#perform",
661 None,
662 None,
663 &ActiveFiles::default(),
664 10,
665 )
666 .unwrap();
667 assert_eq!(m[0].kind, "method");
668 assert_eq!(m[0].parent.as_deref(), Some("Bar::Worker"));
669 let w = search(
671 &store,
672 "Bar::Worker",
673 None,
674 None,
675 &ActiveFiles::default(),
676 10,
677 )
678 .unwrap();
679 assert_eq!(w[0].name, "Worker");
680 assert_eq!(w[0].parent.as_deref(), Some("Bar"));
681 }
682
683 fn hit(name: &str, in_scope: bool) -> Hit {
684 Hit {
685 name: name.into(),
686 kind: "method".into(),
687 language: "ruby".into(),
688 file: "a.rb".into(),
689 line: 1,
690 end_line: Some(1),
691 parent: None,
692 visibility: None,
693 repo_identity: "r".into(),
694 score: 1.0,
695 confidence: 0.0,
696 features: if in_scope {
697 vec![Feature {
698 name: "parent",
699 value: 180.0,
700 }]
701 } else {
702 vec![]
703 },
704 signature: None,
705 body: None,
706 }
707 }
708
709 #[test]
710 fn scope_gate_keeps_only_in_scope_results_when_some_match() {
711 let mut hits = vec![hit("baz", true), hit("baz", false), hit("baz", false)];
712 apply_scope_gate("Foo::Bar#baz", &mut hits);
713 assert_eq!(hits.len(), 1, "out-of-scope baz methods are dropped");
714 assert!(hits[0].features.iter().any(|f| f.name == "parent"));
715 }
716
717 #[test]
718 fn scope_gate_falls_back_when_nothing_matches_the_scope() {
719 let mut hits = vec![hit("baz", false), hit("baz", false)];
721 apply_scope_gate("Foo::Bar#baz", &mut hits);
722 assert_eq!(hits.len(), 2, "fall back rather than return empty");
723 }
724
725 #[test]
726 fn scope_gate_is_a_noop_for_an_unqualified_query() {
727 let mut hits = vec![hit("baz", true), hit("baz", false)];
728 apply_scope_gate("baz", &mut hits);
729 assert_eq!(hits.len(), 2, "no qualifier — nothing to gate on");
730 }
731
732 #[test]
733 fn branch_boost_lifts_an_active_file() {
734 let store = store_with(&[sym("User", Kind::Class)]); let active = ActiveFiles::new(["app/x.rb".to_string()]);
736 let hits = search(&store, "user", None, None, &active, 10).unwrap();
737 assert!(hits[0].features.iter().any(|f| f.name == "branch"));
738 }
739}