Skip to main content

ferritin_common/
search.rs

1pub mod indexer;
2
3use crate::{Navigator, navigator::Suggestion};
4use rayon::prelude::*;
5
6pub use indexer::*;
7
8impl Navigator {
9    /// Search across multiple crates with BM25 scoring
10    ///
11    /// Returns results sorted by score (descending). Empty crate list returns empty results.
12    /// Empty query triggers index loading but returns no matches (useful for prewarming).
13    ///
14    /// Returns Err with suggestions if no crates could be loaded/indexed.
15    pub fn search<'nav, 'query>(
16        &'nav self,
17        query: &'query str,
18        crate_names: &'query [&'query str],
19    ) -> Result<Vec<ScoredResult<'query>>, Vec<Suggestion<'nav>>> {
20        if crate_names.is_empty() {
21            return Ok(vec![]);
22        }
23
24        // Load indexes and search in parallel
25        let results: Vec<_> = crate_names
26            .par_iter()
27            .map(|&crate_name| {
28                self.get_or_build_search_index(crate_name)
29                    .map(|index| (crate_name, index.search(query)))
30            })
31            .collect();
32
33        // Separate successes from failures
34        let mut crate_results = Vec::new();
35        let mut first_error = None;
36
37        for result in results {
38            match result {
39                Ok(data) => crate_results.push(data),
40                Err(suggestions) if first_error.is_none() => first_error = Some(suggestions),
41                Err(_) => {}
42            }
43        }
44
45        // If no crates succeeded, return the first error
46        if crate_results.is_empty() && first_error.is_some() {
47            return Err(first_error.unwrap());
48        }
49
50        // Aggregate results with BM25 scoring
51        let mut scorer = BM25Scorer::new();
52        for (crate_name, results) in crate_results {
53            scorer.add(crate_name, results);
54        }
55
56        Ok(scorer.score())
57    }
58
59    /// Get or build a search index for the given crate
60    ///
61    /// Returns Err with suggestions if the crate cannot be found
62    fn get_or_build_search_index<'nav>(
63        &'nav self,
64        crate_name: &str,
65    ) -> Result<&'nav SearchIndex, Vec<Suggestion<'nav>>> {
66        let crate_name = self.canonicalize(crate_name);
67
68        if let Some(cached) = self.search_indexes.get(&crate_name) {
69            if let Some(index) = cached.as_ref() {
70                return Ok(index);
71            } else {
72                // Permanent failure cached - return empty suggestions
73                return Err(vec![]);
74            }
75        }
76
77        log::info!("Loading search index for {}", crate_name);
78
79        // Use existing SearchIndex::load_or_build which handles disk caching
80        let result = SearchIndex::load_or_build(self, crate_name.as_ref());
81
82        match result {
83            Ok(index) => {
84                let index_ref = self
85                    .search_indexes
86                    .insert(crate_name, Box::new(Some(index)))
87                    .as_ref()
88                    .unwrap();
89                Ok(index_ref)
90            }
91            Err(suggestions) => {
92                // Cache the failure
93                self.search_indexes.insert(crate_name, Box::new(None));
94                Err(suggestions)
95            }
96        }
97    }
98}