use std::collections::HashSet; use std::fs; use std::path::{Path, PathBuf}; use std::sync::Mutex; use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; use tantivy::collector::{Count, TopDocs}; use tantivy::directory::MmapDirectory; use tantivy::query::{ AllQuery, BooleanQuery, BoostQuery, EmptyQuery, Occur, PhraseQuery, Query, TermQuery, }; use tantivy::schema::{ Field, IndexRecordOption, STORED, STRING, Schema, TextFieldIndexing, TextOptions, Value, }; use tantivy::snippet::SnippetGenerator; use tantivy::tokenizer::TokenStream; use tantivy::{Index, IndexReader, IndexWriter, TantivyDocument, Term, doc}; use crate::extract::{extract_file, summary_of, supported}; use crate::message::{EXPOSURE_FULL, ResponseItem}; use crate::tokenizer::{TOKENIZER_NAME, analyzer}; const WRITER_BUDGET: usize = 50_000_000; const MAX_QUERY_TERMS: usize = 10; const MAX_COMBOS: usize = 256; const DEFAULT_MIN_COVERAGE: f64 = 0.4; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Collection { pub name: String, pub path: String, pub shared: bool, pub exposure: String, } #[derive(Debug, Default, Serialize, Deserialize)] struct CollectionsFile { #[serde(default, rename = "collection")] collections: Vec, } pub fn load_collections(manifest: &Path) -> Result> { if !manifest.exists() { return Ok(Vec::new()); } let raw = fs::read_to_string(manifest).context("reading collections manifest")?; let parsed: CollectionsFile = toml::from_str(&raw).context("parsing collections manifest")?; Ok(parsed.collections) } pub fn save_collections(manifest: &Path, collections: &[Collection]) -> Result<()> { if let Some(parent) = manifest.parent() { fs::create_dir_all(parent)?; } let file = CollectionsFile { collections: collections.to_vec(), }; fs::write(manifest, toml::to_string_pretty(&file)?)?; Ok(()) } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SearchHit { pub url: String, pub title: String, pub summary: String, pub published: String, pub exposure: String, pub collection: String, pub shared: bool, pub content: Option, } pub fn response_items(hits: &[SearchHit]) -> Vec { hits.iter() .map(|hit| ResponseItem { url: hit.url.clone(), title: hit.title.clone(), summary: hit.summary.clone(), published: hit.published.clone(), exposure: hit.exposure.clone(), content: if hit.exposure == EXPOSURE_FULL { hit.content.clone() } else { None }, }) .collect() } struct Fields { url: Field, title: Field, body: Field, summary: Field, published: Field, exposure: Field, collection: Field, shared: Field, } pub struct LocalIndex { index: Index, reader: IndexReader, writer: Mutex, fields: Fields, } impl LocalIndex { pub fn open(dir: &Path) -> Result { fs::create_dir_all(dir)?; let schema = build_schema(); let index = Index::open_or_create(MmapDirectory::open(dir)?, schema) .context("opening tantivy index")?; index.tokenizers().register(TOKENIZER_NAME, analyzer()); let schema = index.schema(); let fields = Fields { url: schema.get_field("url")?, title: schema.get_field("title")?, body: schema.get_field("body")?, summary: schema.get_field("summary")?, published: schema.get_field("published")?, exposure: schema.get_field("exposure")?, collection: schema.get_field("collection")?, shared: schema.get_field("shared")?, }; let reader = index.reader()?; let writer = index.writer::(WRITER_BUDGET)?; Ok(Self { index, reader, writer: Mutex::new(writer), fields, }) } pub fn add_collection(&self, collection: &Collection) -> Result { let root = PathBuf::from(&collection.path); { let writer = self.writer.lock().expect("writer lock"); writer.delete_term(Term::from_field_text( self.fields.collection, &collection.name, )); } let mut files = Vec::new(); walk(&root, &mut files); let mut added = 0; for file in files { if self.add_file(collection, &file)? { added += 1; } } self.commit()?; Ok(added) } pub fn add_file(&self, collection: &Collection, path: &Path) -> Result { let Some(extracted) = extract_file(path) else { return Ok(false); }; let canonical = path.canonicalize().unwrap_or_else(|_| path.to_path_buf()); let url = format!("file://{}", canonical.display()); let summary = summary_of(&extracted.body); let shared = if collection.shared { "true" } else { "false" }; let writer = self.writer.lock().expect("writer lock"); writer.delete_term(Term::from_field_text(self.fields.url, &url)); writer.add_document(doc!( self.fields.url => url.as_str(), self.fields.title => extracted.title.as_str(), self.fields.body => extracted.body.as_str(), self.fields.summary => summary.as_str(), self.fields.published => extracted.published.as_str(), self.fields.exposure => collection.exposure.as_str(), self.fields.collection => collection.name.as_str(), self.fields.shared => shared, ))?; Ok(true) } pub fn commit(&self) -> Result<()> { self.writer.lock().expect("writer lock").commit()?; self.reader.reload()?; Ok(()) } pub fn search( &self, text: &str, limit: usize, only_shared: bool, ) -> Result<(Vec, u64)> { self.search_with_coverage(text, limit, only_shared, DEFAULT_MIN_COVERAGE) } pub fn search_with_coverage( &self, text: &str, limit: usize, only_shared: bool, min_coverage: f64, ) -> Result<(Vec, u64)> { let limit = limit.max(1); self.reader.reload()?; let searcher = self.reader.searcher(); let terms = self.query_terms(text); let user_query: Box = if text.trim().is_empty() { Box::new(AllQuery) } else if terms.is_empty() { Box::new(EmptyQuery) } else { let gate = self.coverage_query(&terms, min_coverage); let mut parts: Vec<(Occur, Box)> = vec![(Occur::Must, gate)]; if terms.len() >= 2 { let phrase_terms: Vec<(usize, Term)> = terms .iter() .enumerate() .map(|(index, term)| (index, Term::from_field_text(self.fields.body, term))) .collect(); parts.push(( Occur::Should, Box::new(BoostQuery::new( Box::new(PhraseQuery::new_with_offset_and_slop(phrase_terms, 1)), 3.0, )), )); } Box::new(BooleanQuery::new(parts)) }; let query: Box = if only_shared { let shared_term = TermQuery::new( Term::from_field_text(self.fields.shared, "true"), IndexRecordOption::Basic, ); Box::new(BooleanQuery::new(vec![ (Occur::Must, user_query), (Occur::Must, Box::new(shared_term)), ])) } else { user_query }; let total = searcher.search(&*query, &Count)? as u64; let top = searcher.search(&*query, &TopDocs::with_limit(limit).order_by_score())?; let snippet_generator = SnippetGenerator::create(&searcher, &*query, self.fields.body).ok(); let mut hits = Vec::with_capacity(top.len()); for (_score, address) in top { let document: TantivyDocument = searcher.doc(address)?; let snippet = snippet_generator .as_ref() .map(|generator| generator.snippet_from_doc(&document).fragment().to_string()) .filter(|fragment| !fragment.trim().is_empty()); hits.push(SearchHit { url: text_value(&document, self.fields.url), title: text_value(&document, self.fields.title), summary: snippet.unwrap_or_else(|| text_value(&document, self.fields.summary)), published: text_value(&document, self.fields.published), exposure: text_value(&document, self.fields.exposure), collection: text_value(&document, self.fields.collection), shared: text_value(&document, self.fields.shared) == "true", content: Some(text_value(&document, self.fields.body)), }); } Ok((hits, total)) } fn query_terms(&self, text: &str) -> Vec { let mut terms = Vec::new(); let mut seen = HashSet::new(); if let Some(mut tokenizer) = self.index.tokenizers().get(TOKENIZER_NAME) { let mut stream = tokenizer.token_stream(text); while let Some(token) = stream.next() { let term = token.text.to_string(); if seen.insert(term.clone()) { terms.push(term); } if terms.len() >= MAX_QUERY_TERMS { break; } } } terms } fn coverage_query(&self, terms: &[String], min_coverage: f64) -> Box { let n = terms.len(); let mut required = if n <= 2 { n } else { ((n as f64) * min_coverage).ceil() as usize }; required = required.clamp(1, n); let combos = combinations(n, required, MAX_COMBOS); let mut shoulds: Vec<(Occur, Box)> = Vec::with_capacity(combos.len()); for combo in combos { let musts: Vec<(Occur, Box)> = combo .iter() .map(|index| { let term = terms[*index].as_str(); let term_query: Box = Box::new(BooleanQuery::new(vec![ ( Occur::Should, Box::new(BoostQuery::new( Box::new(TermQuery::new( Term::from_field_text(self.fields.title, term), IndexRecordOption::WithFreqs, )), 2.0, )), ), ( Occur::Should, Box::new(TermQuery::new( Term::from_field_text(self.fields.body, term), IndexRecordOption::WithFreqs, )), ), ])); (Occur::Must, term_query) }) .collect(); shoulds.push((Occur::Should, Box::new(BooleanQuery::new(musts)))); } Box::new(BooleanQuery::new(shoulds)) } pub fn doc_count(&self) -> u64 { self.reader.searcher().num_docs() } } fn binomial(n: usize, k: usize) -> usize { if k > n { return 0; } let k = k.min(n - k); let mut result = 1usize; for index in 0..k { result = result * (n - index) / (index + 1); } result } fn combinations(n: usize, start_required: usize, cap: usize) -> Vec> { let mut required = start_required; while required < n && binomial(n, required) > cap { required += 1; } fn walk( start: usize, remaining: usize, n: usize, current: &mut Vec, out: &mut Vec>, ) { if remaining == 0 { out.push(current.clone()); return; } for index in start..=n - remaining { current.push(index); walk(index + 1, remaining - 1, n, current, out); current.pop(); } } let mut out = Vec::new(); walk(0, required, n, &mut Vec::new(), &mut out); out } fn text_value(document: &TantivyDocument, field: Field) -> String { document .get_first(field) .and_then(|value| value.as_str()) .unwrap_or("") .to_string() } fn build_schema() -> Schema { let mut builder = Schema::builder(); let frx_options = TextOptions::default() .set_indexing_options( TextFieldIndexing::default() .set_tokenizer(TOKENIZER_NAME) .set_index_option(IndexRecordOption::WithFreqsAndPositions), ) .set_stored(); builder.add_text_field("url", STRING | STORED); builder.add_text_field("title", frx_options.clone()); builder.add_text_field("body", frx_options); builder.add_text_field("summary", STORED); builder.add_text_field("published", STORED); builder.add_text_field("exposure", STRING | STORED); builder.add_text_field("collection", STRING | STORED); builder.add_text_field("shared", STRING | STORED); builder.build() } fn walk(root: &Path, out: &mut Vec) { let Ok(entries) = fs::read_dir(root) else { return; }; for entry in entries.flatten() { let Ok(file_type) = entry.file_type() else { continue; }; if file_type.is_symlink() { continue; } let name = entry.file_name(); let name = name.to_string_lossy(); if file_type.is_dir() { if name.starts_with('.') { continue; } walk(&entry.path(), out); } else if file_type.is_file() && supported(&entry.path()) { out.push(entry.path()); } } } #[cfg(test)] mod tests { use super::*; use crate::message::EXPOSURE_METADATA; fn collection(name: &str, root: &Path, shared: bool, exposure: &str) -> Collection { Collection { name: name.to_string(), path: root.display().to_string(), shared, exposure: exposure.to_string(), } } #[test] fn shared_filter_and_exposure_are_enforced() { let temp = tempfile::tempdir().unwrap(); let shared_dir = temp.path().join("shared"); let private_dir = temp.path().join("private"); fs::create_dir_all(&shared_dir).unwrap(); fs::create_dir_all(&private_dir).unwrap(); fs::write(shared_dir.join("a.txt"), "rust ownership rules").unwrap(); fs::write(private_dir.join("b.txt"), "rust secret notes").unwrap(); let index = LocalIndex::open(&temp.path().join("index")).unwrap(); index .add_collection(&collection("shared", &shared_dir, true, EXPOSURE_FULL)) .unwrap(); index .add_collection(&collection( "private", &private_dir, false, EXPOSURE_METADATA, )) .unwrap(); let (hits, total) = index.search("rust", 10, false).unwrap(); assert_eq!(total, 2); assert_eq!(hits.len(), 2); let (shared_hits, shared_total) = index.search("rust", 10, true).unwrap(); assert_eq!(shared_total, 1); assert_eq!(shared_hits.len(), 1); assert!(shared_hits[0].url.contains("shared")); let items = response_items(&shared_hits); assert!(items[0].content.is_some()); } #[test] fn coverage_gate_rejects_weak_single_term_overlap() { let temp = tempfile::tempdir().unwrap(); let dir = temp.path().join("corpus"); fs::create_dir_all(&dir).unwrap(); fs::write(dir.join("errors.txt"), "error codes and diagnostics").unwrap(); let index = LocalIndex::open(&temp.path().join("index")).unwrap(); index .add_collection(&collection("test", &dir, true, EXPOSURE_FULL)) .unwrap(); let (hits, total) = index.search("launch codes", 10, true).unwrap(); assert_eq!(total, 0, "two-term queries require both terms"); assert!(hits.is_empty()); } #[test] fn coverage_gate_allows_partial_long_queries() { let temp = tempfile::tempdir().unwrap(); let dir = temp.path().join("corpus"); fs::create_dir_all(&dir).unwrap(); fs::write( dir.join("manual.txt"), "LG DLEX5555 dryer circuit diagram and wiring", ) .unwrap(); let index = LocalIndex::open(&temp.path().join("index")).unwrap(); index .add_collection(&collection("test", &dir, true, EXPOSURE_FULL)) .unwrap(); let (hits, total) = index .search("LG dryer 5555 circuit diagram", 10, true) .unwrap(); assert_eq!(total, 1, "four of five terms is enough"); assert_eq!(hits.len(), 1); } #[test] fn reindex_removes_deleted_files() { let temp = tempfile::tempdir().unwrap(); let dir = temp.path().join("corpus"); fs::create_dir_all(&dir).unwrap(); fs::write(dir.join("a.txt"), "rust one").unwrap(); fs::write(dir.join("b.txt"), "rust two").unwrap(); let index = LocalIndex::open(&temp.path().join("index")).unwrap(); index .add_collection(&collection("test", &dir, true, EXPOSURE_FULL)) .unwrap(); assert_eq!(index.search("rust", 10, false).unwrap().1, 2); fs::remove_file(dir.join("b.txt")).unwrap(); index .add_collection(&collection("test", &dir, true, EXPOSURE_FULL)) .unwrap(); assert_eq!(index.search("rust", 10, false).unwrap().1, 1); } #[test] fn metadata_exposure_withholds_content() { let temp = tempfile::tempdir().unwrap(); let dir = temp.path().join("corpus"); fs::create_dir_all(&dir).unwrap(); fs::write(dir.join("a.txt"), "metadata only").unwrap(); let index = LocalIndex::open(&temp.path().join("index")).unwrap(); index .add_collection(&collection("test", &dir, true, EXPOSURE_METADATA)) .unwrap(); let (hits, _) = index.search("metadata", 10, true).unwrap(); let items = response_items(&hits); assert!(items[0].content.is_none()); assert_eq!(items[0].exposure, EXPOSURE_METADATA); } }