1use 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#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct TitleHit {
30 pub path: String,
32 pub title: String,
34 pub score: i64,
36}
37
38#[derive(Debug, Clone)]
40pub struct FullTextHit {
41 pub path: String,
43 pub title: String,
45 pub snippet: String,
48 pub score: f32,
50}
51
52pub 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)>, matcher: SkimMatcherV2,
61}
62
63impl SearchIndex {
64 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 #[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 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 return Ok(Vec::new());
152 };
153 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 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#[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}