Skip to main content

rtb_docs/
search.rs

1//! Title + full-text search over the doc tree.
2//!
3//! Two code paths:
4//!
5//! - **Title search** — `fuzzy-matcher::SkimMatcherV2` against the
6//!   index's entry titles. Fast, no preprocessing, suitable for the
7//!   "I know roughly what the page is called" query.
8//! - **Full-text search** — `tantivy` in-memory (`RamDirectory`)
9//!   index built once at [`SearchIndex::build`]. Matches tokens in
10//!   the plain-text projection of every page's body and returns
11//!   top-N matches with a 140-character snippet.
12//!
13//! The `DocsBrowser` TUI and the `DocsServer::search_handler` both
14//! call into [`SearchIndex`].
15
16use fuzzy_matcher::skim::SkimMatcherV2;
17use fuzzy_matcher::FuzzyMatcher as _;
18use tantivy::collector::TopDocs;
19use tantivy::query::QueryParser;
20use tantivy::schema::{Schema, STORED, TEXT};
21use tantivy::{doc, Index, IndexWriter, TantivyDocument};
22
23use crate::error::Result;
24use crate::render;
25
26/// Result of a title search — one row per hit, ordered by descending
27/// fuzzy score.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct TitleHit {
30    /// Doc-tree-relative page path.
31    pub path: String,
32    /// The title as stored on the index entry.
33    pub title: String,
34    /// Fuzzy score (higher = better match).
35    pub score: i64,
36}
37
38/// Result of a full-text search.
39#[derive(Debug, Clone)]
40pub struct FullTextHit {
41    /// Doc-tree-relative page path.
42    pub path: String,
43    /// The page's title.
44    pub title: String,
45    /// Up to 140 chars of the matched body, with the match token
46    /// preserved.
47    pub snippet: String,
48    /// tantivy score (higher = better match).
49    pub score: f32,
50}
51
52/// In-memory tantivy-backed full-text index over the doc tree.
53pub struct SearchIndex {
54    index: Index,
55    query_parser: QueryParser,
56    title_field: tantivy::schema::Field,
57    path_field: tantivy::schema::Field,
58    body_field: tantivy::schema::Field,
59    titles: Vec<(String, String)>, // (path, title) — for title search
60    matcher: SkimMatcherV2,
61}
62
63impl SearchIndex {
64    /// Build an index from a list of `(path, title, body)` tuples.
65    ///
66    /// # Errors
67    ///
68    /// [`DocsError::Search`](crate::error::DocsError::Search) on any
69    /// tantivy-internal failure.
70    pub fn build(pages: &[(String, String, String)]) -> Result<Self> {
71        let mut schema_builder = Schema::builder();
72        let path_field = schema_builder.add_text_field("path", STORED);
73        let title_field = schema_builder.add_text_field("title", TEXT | STORED);
74        let body_field = schema_builder.add_text_field("body", TEXT | STORED);
75        let schema = schema_builder.build();
76
77        let index = Index::create_in_ram(schema);
78        {
79            let mut writer: IndexWriter = index.writer(50_000_000)?;
80            for (path, title, body) in pages {
81                let plain = render::to_plain_text(body);
82                writer.add_document(doc!(
83                    path_field => path.as_str(),
84                    title_field => title.as_str(),
85                    body_field => plain,
86                ))?;
87            }
88            writer.commit()?;
89        }
90
91        let query_parser = QueryParser::for_index(&index, vec![title_field, body_field]);
92        let titles: Vec<(String, String)> =
93            pages.iter().map(|(p, t, _)| (p.clone(), t.clone())).collect();
94
95        Ok(Self {
96            index,
97            query_parser,
98            title_field,
99            path_field,
100            body_field,
101            titles,
102            matcher: SkimMatcherV2::default(),
103        })
104    }
105
106    /// Fuzzy-match `query` against every stored title. Empty query
107    /// returns all titles in insertion order with score 0.
108    #[must_use]
109    pub fn title_search(&self, query: &str) -> Vec<TitleHit> {
110        if query.is_empty() {
111            return self
112                .titles
113                .iter()
114                .map(|(path, title)| TitleHit {
115                    path: path.clone(),
116                    title: title.clone(),
117                    score: 0,
118                })
119                .collect();
120        }
121        let mut hits: Vec<TitleHit> = self
122            .titles
123            .iter()
124            .filter_map(|(path, title)| {
125                self.matcher.fuzzy_match(title, query).map(|score| TitleHit {
126                    path: path.clone(),
127                    title: title.clone(),
128                    score,
129                })
130            })
131            .collect();
132        hits.sort_by_key(|h| std::cmp::Reverse(h.score));
133        hits
134    }
135
136    /// Full-text search. Returns up to `limit` hits, highest-score
137    /// first.
138    ///
139    /// # Errors
140    ///
141    /// Surfaces tantivy parse / search failures.
142    pub fn full_text_search(&self, query: &str, limit: usize) -> Result<Vec<FullTextHit>> {
143        if query.trim().is_empty() || limit == 0 {
144            return Ok(Vec::new());
145        }
146        let reader = self.index.reader()?;
147        let searcher = reader.searcher();
148        let Ok(parsed) = self.query_parser.parse_query(query) else {
149            // Malformed queries produce no hits rather than failing
150            // loudly — the user can't always control what they type.
151            return Ok(Vec::new());
152        };
153        // tantivy 0.26: `TopDocs::with_limit` no longer impls
154        // `Collector` directly — must compose with an ordering. Use
155        // `.order_by_score()` for the same default-relevance behaviour
156        // we had under 0.22.
157        let top_docs = searcher.search(&parsed, &TopDocs::with_limit(limit).order_by_score())?;
158        let mut out = Vec::with_capacity(top_docs.len());
159        for (score, doc_address) in top_docs {
160            let retrieved: TantivyDocument = searcher.doc(doc_address)?;
161            let path = extract_first_text(&retrieved, self.path_field).unwrap_or_default();
162            let title = extract_first_text(&retrieved, self.title_field).unwrap_or_default();
163            let body_text = extract_first_text(&retrieved, self.body_field).unwrap_or_default();
164            let snippet = build_snippet(&body_text, query);
165            out.push(FullTextHit { path, title, snippet, score });
166        }
167        Ok(out)
168    }
169}
170
171fn extract_first_text(document: &TantivyDocument, field: tantivy::schema::Field) -> Option<String> {
172    use tantivy::schema::Value;
173    document.get_all(field).next().and_then(|v| v.as_str().map(str::to_string))
174}
175
176fn build_snippet(body: &str, query: &str) -> String {
177    let lowered = body.to_ascii_lowercase();
178    let q = query.to_ascii_lowercase();
179    let Some(idx) = lowered.find(&q) else {
180        // No literal match — return the head of the body.
181        return truncate(body, 140);
182    };
183    let start = idx.saturating_sub(30);
184    let end = (idx + q.len() + 110).min(body.len());
185    let prefix = if start > 0 { "…" } else { "" };
186    let suffix = if end < body.len() { "…" } else { "" };
187    format!("{prefix}{}{suffix}", &body[start..end])
188}
189
190fn truncate(s: &str, max_bytes: usize) -> String {
191    if s.len() <= max_bytes {
192        return s.to_string();
193    }
194    let mut end = max_bytes;
195    while !s.is_char_boundary(end) && end > 0 {
196        end -= 1;
197    }
198    format!("{}…", &s[..end])
199}
200
201// ---------------------------------------------------------------------
202// Tests
203// ---------------------------------------------------------------------
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    fn sample_pages() -> Vec<(String, String, String)> {
210        vec![
211            (
212                "intro.md".into(),
213                "Introduction".into(),
214                "# Introduction\n\nWelcome to the tool.".into(),
215            ),
216            (
217                "config.md".into(),
218                "Configuration".into(),
219                "# Configuration\n\nThe tool reads `config.yaml` for settings.".into(),
220            ),
221            (
222                "shutdown.md".into(),
223                "App context".into(),
224                "# App context\n\nUse a cancellation token to coordinate shutdown.".into(),
225            ),
226        ]
227    }
228
229    #[test]
230    fn build_succeeds_on_empty_input() {
231        let idx = SearchIndex::build(&[]).expect("empty build");
232        assert!(idx.title_search("anything").is_empty());
233    }
234
235    #[test]
236    fn title_search_is_fuzzy() {
237        let idx = SearchIndex::build(&sample_pages()).expect("build");
238        let hits = idx.title_search("conf");
239        assert!(!hits.is_empty(), "'conf' should fuzzy-match Configuration");
240        assert_eq!(hits[0].title, "Configuration");
241    }
242
243    #[test]
244    fn title_search_ranks_best_match_first() {
245        let idx = SearchIndex::build(&sample_pages()).expect("build");
246        let hits = idx.title_search("intro");
247        assert_eq!(hits[0].title, "Introduction");
248    }
249
250    #[test]
251    fn full_text_search_finds_body_tokens() {
252        let idx = SearchIndex::build(&sample_pages()).expect("build");
253        let hits = idx.full_text_search("cancellation", 5).expect("search");
254        assert_eq!(hits.len(), 1, "'cancellation' should hit shutdown.md only");
255        assert_eq!(hits[0].path, "shutdown.md");
256        assert!(
257            hits[0].snippet.to_lowercase().contains("cancellation"),
258            "snippet should contain the query: {}",
259            hits[0].snippet
260        );
261    }
262
263    #[test]
264    fn full_text_search_honours_limit() {
265        let idx = SearchIndex::build(&sample_pages()).expect("build");
266        let hits = idx.full_text_search("tool", 1).expect("search");
267        assert_eq!(hits.len(), 1);
268    }
269
270    #[test]
271    fn full_text_search_rejects_empty_query() {
272        let idx = SearchIndex::build(&sample_pages()).expect("build");
273        let hits = idx.full_text_search("", 5).expect("search");
274        assert!(hits.is_empty());
275    }
276
277    #[test]
278    fn build_snippet_centers_on_match() {
279        let body = "The quick brown fox jumps over the lazy dog for testing purposes";
280        let snippet = build_snippet(body, "jumps");
281        assert!(snippet.contains("jumps"));
282    }
283}