Full-Text Search Engine

const STOP_WORDS: &[&str] = &["a", "an", "and", "are", "in", "is", "it", "of", "on", "or", "the", "to", "with"];

/// A tiny stemmer: the plural rules from the first step of the Porter algorithm.
pub fn stem(word: &str) -> String {
    if let Some(base) = word.strip_suffix("sses") {
        format!("{base}ss")
    } else if let Some(base) = word.strip_suffix("ies") {
        format!("{base}y")
    } else if word.ends_with("ss") || word.len() <= 3 {
        word.to_string()
    } else if let Some(base) = word.strip_suffix('s') {
        base.to_string()
    } else {
        word.to_string()
    }
}

pub fn tokenize(text: &str) -> impl Iterator<Item = String> + '_ {
    text.split(|c: char| !c.is_alphanumeric())
        .filter(|word| word.len() > 1) // single letters, like the "s" in "runtime's", carry no meaning
        .map(str::to_lowercase)
        .filter(|word| !STOP_WORDS.contains(&word.as_str()))
        .map(|word| stem(&word))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn terms(text: &str) -> Vec<String> {
        tokenize(text).collect()
    }

    #[test]
    fn lowercases_splits_and_drops_stop_words() {
        assert_eq!(terms("The Runtime's Garbage-Collector"), ["runtime", "garbage", "collector"]);
    }

    #[test]
    fn stems_plurals() {
        assert_eq!(terms("threads classes libraries futures"), ["thread", "class", "library", "future"]);
    }

    #[test]
    fn leaves_short_words_and_double_s_alone() {
        assert_eq!(terms("gas bus process"), ["gas", "bus", "process"]);
    }
}

use std::collections::HashMap;

fn tokenize(text: &str) -> impl Iterator<Item = String> + '_ {
    text.split(|c: char| !c.is_alphanumeric()).filter(|w| !w.is_empty()).map(str::to_lowercase)
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct Posting {
    doc: usize,
    term_frequency: u32,
}

#[derive(Debug, Default)]
struct Index {
    postings: HashMap<String, Vec<Posting>>,
    doc_lengths: HashMap<usize, u32>,
}

impl Index {
    fn add(&mut self, doc: usize, text: &str) {
        let mut counts: HashMap<String, u32> = HashMap::new();
        for term in tokenize(text) {
            *counts.entry(term).or_insert(0) += 1;
        }
        self.doc_lengths.insert(doc, counts.values().sum());
        for (term, term_frequency) in counts {
            self.postings.entry(term).or_default().push(Posting { doc, term_frequency });
        }
    }

    fn lookup(&self, term: &str) -> &[Posting] {
        self.postings.get(term).map_or(&[], Vec::as_slice)
    }
}

fn main() {
    let mut index = Index::default();
    index.add(0, "Rust uses ownership. Ownership replaces the garbage collector.");
    index.add(1, "The Go garbage collector frees memory.");
    index.add(2, "Traits are like interfaces.");

    for term in ["ownership", "garbage", "interfaces", "python"] {
        println!("{term:<11} {:?}", index.lookup(term));
    }
    println!("{} distinct terms, doc 0 has {} terms", index.postings.len(), index.doc_lengths[&0]);
}

use std::collections::HashMap;

struct Posting {
    doc: usize,
    term_frequency: u32,
}

const K1: f64 = 1.2;
const B: f64 = 0.75;

fn bm25(postings: &HashMap<&str, Vec<Posting>>, doc_lengths: &[u32], query: &[&str]) -> Vec<(usize, f64)> {
    let docs = doc_lengths.len() as f64;
    let avg_len = doc_lengths.iter().sum::<u32>() as f64 / docs;
    let mut scores: HashMap<usize, f64> = HashMap::new();

    for term in query {
        let Some(list) = postings.get(term) else { continue };
        let idf = (1.0 + (docs - list.len() as f64 + 0.5) / (list.len() as f64 + 0.5)).ln();
        for posting in list {
            let tf = posting.term_frequency as f64;
            let length_ratio = doc_lengths[posting.doc] as f64 / avg_len;
            let score = idf * tf * (K1 + 1.0) / (tf + K1 * (1.0 - B + B * length_ratio));
            *scores.entry(posting.doc).or_insert(0.0) += score;
        }
    }

    let mut ranked: Vec<(usize, f64)> = scores.into_iter().collect();
    // f64 isn't Ord (NaN), so sort with total_cmp; ties by document id for stable output.
    ranked.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
    ranked
}

fn main() {
    // Five documents: 0 and 1 mention "rust"; only 3 mentions "borrow".
    let doc_lengths = [8, 40, 10, 12, 9];
    let postings = HashMap::from([
        ("rust", vec![Posting { doc: 0, term_frequency: 1 }, Posting { doc: 1, term_frequency: 3 }]),
        ("borrow", vec![Posting { doc: 3, term_frequency: 1 }]),
        ("code", vec![
            Posting { doc: 0, term_frequency: 1 },
            Posting { doc: 1, term_frequency: 1 },
            Posting { doc: 2, term_frequency: 1 },
            Posting { doc: 3, term_frequency: 1 },
        ]),
    ]);

    for query in [&["rust"][..], &["code"], &["borrow", "code"]] {
        let ranked: Vec<String> = bm25(&postings, &doc_lengths, query).iter().map(|(d, s)| format!("doc {d}: {s:.3}")).collect();
        println!("{query:?} -> {}", ranked.join(", "));
    }
}

#[derive(Debug, Default, PartialEq)]
struct Query {
    optional: Vec<String>,
    required: Vec<String>,
    excluded: Vec<String>,
}

fn parse_query(text: &str) -> Query {
    let mut query = Query::default();
    for word in text.split_whitespace() {
        let (list, term) = match word.as_bytes()[0] {
            b'+' => (&mut query.required, &word[1..]),
            b'-' => (&mut query.excluded, &word[1..]),
            _ => (&mut query.optional, word),
        };
        if !term.is_empty() {
            list.push(term.to_lowercase());
        }
    }
    query
}

struct Document {
    title: &'static str,
    body: &'static str,
}

struct Hit<'a> {
    doc: &'a Document,
    snippet: String,
}

/// Up to `width` words around the first word that matches, with matches in [brackets].
fn snippet(body: &str, terms: &[String], width: usize) -> String {
    let words: Vec<&str> = body.split_whitespace().collect();
    let matches = |word: &str| {
        let clean = word.trim_matches(|c: char| !c.is_alphanumeric()).to_lowercase();
        terms.iter().any(|t| clean.starts_with(t.as_str()))
    };
    let first = words.iter().position(|w| matches(w)).unwrap_or(0);
    let start = first.saturating_sub(width / 2);
    let end = (start + width).min(words.len());
    let shown: Vec<String> =
        words[start..end].iter().map(|w| if matches(w) { format!("[{w}]") } else { w.to_string() }).collect();
    let prefix = if start > 0 { "... " } else { "" };
    let suffix = if end < words.len() { " ..." } else { "" };
    format!("{prefix}{}{suffix}", shown.join(" "))
}

fn search<'a>(docs: &'a [Document], query: &Query) -> Vec<Hit<'a>> {
    let contains = |doc: &Document, term: &str| format!("{} {}", doc.title, doc.body).to_lowercase().contains(term);
    let terms: Vec<String> = query.required.iter().chain(&query.optional).cloned().collect();
    docs.iter()
        .filter(|doc| query.required.iter().all(|t| contains(doc, t)))
        .filter(|doc| !query.excluded.iter().any(|t| contains(doc, t)))
        .filter(|doc| terms.iter().any(|t| contains(doc, t)))
        .map(|doc| Hit { doc, snippet: snippet(doc.body, &terms, 8) })
        .collect()
}

fn main() {
    println!("{:?}", parse_query("threads +rust -go"));

    let docs = [
        Document { title: "Async Rust", body: "Async functions return futures, and Tokio runs those futures on a pool of worker threads." },
        Document { title: "Go goroutines", body: "Goroutines make blocking code cheap because the Go runtime schedules them on system threads." },
        Document { title: "Rust threads", body: "std::thread::spawn starts an operating system thread that runs a closure." },
    ];
    for text in ["threads", "threads -go", "+rust future"] {
        println!("query: {text}");
        for hit in search(&docs, &parse_query(text)) {
            println!("  {:<22}{}", hit.doc.title, hit.snippet);
        }
    }
}

use rayon::prelude::*;
use std::collections::{HashMap, HashSet};

// ---------- tokenize.rs ----------

const STOP_WORDS: &[&str] = &["a", "an", "and", "are", "by", "in", "is", "it", "of", "on", "or", "the", "them", "they", "to", "with"];

fn stem(word: &str) -> String {
    if let Some(base) = word.strip_suffix("sses") {
        format!("{base}ss")
    } else if let Some(base) = word.strip_suffix("ies") {
        format!("{base}y")
    } else if word.ends_with("ss") || word.len() <= 3 {
        word.to_string()
    } else if let Some(base) = word.strip_suffix('s') {
        base.to_string()
    } else {
        word.to_string()
    }
}

fn tokenize(text: &str) -> impl Iterator<Item = String> + '_ {
    text.split(|c: char| !c.is_alphanumeric())
        .filter(|word| word.len() > 1) // single letters, like the "s" in "runtime's", carry no meaning
        .map(str::to_lowercase)
        .filter(|word| !STOP_WORDS.contains(&word.as_str()))
        .map(|word| stem(&word))
}

// ---------- index.rs ----------

struct Document {
    id: usize,
    title: String,
    body: String,
}

#[derive(Clone, Copy)]
struct Posting {
    doc: usize,
    term_frequency: u32,
}

#[derive(Default)]
struct Index {
    postings: HashMap<String, Vec<Posting>>,
    doc_lengths: HashMap<usize, u32>,
}

impl Index {
    fn build(docs: &[Document]) -> Index {
        let mut index = Index::default();
        for doc in docs {
            let mut counts: HashMap<String, u32> = HashMap::new();
            for term in tokenize(&doc.title).chain(tokenize(&doc.body)) {
                *counts.entry(term).or_insert(0) += 1;
            }
            index.doc_lengths.insert(doc.id, counts.values().sum());
            for (term, term_frequency) in counts {
                index.postings.entry(term).or_default().push(Posting { doc: doc.id, term_frequency });
            }
        }
        index
    }

    fn merge(mut self, other: Index) -> Index {
        for (term, list) in other.postings {
            self.postings.entry(term).or_default().extend(list);
        }
        self.doc_lengths.extend(other.doc_lengths);
        self
    }

    fn lookup(&self, term: &str) -> &[Posting] {
        self.postings.get(term).map_or(&[], Vec::as_slice)
    }
}

// ---------- search.rs ----------

#[derive(Default)]
struct Query {
    optional: Vec<String>,
    required: Vec<String>,
    excluded: Vec<String>,
}

fn parse_query(text: &str) -> Query {
    let mut query = Query::default();
    for word in text.split_whitespace() {
        let (list, rest) = match word.as_bytes()[0] {
            b'+' => (&mut query.required, &word[1..]),
            b'-' => (&mut query.excluded, &word[1..]),
            _ => (&mut query.optional, word),
        };
        list.extend(tokenize(rest)); // query words are normalized exactly like documents
    }
    query
}

struct Hit<'a> {
    doc: &'a Document,
    score: f64,
}

struct Engine {
    docs: Vec<Document>,
    index: Index,
}

const K1: f64 = 1.2;
const B: f64 = 0.75;

impl Engine {
    fn new(docs: Vec<Document>) -> Engine {
        let index = docs.par_chunks(3).map(Index::build).reduce(Index::default, Index::merge);
        Engine { docs, index }
    }

    fn docs_with(&self, term: &str) -> HashSet<usize> {
        self.index.lookup(term).iter().map(|p| p.doc).collect()
    }

    fn search(&self, text: &str) -> Vec<Hit<'_>> {
        let query = parse_query(text);
        let count = self.docs.len() as f64;
        let avg_len = self.index.doc_lengths.values().sum::<u32>() as f64 / count;

        let mut scores: HashMap<usize, f64> = HashMap::new();
        for term in query.required.iter().chain(&query.optional) {
            let list = self.index.lookup(term);
            let idf = (1.0 + (count - list.len() as f64 + 0.5) / (list.len() as f64 + 0.5)).ln();
            for posting in list {
                let tf = posting.term_frequency as f64;
                let length_ratio = self.index.doc_lengths[&posting.doc] as f64 / avg_len;
                *scores.entry(posting.doc).or_insert(0.0) += idf * tf * (K1 + 1.0) / (tf + K1 * (1.0 - B + B * length_ratio));
            }
        }

        for term in &query.required {
            let allowed = self.docs_with(term);
            scores.retain(|doc, _| allowed.contains(doc));
        }
        for term in &query.excluded {
            let banned = self.docs_with(term);
            scores.retain(|doc, _| !banned.contains(doc));
        }

        let mut hits: Vec<Hit<'_>> = scores.into_iter().map(|(id, score)| Hit { doc: &self.docs[id], score }).collect();
        hits.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.doc.id.cmp(&b.doc.id)));
        hits
    }
}

/// Up to `width` words around the first match, with matching words in [brackets].
fn snippet(body: &str, query: &str, width: usize) -> String {
    let wanted: HashSet<String> = tokenize(query).collect();
    let words: Vec<&str> = body.split_whitespace().collect();
    let matches = |word: &str| tokenize(word).any(|term| wanted.contains(&term));
    let first = words.iter().position(|w| matches(w)).unwrap_or(0);
    let start = first.saturating_sub(width / 2);
    let end = (start + width).min(words.len());
    let shown: Vec<String> = words[start..end].iter().map(|w| if matches(w) { format!("[{w}]") } else { w.to_string() }).collect();
    format!("{}{}{}", if start > 0 { "... " } else { "" }, shown.join(" "), if end < words.len() { " ..." } else { "" })
}

// ---------- main.rs ----------

const CORPUS: &[(&str, &str)] = &[
    ("Ownership in Rust", "Rust uses ownership and borrowing instead of a garbage collector. Each value has exactly one owner."),
    ("Garbage collection in Go", "The Go garbage collector frees memory automatically, so Go developers rarely manage memory by hand."),
    ("Rust traits", "Traits define shared behavior. In practice a trait is similar to an interface with default methods."),
    ("Go interfaces", "An interface declares methods that types implement. Go interfaces are satisfied implicitly, without a declaration."),
    ("Async Rust", "Async functions return futures. Tokio runs futures on a pool of worker threads."),
    ("Go goroutines", "Goroutines make blocking code cheap. The Go runtime schedules them on system threads."),
    ("Rust error handling", "Rust returns errors as values with Result instead of throwing exceptions."),
    ("Python exceptions", "Python code raises exceptions and catches them in try blocks. Exceptions are not part of a function signature."),
    ("Memory safety", "Rust guarantees memory safety at compile time, without a garbage collector or runtime checks."),
];

fn main() {
    let docs: Vec<Document> =
        CORPUS.iter().enumerate().map(|(id, (title, body))| Document { id, title: title.to_string(), body: body.to_string() }).collect();
    let engine = Engine::new(docs);
    println!("indexed {} documents, {} distinct terms\n", engine.docs.len(), engine.index.postings.len());

    for query in ["garbage collector", "interface trait", "threads -go", "+rust errors", "memory", "haskell"] {
        println!("query: {query}");
        let hits = engine.search(query);
        if hits.is_empty() {
            println!("  no results");
        }
        for hit in hits.iter().take(3) {
            println!("  {:.2}  {:<28}{}", hit.score, hit.doc.title, snippet(&hit.doc.body, query, 9));
        }
        println!();
    }
}