lb_rs/search/

content.rs

use super::path::split_path;
use super::{ContentMatch, SearchFilter, SearchResult, build_descendants};
use crate::blocking::Lb;
use crate::model::file::File;
use std::cmp::Reverse;
use std::collections::{HashMap, HashSet};
use std::ops::Range;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};
use uuid::Uuid;

struct Document {
    file: File,
    filename: String,
    parent_path: String,
    lowercased_content: String,
}

pub struct ContentSearcher {
    documents: Vec<Document>,
    results: Vec<SearchResult>,
    submitted_query: String,
    build_time: Duration,
    descendants: HashMap<Uuid, Vec<Uuid>>,
    path_to_id: HashMap<String, Uuid>,
    filter_ids: Option<HashSet<Uuid>>,
}

impl ContentSearcher {
    pub fn new(lb: &Lb) -> Self {
        let start = Instant::now();
        let metas = lb.list_metadatas().unwrap_or_default();
        let paths = Arc::new(lb.list_paths_with_ids(None).unwrap_or_default());

        let descendants = build_descendants(&metas);
        let path_to_id = paths.iter().map(|(id, path)| (path.clone(), *id)).collect();

        let md_files: Vec<File> = metas
            .into_iter()
            .filter(|m| m.is_document() && m.name.ends_with(".md"))
            .collect();

        let queue = Arc::new(Mutex::new(md_files));
        let documents = Arc::new(Mutex::new(Vec::<Document>::new()));

        let handles: Vec<_> = (0..thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(4))
            .map(|_| {
                let queue = queue.clone();
                let paths = paths.clone();
                let documents = documents.clone();
                let lb = lb.clone();
                thread::spawn(move || {
                    loop {
                        let Some(meta) = queue.lock().unwrap().pop() else {
                            return;
                        };

                        let id = meta.id;
                        let doc = lb
                            .read_document(meta.id, false)
                            .ok()
                            .and_then(|bytes| String::from_utf8(bytes).ok());

                        if let Some(content) = doc {
                            let path = paths
                                .iter()
                                .find(|(i, _)| *i == id)
                                .map(|(_, p)| p.clone())
                                .unwrap_or_default();

                            let (parent, name) = split_path(&path);

                            documents.lock().unwrap().push(Document {
                                file: meta,
                                filename: name.to_string(),
                                parent_path: parent.to_string(),
                                lowercased_content: content.to_lowercase(),
                            });
                        }
                    }
                })
            })
            .collect();

        for h in handles {
            h.join().unwrap();
        }

        let documents = Arc::try_unwrap(documents)
            .ok()
            .expect("all workers joined")
            .into_inner()
            .unwrap();

        Self {
            documents,
            results: Vec::new(),
            submitted_query: String::new(),
            build_time: start.elapsed(),
            descendants,
            path_to_id,
            filter_ids: None,
        }
    }

    /// Update the active filter and refresh results for the current query.
    pub fn update_filter(&mut self, filter: Option<SearchFilter>) {
        self.filter_ids = super::resolve_filter(filter, &self.path_to_id, &self.descendants);
        self.rebuild();
    }

    /// Update the search query. Results available via `results()`.
    pub fn query(&mut self, input: &str) {
        let query = input.trim().to_lowercase();
        if self.submitted_query == query {
            return;
        }
        self.submitted_query = query;
        self.rebuild();
    }

    /// Rebuild `results` from the current query and filter. Shared by `query` and
    /// `update_filter`.
    fn rebuild(&mut self) {
        self.results.clear();
        if self.submitted_query.is_empty() {
            return;
        }

        let query = &self.submitted_query;
        let words: Vec<&str> = query.split_whitespace().collect();

        let mut scored = Vec::new();

        for doc in &self.documents {
            if let Some(ids) = &self.filter_ids {
                if !ids.contains(&doc.file.id) {
                    continue;
                }
            }

            let parent = doc.parent_path.to_lowercase();
            let name = doc.filename.to_lowercase();
            let (path, name_start) = if parent == "/" {
                (format!("/{name}"), 1)
            } else {
                let name_start = parent.len() + 1;
                (format!("{parent}/{name}"), name_start)
            };

            let mut matched_words = vec![false; words.len()];
            let path_matches = collect_matches(&path, query, &words, &mut matched_words);
            let content_matches =
                collect_matches(&doc.lowercased_content, query, &words, &mut matched_words);

            let all_words_matched = matched_words.iter().all(|&m| m);
            let has_match = !path_matches.is_empty() || !content_matches.is_empty();

            if all_words_matched && has_match {
                let filename_exact = path_matches
                    .iter()
                    .filter(|m| m.exact && m.range.end > name_start)
                    .count();
                let filename_partial = path_matches
                    .iter()
                    .filter(|m| !m.exact && m.range.end > name_start)
                    .count();
                let content_exact = content_matches.iter().filter(|m| m.exact).count();

                let tier = if filename_exact > 0 {
                    0
                } else if content_exact > 0 {
                    1
                } else if filename_partial > 0 {
                    2
                } else {
                    3
                };

                let key = (
                    tier,
                    Reverse(filename_exact),
                    Reverse(content_exact),
                    Reverse(filename_partial),
                    Reverse(path_matches.len()),
                    Reverse(content_matches.len()),
                    path,
                );

                scored.push((
                    key,
                    SearchResult {
                        id: doc.file.id,
                        filename: doc.filename.clone(),
                        parent_path: doc.parent_path.clone(),
                        is_folder: false,
                        path_indices: Vec::new(),
                        path_matches,
                        content_matches,
                    },
                ));
            }
        }

        scored.sort_by(|a, b| a.0.cmp(&b.0));
        self.results = scored.into_iter().map(|(_, r)| r).collect();
    }

    /// Get search results.
    pub fn results(&self) -> &[SearchResult] {
        &self.results
    }

    /// How long it took to build the search index.
    pub fn build_time(&self) -> Duration {
        self.build_time
    }

    /// Get document content by ID.
    pub fn content(&self, id: Uuid) -> Option<&str> {
        self.documents
            .iter()
            .find(|d| d.file.id == id)
            .map(|d| d.lowercased_content.as_str())
    }

    /// Extract snippet with context. Returns (prefix, matched, suffix).
    pub fn snippet(
        &self, id: Uuid, range: &Range<usize>, context_chars: usize,
    ) -> Option<(&str, &str, &str)> {
        let content = self.content(id)?;
        let char_indices: Vec<(usize, char)> = content.char_indices().collect();

        let match_start_ci = char_indices
            .iter()
            .position(|(byte, _)| *byte >= range.start)
            .unwrap_or(char_indices.len());
        let match_end_ci = char_indices
            .iter()
            .position(|(byte, _)| *byte >= range.end)
            .unwrap_or(char_indices.len());

        let start_ci = match_start_ci.saturating_sub(context_chars);
        let end_ci = (match_end_ci + context_chars).min(char_indices.len());

        // Snap to whitespace
        let mut start = start_ci;
        if start > 0 {
            for (i, &(_, c)) in char_indices
                .iter()
                .enumerate()
                .take(match_start_ci)
                .skip(start)
            {
                if c.is_whitespace() {
                    start = i + 1;
                    break;
                }
            }
        }

        let mut end = end_ci;
        if end < char_indices.len() {
            for i in (match_end_ci..end).rev() {
                if char_indices[i].1.is_whitespace() {
                    end = i;
                    break;
                }
            }
        }

        let get_byte = |ci: usize| {
            char_indices
                .get(ci)
                .map(|(b, _)| *b)
                .unwrap_or(content.len())
        };

        let start_byte = get_byte(start);
        let match_start_byte = get_byte(match_start_ci);
        let match_end_byte = get_byte(match_end_ci);
        let end_byte = get_byte(end);

        Some((
            &content[start_byte..match_start_byte],
            &content[match_start_byte..match_end_byte],
            &content[match_end_byte..end_byte],
        ))
    }
}

fn collect_matches(
    haystack: &str, query: &str, words: &[&str], matched_words: &mut [bool],
) -> Vec<ContentMatch> {
    let mut matches = Vec::new();

    for (idx, _) in haystack.match_indices(query) {
        matches.push(ContentMatch { range: idx..idx + query.len(), exact: true });
    }

    for (wi, word) in words.iter().enumerate() {
        for (idx, _) in haystack.match_indices(word) {
            matched_words[wi] = true;
            if matches.iter().any(|m| m.range.contains(&idx)) {
                continue;
            }
            matches.push(ContentMatch { range: idx..idx + word.len(), exact: false });
        }
    }

    matches
}