1mod score;
9
10pub(crate) use score::{Boosts, Feature, 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(crate) struct ActiveFiles {
38 files: HashSet<String>,
39 dirs: HashSet<String>,
40}
41
42impl ActiveFiles {
43 pub(crate) 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(crate) 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 #[serde(skip_serializing_if = "is_one")]
119 pub declarations: usize,
120 #[serde(skip_serializing_if = "Vec::is_empty")]
123 pub also_in: Vec<String>,
124 pub total: usize,
127 #[serde(skip_serializing_if = "Option::is_none")]
131 pub explain: Option<std::collections::BTreeMap<String, f64>>,
132}
133
134fn serialize_feature_names<S: serde::Serializer>(
137 features: &[Feature],
138 s: S,
139) -> Result<S::Ok, S::Error> {
140 use serde::Serialize;
141 let mut sorted: Vec<&Feature> = features.iter().collect();
142 sorted.sort_by(|a, b| b.value.total_cmp(&a.value));
143 let names: Vec<&str> = sorted.iter().map(|f| f.name).collect();
144 names.serialize(s)
145}
146
147pub(crate) struct Matches {
149 pub hits: Vec<Hit>,
150 pub total: usize,
152}
153
154fn is_one(n: &usize) -> bool {
156 *n <= 1
157}
158
159impl std::ops::Deref for Matches {
162 type Target = [Hit];
163 fn deref(&self) -> &[Hit] {
164 &self.hits
165 }
166}
167
168pub(crate) fn search(
174 store: &Store,
175 query: &str,
176 current_repo_id: Option<i64>,
177 only_repo: Option<i64>,
178 active: &ActiveFiles,
179 limit: usize,
180) -> crate::store::Result<Matches> {
181 let (leaf, _) = score::parse_qualified(query);
186 let stripped;
187 let recall = if score::has_wildcard(leaf) {
188 stripped = score::strip_wildcards(leaf);
189 stripped.as_str()
190 } else {
191 leaf
192 };
193 let trace_on = crate::trace::enabled();
194 let t = std::time::Instant::now();
195 let candidates = store.search_candidates(recall, CANDIDATE_LIMIT, score::has_wildcard(leaf))?;
196 let n_candidates = candidates.len();
197 let t_recall = t.elapsed();
198 let t = std::time::Instant::now();
199 let now = now_unix();
200 let learned = learned_boosts(store, query, now)?;
201
202 let rank = |candidates: &[SymbolRow], near_miss: bool| -> Vec<Hit> {
205 candidates
206 .iter()
207 .filter_map(|c| {
208 if only_repo.is_some_and(|r| r != c.repository_id) {
211 return None;
212 }
213 let learned_boost = if learned.is_empty() {
216 0.0
217 } else {
218 let key = (c.repository_id, c.file.clone(), c.name.clone());
219 learned.get(&key).copied().unwrap_or(0.0)
220 };
221 let boosts = Boosts {
222 learned: learned_boost,
223 recency: recency_boost(c.git_ts.max(c.mtime.map(|n| n / 1_000_000_000)), now),
227 branch: if active.is_empty() {
228 0.0
229 } else {
230 active.boost(&c.file)
231 },
232 };
233 rank_one(query, c, current_repo_id, boosts, near_miss)
234 })
235 .collect()
236 };
237 let mut hits = rank(&candidates, false);
241 if hits.iter().all(|h| h.score <= 0.0) {
246 let near: Vec<SymbolRow> = candidates
250 .into_iter()
251 .filter(|c| score::near_miss_possible(query, &c.name))
252 .collect();
253 let retried = rank(&near, true);
254 if !retried.is_empty() {
256 hits = retried;
257 }
258 }
259 let n_hits = hits.len();
260 let t_score = t.elapsed();
261
262 let t = std::time::Instant::now();
263 let total = sort_and_truncate(&mut hits, limit);
266 crate::profile::record("recall", t_recall, || format!("{n_candidates} candidates"));
269 crate::profile::record("score", t_score, || format!("{n_hits} hits"));
270 crate::profile::record("sort", t.elapsed(), || format!("top {limit}"));
271 if trace_on {
272 crate::trace!(
273 "search {query:?}: recall {n_candidates} cand in {} ms, score→{n_hits} hits in {} ms, sort {} ms",
274 t_recall.as_millis(),
275 t_score.as_millis(),
276 t.elapsed().as_millis(),
277 );
278 }
279 Ok(Matches { hits, total })
280}
281
282fn recency_boost(mtime: Option<i64>, now: i64) -> f64 {
285 let Some(mtime) = mtime else {
286 return 0.0;
287 };
288 let age_days = (now - mtime).max(0) as f64 / 86_400.0;
289 let boost = 120.0 * 0.5_f64.powf(age_days / 14.0);
290 if boost < 1.0 { 0.0 } else { boost }
291}
292
293fn learned_boosts(
295 store: &Store,
296 query: &str,
297 now: i64,
298) -> crate::store::Result<HashMap<(i64, String, String), f64>> {
299 let q = query.to_ascii_lowercase();
300 let mut map: HashMap<(i64, String, String), f64> = HashMap::new();
301 for s in store.selections_for(&q)? {
302 let boost = learned_boost(s.selections, s.last_selected_at, now);
305 let entry = map.entry((s.repository_id, s.file, s.name)).or_insert(0.0);
306 *entry = entry.max(boost);
307 }
308 Ok(map)
309}
310
311fn learned_boost(selections: i64, last_selected_at: i64, now: i64) -> f64 {
322 if selections <= 0 {
323 return 0.0;
324 }
325 let strength = (selections.min(5) as f64) / 5.0;
326 let age_days = (now - last_selected_at).max(0) as f64 / 86_400.0;
327 let recency = 0.5_f64.powf(age_days / 30.0);
328 260.0 * strength * recency
329}
330
331fn now_unix() -> i64 {
332 SystemTime::now()
333 .duration_since(UNIX_EPOCH)
334 .map(|d| d.as_secs() as i64)
335 .unwrap_or(0)
336}
337
338pub(crate) fn live_search(
346 root: &Path,
347 query: &str,
348 limit: usize,
349 skip: &HashSet<String>,
350 deadline: Option<Instant>,
351 prefilter: bool,
352) -> Vec<Hit> {
353 let needle = prefilter.then_some(query.as_bytes());
354 let identity = crate::index::detect_identity(root).to_string();
355 let mut hits: Vec<Hit> = crate::index::scan(root, skip, deadline, needle)
356 .into_iter()
357 .flat_map(|fs| fs.symbols)
358 .filter_map(|s| {
359 let row = SymbolRow {
360 name: s.name,
361 kind: s.kind.as_str().to_string(),
362 language: s.language,
363 file: s.file,
364 line: s.line as i64,
365 end_line: Some(s.end_line as i64),
366 parent: s.parent,
367 repository_id: LIVE_REPO_ID,
368 repo_identity: identity.clone(),
369 mtime: None,
370 git_ts: None,
371 visibility: s.visibility.map(str::to_string),
372 };
373 rank_one(query, &row, Some(LIVE_REPO_ID), Boosts::default(), false)
374 })
375 .collect();
376 sort_and_truncate(&mut hits, limit);
377 hits
378}
379
380pub(crate) fn merge(a: Vec<Hit>, b: Vec<Hit>, limit: usize) -> Vec<Hit> {
384 use std::collections::HashMap;
385 let mut by_key: HashMap<(String, i64, String), Hit> = HashMap::new();
386 for hit in a.into_iter().chain(b) {
387 let key = (hit.file.clone(), hit.line, hit.name.clone());
388 match by_key.get(&key) {
389 Some(existing) if existing.score >= hit.score => {}
390 _ => {
391 by_key.insert(key, hit);
392 }
393 }
394 }
395 let mut hits: Vec<Hit> = by_key.into_values().collect();
396 sort_and_truncate(&mut hits, limit);
397 hits
398}
399
400pub(crate) fn apply_scope_gate(query: &str, hits: &mut Vec<Hit>) {
411 if score::parse_qualified(query).1.is_none() {
412 return; }
414 let in_scope = |h: &Hit| h.features.iter().any(|f| f.name == "parent");
415 if hits.iter().any(in_scope) {
416 hits.retain(in_scope);
417 }
418}
419
420fn sort_and_truncate(hits: &mut Vec<Hit>, limit: usize) -> usize {
431 hits.sort_by(|a, b| {
432 b.score
433 .partial_cmp(&a.score)
434 .unwrap_or(std::cmp::Ordering::Equal)
435 .then_with(|| a.name.len().cmp(&b.name.len()))
436 .then_with(|| a.name.cmp(&b.name))
437 .then_with(|| (&a.file, a.line).cmp(&(&b.file, b.line)))
438 });
439 collapse_declarations(hits);
440 let total = hits.len();
441 hits.truncate(limit);
442 total
443}
444
445fn collapse_declarations(hits: &mut Vec<Hit>) {
459 use std::collections::HashMap;
460 let mut first: HashMap<(String, String, String, String), usize> = HashMap::new();
461 let mut folded: Vec<Vec<String>> = vec![Vec::new(); hits.len()];
462 let mut keep = Vec::with_capacity(hits.len());
463 for (i, hit) in hits.iter().enumerate() {
464 let Some(parent) = hit.parent.clone() else {
465 keep.push(true);
466 continue;
467 };
468 let key = (
469 hit.repo_identity.clone(),
470 parent,
471 hit.name.clone(),
472 hit.kind.clone(),
473 );
474 match first.get(&key) {
475 Some(&at) => {
476 folded[at].push(format!("{}:{}", hit.file, hit.line));
477 keep.push(false);
478 }
479 None => {
480 first.insert(key, i);
481 keep.push(true);
482 }
483 }
484 }
485 let mut i = 0;
486 hits.retain(|_| {
487 let k = keep[i];
488 i += 1;
489 k
490 });
491 let mut survivors = keep.iter().enumerate().filter(|(_, k)| **k).map(|(i, _)| i);
493 for hit in hits.iter_mut() {
494 let Some(src) = survivors.next() else { break };
495 if !folded[src].is_empty() {
496 hit.declarations = 1 + folded[src].len();
497 hit.also_in = std::mem::take(&mut folded[src]);
498 }
499 }
500}
501
502fn rank_one(
503 query: &str,
504 c: &SymbolRow,
505 current_repo_id: Option<i64>,
506 boosts: Boosts,
507 near_miss: bool,
508) -> Option<Hit> {
509 let scored = score::score(query, c, current_repo_id, boosts, near_miss)?;
512 Some(Hit {
513 name: c.name.clone(),
514 kind: c.kind.clone(),
515 language: c.language.clone(),
516 file: c.file.clone(),
517 line: c.line,
518 end_line: c.end_line,
519 parent: c.parent.clone(),
520 visibility: c.visibility.clone(),
521 repo_identity: c.repo_identity.clone(),
522 score: scored.total,
523 confidence: 0.0, features: scored.features,
525 signature: None,
526 body: None,
527 declarations: 1,
528 also_in: Vec::new(),
529 total: 0, explain: None,
531 })
532}
533
534#[cfg(test)]
535mod tests {
536 #[test]
537 fn identical_names_rank_in_a_stable_order() {
538 let hit = |file: &str, line: i64| Hit {
543 name: "Transaction".into(),
544 kind: "class".into(),
545 language: "ruby".into(),
546 file: file.into(),
547 line,
548 end_line: None,
549 parent: None,
550 visibility: None,
551 score: 1.0,
552 confidence: 0.5,
553 signature: None,
554 repo_identity: "local:/tmp/x".into(),
555 features: Vec::new(),
556 body: None,
557 declarations: 1,
558 also_in: Vec::new(),
559 total: 0,
560 explain: None,
561 };
562 let ordered = |mut hits: Vec<Hit>| {
563 sort_and_truncate(&mut hits, 10);
564 hits.into_iter()
565 .map(|h| (h.file, h.line))
566 .collect::<Vec<_>>()
567 };
568
569 let a = ordered(vec![
570 hit("app/models/b.rb", 1),
571 hit("app/models/a.rb", 9),
572 hit("app/models/a.rb", 2),
573 ]);
574 let b = ordered(vec![
576 hit("app/models/a.rb", 2),
577 hit("app/models/b.rb", 1),
578 hit("app/models/a.rb", 9),
579 ]);
580 assert_eq!(a, b, "ranking must not depend on row order");
581 assert_eq!(
582 a,
583 vec![
584 ("app/models/a.rb".to_string(), 2),
585 ("app/models/a.rb".to_string(), 9),
586 ("app/models/b.rb".to_string(), 1),
587 ]
588 );
589 }
590
591 use super::*;
592 use crate::core::{Kind, Symbol};
593
594 fn sym(name: &str, kind: Kind) -> Symbol {
595 Symbol {
596 name: name.into(),
597 kind,
598 language: "ruby".into(),
599 file: "app/x.rb".into(),
600 line: 1,
601 end_line: 1,
602 parent: None,
603 visibility: None,
604 }
605 }
606
607 fn store_with(symbols: &[Symbol]) -> Store {
608 let mut store = Store::open_in_memory().unwrap();
609 let repo = store
610 .upsert_repository(&crate::core::RepoIdentity::local("/tmp/x"), None)
611 .unwrap();
612 store
613 .replace_file_symbols(repo, "app/x.rb", "ruby", None, "h", symbols)
614 .unwrap();
615 store
616 }
617
618 fn names(hits: &[Hit]) -> Vec<&str> {
619 hits.iter().map(|h| h.name.as_str()).collect()
620 }
621
622 fn store_two_repos() -> (Store, i64, i64) {
624 let mut store = Store::open_in_memory().unwrap();
625 let a = store
626 .upsert_repository(&crate::core::RepoIdentity::local("/tmp/a"), None)
627 .unwrap();
628 let b = store
629 .upsert_repository(&crate::core::RepoIdentity::local("/tmp/b"), None)
630 .unwrap();
631 store
632 .replace_file_symbols(a, "a.rb", "ruby", None, "h", &[sym("Widget", Kind::Class)])
633 .unwrap();
634 store
635 .replace_file_symbols(b, "b.rb", "ruby", None, "h", &[sym("Widget", Kind::Class)])
636 .unwrap();
637 (store, a, b)
638 }
639
640 #[test]
641 fn only_repo_scopes_results_to_that_repo() {
642 let (store, a, b) = store_two_repos();
643 let hits = search(
645 &store,
646 "Widget",
647 Some(a),
648 Some(a),
649 &ActiveFiles::default(),
650 10,
651 )
652 .unwrap();
653 assert_eq!(hits.hits.len(), 1);
654 assert_eq!(hits.hits[0].repo_identity, "local:/tmp/a");
655 let all = search(&store, "Widget", Some(a), None, &ActiveFiles::default(), 10).unwrap();
657 assert_eq!(all.hits.len(), 2);
658 let _ = b;
659 }
660
661 #[test]
662 fn scoped_search_reports_no_match_rather_than_leaking_another_repo() {
663 let (store, a, _b) = store_two_repos();
664 let hits = search(
666 &store,
667 "Gadget",
668 Some(a),
669 Some(a),
670 &ActiveFiles::default(),
671 10,
672 )
673 .unwrap();
674 assert!(hits.is_empty());
675 }
676
677 #[test]
678 fn ranks_exact_match_first() {
679 let store = store_with(&[
680 sym("Users", Kind::Class),
681 sym("User", Kind::Class),
682 sym("UserMailer", Kind::Class),
683 ]);
684 let hits = search(&store, "user", None, None, &ActiveFiles::default(), 10).unwrap();
685 assert_eq!(hits[0].name, "User");
686 }
687
688 #[test]
689 fn abbreviation_finds_the_intended_symbol() {
690 let store = store_with(&[
691 sym("RefundProcessor", Kind::Class),
692 sym("Refund", Kind::Class),
693 sym("Payment", Kind::Class),
694 ]);
695 let hits = search(
696 &store,
697 "refundproc",
698 None,
699 None,
700 &ActiveFiles::default(),
701 10,
702 )
703 .unwrap();
704 assert_eq!(hits[0].name, "RefundProcessor");
705 assert!(!names(&hits).contains(&"Payment"));
706 }
707
708 #[test]
709 fn short_fuzzy_query_still_resolves() {
710 let store = store_with(&[sym("User", Kind::Class), sym("Account", Kind::Class)]);
711 let hits = search(&store, "usr", None, None, &ActiveFiles::default(), 10).unwrap();
712 assert_eq!(hits[0].name, "User");
713 }
714
715 #[test]
716 fn no_match_returns_empty() {
717 let store = store_with(&[sym("User", Kind::Class)]);
718 let hits = search(&store, "zzzzz", None, None, &ActiveFiles::default(), 10).unwrap();
719 assert!(hits.is_empty());
720 }
721
722 #[test]
723 fn merge_dedups_by_location_keeping_higher_score() {
724 let mk = |name: &str, score: f64| Hit {
725 name: name.into(),
726 kind: "class".into(),
727 language: "ruby".into(),
728 file: "a.rb".into(),
729 line: 1,
730 end_line: Some(1),
731 parent: None,
732 visibility: None,
733 repo_identity: "r".into(),
734 score,
735 confidence: 0.0,
736 features: vec![],
737 signature: None,
738 body: None,
739 declarations: 1,
740 also_in: Vec::new(),
741 total: 0,
742 explain: None,
743 };
744 let from_index = vec![mk("User", 100.0)];
745 let from_live = vec![mk("User", 500.0), mk("Account", 200.0)];
746 let merged = merge(from_index, from_live, 10);
747 assert_eq!(merged.len(), 2, "the duplicate User is collapsed");
748 assert_eq!(merged[0].name, "User");
749 assert_eq!(merged[0].score, 500.0, "the higher-scored duplicate wins");
750 }
751
752 #[test]
753 fn active_files_boosts_the_file_and_its_neighbors() {
754 let active = ActiveFiles::new(["app/services/refund.rb".to_string()]);
755 assert_eq!(active.boost("app/services/refund.rb"), BRANCH_FILE_BOOST);
757 assert_eq!(active.boost("app/services/charge.rb"), BRANCH_DIR_BOOST);
759 assert_eq!(active.boost("app/models/user.rb"), 0.0);
761 }
762
763 fn nested(name: &str, kind: Kind, parent: &str) -> Symbol {
764 Symbol {
765 parent: Some(parent.into()),
766 ..sym(name, kind)
767 }
768 }
769
770 #[test]
771 fn qualified_query_ranks_the_definition_in_the_named_scope() {
772 let store = store_with(&[
773 nested("Config", Kind::Class, "Baz"),
774 nested("Config", Kind::Class, "Foo"),
775 nested("Config", Kind::Class, "Qux"),
776 ]);
777 let hits = search(
779 &store,
780 "Foo::Config",
781 None,
782 None,
783 &ActiveFiles::default(),
784 10,
785 )
786 .unwrap();
787 assert_eq!(hits[0].parent.as_deref(), Some("Foo"));
788 assert!(hits[0].features.iter().any(|f| f.name == "parent"));
789 }
790
791 #[test]
792 fn qualifier_resolves_modules_and_methods_too() {
793 let store = store_with(&[
794 nested("perform", Kind::Method, "Bar::Worker"),
795 nested("perform", Kind::Method, "Other::Worker"),
796 nested("Worker", Kind::Module, "Bar"),
797 ]);
798 let m = search(
800 &store,
801 "Bar::Worker#perform",
802 None,
803 None,
804 &ActiveFiles::default(),
805 10,
806 )
807 .unwrap();
808 assert_eq!(m[0].kind, "method");
809 assert_eq!(m[0].parent.as_deref(), Some("Bar::Worker"));
810 let w = search(
812 &store,
813 "Bar::Worker",
814 None,
815 None,
816 &ActiveFiles::default(),
817 10,
818 )
819 .unwrap();
820 assert_eq!(w[0].name, "Worker");
821 assert_eq!(w[0].parent.as_deref(), Some("Bar"));
822 }
823
824 fn hit(name: &str, in_scope: bool) -> Hit {
825 Hit {
826 name: name.into(),
827 kind: "method".into(),
828 language: "ruby".into(),
829 file: "a.rb".into(),
830 line: 1,
831 end_line: Some(1),
832 parent: None,
833 visibility: None,
834 repo_identity: "r".into(),
835 score: 1.0,
836 confidence: 0.0,
837 features: if in_scope {
838 vec![Feature {
839 name: "parent",
840 value: 180.0,
841 }]
842 } else {
843 vec![]
844 },
845 signature: None,
846 body: None,
847 declarations: 1,
848 also_in: Vec::new(),
849 total: 0,
850 explain: None,
851 }
852 }
853
854 #[test]
855 fn scope_gate_keeps_only_in_scope_results_when_some_match() {
856 let mut hits = vec![hit("baz", true), hit("baz", false), hit("baz", false)];
857 apply_scope_gate("Foo::Bar#baz", &mut hits);
858 assert_eq!(hits.len(), 1, "out-of-scope baz methods are dropped");
859 assert!(hits[0].features.iter().any(|f| f.name == "parent"));
860 }
861
862 #[test]
863 fn scope_gate_falls_back_when_nothing_matches_the_scope() {
864 let mut hits = vec![hit("baz", false), hit("baz", false)];
866 apply_scope_gate("Foo::Bar#baz", &mut hits);
867 assert_eq!(hits.len(), 2, "fall back rather than return empty");
868 }
869
870 #[test]
871 fn scope_gate_is_a_noop_for_an_unqualified_query() {
872 let mut hits = vec![hit("baz", true), hit("baz", false)];
873 apply_scope_gate("baz", &mut hits);
874 assert_eq!(hits.len(), 2, "no qualifier — nothing to gate on");
875 }
876
877 #[test]
878 fn branch_boost_lifts_an_active_file() {
879 let store = store_with(&[sym("User", Kind::Class)]); let active = ActiveFiles::new(["app/x.rb".to_string()]);
881 let hits = search(&store, "user", None, None, &active, 10).unwrap();
882 assert!(hits[0].features.iter().any(|f| f.name == "branch"));
883 }
884}