searchez 1.0.0

// searchez/src/memory.rs
//
// The batteries-included backend: an in-memory inverted index with BM25 ranking.
// No server, no external service — the "runs on your laptop" default, and the
// reference implementation of `Backend`. Real relevance (BM25 is what Lucene /
// Elasticsearch / Tantivy use), so search results are ordered sensibly, not by
// naive term counts.
//
// Scope: it holds everything in memory and is not persistent, so it fits tests,
// development, and small single-process datasets. Point a server backend at the
// same `Searchable` models when you outgrow it — nothing above the trait changes.

use crate::backend::{Backend, IndexDoc};
use crate::document::Document;
use crate::error::Result;
use crate::query::{Hit, Query};
use crate::token::{tokenize, tokenize_value};
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::RwLock;

/// BM25 term-frequency saturation. The Lucene/Elasticsearch default.
const K1: f32 = 1.2;
/// BM25 length-normalization strength. The Lucene/Elasticsearch default.
const B: f32 = 0.75;

/// One stored record: the document as given, plus its precomputed text terms and
/// length, so search doesn't re-tokenize on every query.
struct Stored {
    document: Document,
    /// term -> frequency in this document.
    term_freq: HashMap<String, u32>,
    /// total term count (document length), for BM25 length normalization.
    len: u32,
}

impl Stored {
    fn build(document: Document) -> Self {
        let mut tokens = Vec::new();
        for (_field, value) in document.iter() {
            tokenize_value(value, &mut tokens);
        }
        let len = tokens.len() as u32;
        let mut term_freq = HashMap::new();
        for t in tokens {
            *term_freq.entry(t).or_insert(0) += 1;
        }
        Self { document, term_freq, len }
    }
}

#[derive(Default)]
struct Index {
    /// id -> stored document.
    docs: HashMap<String, Stored>,
    /// term -> set of ids containing it (the inverted index, for IDF + candidate
    /// lookup without scanning every document).
    postings: HashMap<String, std::collections::HashSet<String>>,
    /// running sum of document lengths, for average-length in BM25.
    total_len: u64,
}

impl Index {
    fn avgdl(&self) -> f32 {
        if self.docs.is_empty() {
            0.0
        } else {
            self.total_len as f32 / self.docs.len() as f32
        }
    }

    fn remove(&mut self, id: &str) {
        if let Some(old) = self.docs.remove(id) {
            self.total_len -= old.len as u64;
            for term in old.term_freq.keys() {
                if let Some(set) = self.postings.get_mut(term) {
                    set.remove(id);
                    if set.is_empty() {
                        self.postings.remove(term);
                    }
                }
            }
        }
    }

    fn insert(&mut self, id: String, document: Document) {
        // Replace supersedes: clear the old posting entries first.
        self.remove(&id);
        let stored = Stored::build(document);
        self.total_len += stored.len as u64;
        for term in stored.term_freq.keys() {
            self.postings.entry(term.clone()).or_default().insert(id.clone());
        }
        self.docs.insert(id, stored);
    }
}

/// In-memory, BM25-ranked search backend. Construct with [`MemoryBackend::new`].
#[derive(Default)]
pub struct MemoryBackend {
    // One lock over all indexes keeps the code simple; writes are short and this
    // backend targets modest datasets. A server backend is the answer at scale.
    indexes: RwLock<HashMap<String, Index>>,
}

impl MemoryBackend {
    pub fn new() -> Self {
        Self::default()
    }
}

/// Whether a stored document satisfies every filter (exact field equality).
fn passes_filters(document: &Document, query: &Query) -> bool {
    query
        .filters
        .iter()
        .all(|f| document.get(&f.field) == Some(&f.value))
}

#[async_trait]
impl Backend for MemoryBackend {
    async fn upsert(&self, index: &str, docs: Vec<IndexDoc>) -> Result<()> {
        let mut map = self.indexes.write().unwrap();
        let idx = map.entry(index.to_string()).or_default();
        for d in docs {
            idx.insert(d.id, d.document);
        }
        Ok(())
    }

    async fn delete(&self, index: &str, ids: &[String]) -> Result<()> {
        let mut map = self.indexes.write().unwrap();
        if let Some(idx) = map.get_mut(index) {
            for id in ids {
                idx.remove(id);
            }
        }
        Ok(())
    }

    async fn search(&self, index: &str, query: &Query) -> Result<Vec<Hit>> {
        let map = self.indexes.read().unwrap();
        let Some(idx) = map.get(index) else {
            return Ok(Vec::new());
        };

        let query_terms: Vec<String> = query
            .text
            .as_deref()
            .map(tokenize)
            .unwrap_or_default();

        let n = idx.docs.len() as f32;
        let avgdl = idx.avgdl();

        // Precompute IDF for each distinct query term. BM25's IDF with the
        // +0.5/+1 smoothing that keeps it non-negative.
        let mut idf: HashMap<&str, f32> = HashMap::new();
        for term in &query_terms {
            if idf.contains_key(term.as_str()) {
                continue;
            }
            let df = idx.postings.get(term).map(|s| s.len()).unwrap_or(0) as f32;
            let value = (1.0 + (n - df + 0.5) / (df + 0.5)).ln();
            idf.insert(term.as_str(), value);
        }

        let mut hits: Vec<Hit> = Vec::new();
        for (id, stored) in &idx.docs {
            if !passes_filters(&stored.document, query) {
                continue;
            }

            let score = if query_terms.is_empty() {
                // Pure filter query: no ranking signal, everyone scores 0.
                0.0
            } else {
                let mut s = 0.0f32;
                for term in &query_terms {
                    let tf = *stored.term_freq.get(term).unwrap_or(&0) as f32;
                    if tf == 0.0 {
                        continue;
                    }
                    let norm = tf * (K1 + 1.0)
                        / (tf + K1 * (1.0 - B + B * stored.len as f32 / avgdl.max(1.0)));
                    s += idf[term.as_str()] * norm;
                }
                s
            };

            // A text query keeps only documents that matched at least one term;
            // a filter-only query keeps everything that passed the filters.
            if query_terms.is_empty() || score > 0.0 {
                hits.push(Hit {
                    id: id.clone(),
                    score,
                    document: stored.document.clone(),
                });
            }
        }

        // Rank by score desc; break ties by id so paging is stable and
        // deterministic (float scores have no total order otherwise).
        hits.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.id.cmp(&b.id))
        });

        let limit = query.effective_limit();
        Ok(hits.into_iter().skip(query.offset).take(limit).collect())
    }

    async fn clear(&self, index: &str) -> Result<()> {
        self.indexes.write().unwrap().remove(index);
        Ok(())
    }

    async fn count(&self, index: &str) -> Result<usize> {
        Ok(self
            .indexes
            .read()
            .unwrap()
            .get(index)
            .map(|i| i.docs.len())
            .unwrap_or(0))
    }
}