markon_core/
search.rs

1use serde::{Deserialize, Serialize};
2use std::{
3    fs,
4    path::{Path, PathBuf},
5    sync::{Arc, Mutex},
6};
7use tantivy::{
8    collector::TopDocs, query::QueryParser, schema::*, snippet::SnippetGenerator, Index,
9    IndexReader, IndexWriter, TantivyDocument,
10};
11use tantivy_jieba::JiebaTokenizer;
12use walkdir::WalkDir;
13
14#[derive(Deserialize)]
15pub struct SearchQuery {
16    pub q: String,
17}
18
19#[derive(Serialize)]
20pub struct SearchResult {
21    pub file_path: String,
22    pub file_name: String,
23    pub title: String,
24    pub snippet: String,
25}
26
27pub struct SearchIndex {
28    index: Index,
29    reader: IndexReader,
30    writer: Arc<Mutex<IndexWriter>>,
31    #[allow(dead_code)]
32    schema: Schema,
33    field_path: Field,
34    field_file_name: Field,
35    field_title: Field,
36    field_content: Field,
37    start_dir: PathBuf,
38}
39
40impl SearchIndex {
41    pub fn new(start_dir: &Path) -> tantivy::Result<Self> {
42        // Build schema
43        let mut schema_builder = Schema::builder();
44
45        let text_options = TextOptions::default()
46            .set_indexing_options(
47                TextFieldIndexing::default()
48                    .set_tokenizer("jieba")
49                    .set_index_option(IndexRecordOption::WithFreqsAndPositions),
50            )
51            .set_stored();
52
53        // Use STRING for path field - indexed but not tokenized, so we can delete by exact match
54        let field_path = schema_builder.add_text_field("path", STRING | STORED);
55        let field_file_name = schema_builder.add_text_field("file_name", text_options.clone());
56        let field_title = schema_builder.add_text_field("title", text_options.clone());
57        let field_content = schema_builder.add_text_field("content", text_options);
58
59        let schema = schema_builder.build();
60
61        // Create index in RAM
62        let index = Index::create_in_ram(schema.clone());
63
64        // Register jieba tokenizer
65        index.tokenizers().register("jieba", JiebaTokenizer {});
66
67        // Create writer and reader
68        let writer = index.writer(50_000_000)?;
69        let reader = index.reader()?;
70
71        let search_index = Self {
72            index,
73            reader,
74            writer: Arc::new(Mutex::new(writer)),
75            schema,
76            field_path,
77            field_file_name,
78            field_title,
79            field_content,
80            start_dir: start_dir.to_path_buf(),
81        };
82
83        // Index all markdown files
84        search_index.index_directory(start_dir)?;
85
86        Ok(search_index)
87    }
88
89    fn index_directory(&self, dir: &Path) -> tantivy::Result<()> {
90        use rayon::prelude::*;
91
92        println!("Indexing markdown files in {dir:?}...");
93
94        let paths: Vec<PathBuf> = WalkDir::new(dir)
95            .into_iter()
96            .filter_map(Result::ok)
97            .filter(|e| e.path().extension().is_some_and(|ext| ext == "md"))
98            .map(|e| e.into_path())
99            .collect();
100
101        paths.par_iter().for_each(|path| {
102            if let Ok(content) = fs::read_to_string(path) {
103                let _ = self.index_file(path, &content);
104            }
105        });
106
107        let mut writer = self.writer.lock().unwrap();
108        writer.commit()?;
109        drop(writer);
110
111        self.reader.reload()?;
112        println!("Indexing complete!");
113
114        Ok(())
115    }
116
117    fn index_file(&self, path: &Path, content: &str) -> tantivy::Result<()> {
118        // Calculate relative path from start_dir
119        let relative_path = path
120            .strip_prefix(&self.start_dir)
121            .unwrap_or(path)
122            .to_string_lossy()
123            .to_string();
124
125        let file_name = path
126            .file_stem()
127            .and_then(|s| s.to_str())
128            .unwrap_or("")
129            .to_string();
130
131        // Extract title from first heading or filename
132        let title = content
133            .lines()
134            .find(|line| line.starts_with('#'))
135            .map(|line| line.trim_start_matches('#').trim().to_string())
136            .unwrap_or_else(|| file_name.clone());
137
138        let mut doc = TantivyDocument::default();
139        doc.add_text(self.field_path, &relative_path);
140        doc.add_text(self.field_file_name, &file_name);
141        doc.add_text(self.field_title, &title);
142        doc.add_text(self.field_content, content);
143
144        let writer = self.writer.lock().unwrap();
145        writer.add_document(doc)?;
146
147        Ok(())
148    }
149
150    pub fn search(&self, query_str: &str, limit: usize) -> tantivy::Result<Vec<SearchResult>> {
151        let searcher = self.reader.searcher();
152
153        // Search across file_name, title, and content
154        let query_parser = QueryParser::for_index(
155            &self.index,
156            vec![self.field_file_name, self.field_title, self.field_content],
157        );
158
159        let query = query_parser.parse_query(query_str)?;
160        let top_docs = searcher.search(&query, &TopDocs::with_limit(limit))?;
161
162        let mut results = Vec::new();
163        let snippet_generator = SnippetGenerator::create(&searcher, &query, self.field_content)?;
164
165        for (_score, doc_address) in top_docs {
166            let retrieved_doc: TantivyDocument = searcher.doc(doc_address)?;
167
168            let file_path = retrieved_doc
169                .get_first(self.field_path)
170                .and_then(|v| v.as_str())
171                .unwrap_or("")
172                .to_string();
173
174            let file_name = retrieved_doc
175                .get_first(self.field_file_name)
176                .and_then(|v| v.as_str())
177                .unwrap_or("")
178                .to_string();
179
180            let title = retrieved_doc
181                .get_first(self.field_title)
182                .and_then(|v| v.as_str())
183                .unwrap_or("")
184                .to_string();
185
186            let snippet = snippet_generator.snippet_from_doc(&retrieved_doc);
187            let snippet_html = snippet.to_html();
188
189            results.push(SearchResult {
190                file_path,
191                file_name,
192                title,
193                snippet: snippet_html,
194            });
195        }
196
197        Ok(results)
198    }
199
200    pub fn update_file(&self, path: &Path) -> tantivy::Result<()> {
201        if path.extension().is_none_or(|ext| ext != "md") {
202            return Ok(());
203        }
204
205        let content = fs::read_to_string(path)?;
206        // Calculate relative path from start_dir
207        let relative_path = path
208            .strip_prefix(&self.start_dir)
209            .unwrap_or(path)
210            .to_string_lossy()
211            .to_string();
212
213        let file_name = path
214            .file_stem()
215            .and_then(|s| s.to_str())
216            .unwrap_or("")
217            .to_string();
218
219        let title = content
220            .lines()
221            .find(|line| line.starts_with('#'))
222            .map(|line| line.trim_start_matches('#').trim().to_string())
223            .unwrap_or_else(|| file_name.clone());
224
225        let mut doc = TantivyDocument::default();
226        doc.add_text(self.field_path, &relative_path);
227        doc.add_text(self.field_file_name, &file_name);
228        doc.add_text(self.field_title, &title);
229        doc.add_text(self.field_content, &content);
230
231        // Delete and re-add in same transaction
232        let mut writer = self.writer.lock().unwrap();
233        let term = Term::from_field_text(self.field_path, &relative_path);
234        writer.delete_term(term);
235        writer.add_document(doc)?;
236        writer.commit()?;
237        drop(writer);
238
239        // Reload reader to see the changes
240        self.reader.reload()?;
241
242        println!("Updated index: {}", relative_path);
243        Ok(())
244    }
245
246    pub fn delete_file(&self, path: &Path) -> tantivy::Result<()> {
247        // Calculate relative path from start_dir
248        let relative_path = path
249            .strip_prefix(&self.start_dir)
250            .unwrap_or(path)
251            .to_string_lossy()
252            .to_string();
253
254        let mut writer = self.writer.lock().unwrap();
255        let term = Term::from_field_text(self.field_path, &relative_path);
256        writer.delete_term(term);
257        writer.commit()?;
258        drop(writer);
259
260        // Reload reader to see the changes
261        self.reader.reload()?;
262
263        println!("Removed from index: {}", relative_path);
264        Ok(())
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271    use std::fs;
272    use tempfile::TempDir;
273
274    fn create_test_file(dir: &Path, name: &str, content: &str) -> std::io::Result<()> {
275        let file_path = dir.join(name);
276        fs::write(file_path, content)
277    }
278
279    #[test]
280    fn test_search_index_creation() {
281        let temp_dir = TempDir::new().unwrap();
282        let dir_path = temp_dir.path();
283
284        // Create test markdown files
285        create_test_file(dir_path, "test1.md", "# Test Title\nThis is test content.").unwrap();
286        create_test_file(dir_path, "test2.md", "# Another Title\nMore content here.").unwrap();
287
288        let index = SearchIndex::new(dir_path).unwrap();
289
290        // Verify index was created
291        assert!(index.reader.searcher().num_docs() >= 2);
292    }
293
294    #[test]
295    fn test_search_basic() {
296        let temp_dir = TempDir::new().unwrap();
297        let dir_path = temp_dir.path();
298
299        create_test_file(
300            dir_path,
301            "rust.md",
302            "# Rust Programming\nRust is a systems programming language.",
303        )
304        .unwrap();
305        create_test_file(
306            dir_path,
307            "python.md",
308            "# Python Guide\nPython is easy to learn.",
309        )
310        .unwrap();
311
312        let index = SearchIndex::new(dir_path).unwrap();
313
314        let results = index.search("Rust", 10).unwrap();
315        assert_eq!(results.len(), 1);
316        assert!(results[0].title.contains("Rust"));
317    }
318
319    #[test]
320    fn test_search_chinese() {
321        let temp_dir = TempDir::new().unwrap();
322        let dir_path = temp_dir.path();
323
324        create_test_file(
325            dir_path,
326            "chinese.md",
327            "# 中文测试\n这是一个中文搜索测试文档。",
328        )
329        .unwrap();
330        create_test_file(
331            dir_path,
332            "english.md",
333            "# English Test\nThis is an English document.",
334        )
335        .unwrap();
336
337        let index = SearchIndex::new(dir_path).unwrap();
338
339        // Test Chinese search
340        let results = index.search("中文", 10).unwrap();
341        assert_eq!(results.len(), 1);
342        assert!(results[0].title.contains("中文"));
343
344        // Test mixed search
345        let results = index.search("测试", 10).unwrap();
346        assert_eq!(results.len(), 1);
347    }
348
349    #[test]
350    fn test_search_multi_field() {
351        let temp_dir = TempDir::new().unwrap();
352        let dir_path = temp_dir.path();
353
354        create_test_file(
355            dir_path,
356            "hello.md",
357            "# Hello World\nContent about greetings.",
358        )
359        .unwrap();
360        create_test_file(dir_path, "world.md", "# Universe\nThe world is vast.").unwrap();
361        create_test_file(dir_path, "test.md", "# Testing\nAnother document.").unwrap();
362
363        let index = SearchIndex::new(dir_path).unwrap();
364
365        // Search should find "greetings" in content only
366        let results = index.search("greetings", 10).unwrap();
367        assert_eq!(results.len(), 1);
368
369        // Search for "Hello" should find in title
370        let results = index.search("Hello", 10).unwrap();
371        assert!(!results.is_empty());
372    }
373
374    #[test]
375    fn test_update_file() {
376        let temp_dir = TempDir::new().unwrap();
377        let dir_path = temp_dir.path();
378        let file_path = dir_path.join("update.md");
379
380        // Create initial file
381        create_test_file(dir_path, "update.md", "# Original Title\nOriginal content.").unwrap();
382
383        let index = SearchIndex::new(dir_path).unwrap();
384
385        // Verify original content
386        let results = index.search("Original", 10).unwrap();
387        assert_eq!(results.len(), 1);
388
389        // Update file
390        fs::write(&file_path, "# Updated Title\nUpdated content.").unwrap();
391        index.update_file(&file_path).unwrap();
392
393        // Search for new content
394        let results = index.search("Updated", 10).unwrap();
395        assert_eq!(results.len(), 1);
396
397        // Old content should not be found
398        let results = index.search("Original", 10).unwrap();
399        assert_eq!(results.len(), 0);
400    }
401
402    #[test]
403    fn test_delete_file() {
404        let temp_dir = TempDir::new().unwrap();
405        let dir_path = temp_dir.path();
406        let file_path = dir_path.join("delete.md");
407
408        create_test_file(dir_path, "delete.md", "# Delete Me\nThis will be deleted.").unwrap();
409
410        let index = SearchIndex::new(dir_path).unwrap();
411
412        // Verify file is indexed
413        let results = index.search("Delete", 10).unwrap();
414        assert_eq!(results.len(), 1);
415
416        // Delete file
417        index.delete_file(&file_path).unwrap();
418
419        // Verify file is removed from index
420        let results = index.search("Delete", 10).unwrap();
421        assert_eq!(results.len(), 0);
422    }
423
424    #[test]
425    fn test_search_snippet_generation() {
426        let temp_dir = TempDir::new().unwrap();
427        let dir_path = temp_dir.path();
428
429        create_test_file(
430            dir_path,
431            "snippet.md",
432            "# Snippet Test\nThis document contains important information about snippets. Snippets are useful."
433        ).unwrap();
434
435        let index = SearchIndex::new(dir_path).unwrap();
436
437        let results = index.search("snippets", 10).unwrap();
438        assert_eq!(results.len(), 1);
439        assert!(!results[0].snippet.is_empty());
440        // Snippet should contain highlighted text (HTML tags)
441        assert!(results[0].snippet.contains("<b>") || results[0].snippet.contains("snippet"));
442    }
443
444    #[test]
445    fn test_ignore_non_markdown_files() {
446        let temp_dir = TempDir::new().unwrap();
447        let dir_path = temp_dir.path();
448
449        create_test_file(dir_path, "test.md", "# Markdown\nThis is markdown.").unwrap();
450        create_test_file(dir_path, "test.txt", "# Not Markdown\nThis is text.").unwrap();
451
452        let index = SearchIndex::new(dir_path).unwrap();
453
454        // Should find markdown file
455        let results = index.search("Markdown", 10).unwrap();
456        assert_eq!(results.len(), 1);
457
458        // Should not find text file
459        let results = index.search("text", 10).unwrap();
460        assert_eq!(results.len(), 0);
461    }
462
463    #[test]
464    fn test_empty_query() {
465        let temp_dir = TempDir::new().unwrap();
466        let dir_path = temp_dir.path();
467
468        create_test_file(dir_path, "test.md", "# Test\nContent").unwrap();
469
470        let index = SearchIndex::new(dir_path).unwrap();
471
472        // Empty query should return error or empty results
473        let results = index.search("", 10);
474        assert!(results.is_err() || results.unwrap().is_empty());
475    }
476
477    #[test]
478    fn test_title_extraction() {
479        let temp_dir = TempDir::new().unwrap();
480        let dir_path = temp_dir.path();
481
482        // File with heading
483        create_test_file(dir_path, "with-title.md", "# Main Title\nContent").unwrap();
484
485        // File without heading (should use filename)
486        create_test_file(dir_path, "no-title.md", "Just content without heading").unwrap();
487
488        let index = SearchIndex::new(dir_path).unwrap();
489
490        let results = index.search("Main", 10).unwrap();
491        assert_eq!(results.len(), 1);
492        assert_eq!(results[0].title, "Main Title");
493
494        let results = index.search("no-title", 10).unwrap();
495        assert_eq!(results.len(), 1);
496        assert_eq!(results[0].title, "no-title");
497    }
498
499    #[test]
500    fn test_search_limit() {
501        let temp_dir = TempDir::new().unwrap();
502        let dir_path = temp_dir.path();
503
504        // Create multiple files with common content
505        for i in 0..10 {
506            create_test_file(
507                dir_path,
508                &format!("file{}.md", i),
509                "# Document\nCommon content",
510            )
511            .unwrap();
512        }
513
514        let index = SearchIndex::new(dir_path).unwrap();
515
516        // Test limit works
517        let results = index.search("Common", 5).unwrap();
518        assert_eq!(results.len(), 5);
519
520        let results = index.search("Common", 20).unwrap();
521        assert_eq!(results.len(), 10);
522    }
523
524    #[test]
525    fn test_subdirectory_relative_paths() {
526        let temp_dir = TempDir::new().unwrap();
527        let sub = temp_dir.path().join("notes").join("deep");
528        fs::create_dir_all(&sub).unwrap();
529        create_test_file(&sub, "nested.md", "# Nested\nDeep content here").unwrap();
530
531        let index = SearchIndex::new(temp_dir.path()).unwrap();
532        let results = index.search("Deep", 10).unwrap();
533        assert_eq!(results.len(), 1);
534        // Path should be relative, using forward slashes
535        let path = &results[0].file_path;
536        assert!(
537            path.starts_with("notes"),
538            "expected relative path, got: {path}"
539        );
540        assert!(
541            path.contains("nested"),
542            "expected file name in path, got: {path}"
543        );
544        assert!(
545            !path.starts_with('/'),
546            "path should be relative, got: {path}"
547        );
548    }
549
550    #[test]
551    fn test_update_file_ignores_non_markdown() {
552        let temp_dir = TempDir::new().unwrap();
553        create_test_file(temp_dir.path(), "test.md", "# Original\nMarkdown").unwrap();
554        let index = SearchIndex::new(temp_dir.path()).unwrap();
555
556        // Write a .txt file and try to update — should be no-op
557        let txt_path = temp_dir.path().join("notes.txt");
558        fs::write(&txt_path, "Some text content").unwrap();
559        // Should not error
560        index.update_file(&txt_path).unwrap();
561
562        // Searching for the txt content should yield nothing
563        let results = index.search("Some text content", 10).unwrap();
564        assert!(results.is_empty());
565    }
566
567    #[test]
568    fn test_update_file_no_extension() {
569        let temp_dir = TempDir::new().unwrap();
570        create_test_file(temp_dir.path(), "test.md", "# Doc\nContent").unwrap();
571        let index = SearchIndex::new(temp_dir.path()).unwrap();
572
573        let no_ext = temp_dir.path().join("README");
574        fs::write(&no_ext, "Plain text").unwrap();
575        index.update_file(&no_ext).unwrap();
576
577        let results = index.search("Plain text", 10).unwrap();
578        assert!(results.is_empty());
579    }
580}
markon_core/search.rs

markon_core/
search.rs
