path: root/src/suggest.rs

//! Offline suggestions: «¿quiso decir…?» ("did you mean…?").
//!
//! Uses the hunspell dictionary that already ships with the system
//! (`/usr/share/hunspell/es_*.dic`): about 57,000 stems, one per line, with
//! the inflection flags after a slash (`anhelar/REDÀA`).
//!
//! Matching ignores accents, so «cafe» finds «café» and «arbol» finds
//! «árbol», the most common typing mistake in Spanish.

use std::path::{Path, PathBuf};

use anyhow::Result;

/// Where the dictionary usually lives, in order of preference.
const CANDIDATES: &[&str] = &[
    "/usr/share/hunspell",
    "/usr/share/myspell/dicts",
    "/usr/local/share/hunspell",
];

/// Finds an installed Spanish `.dic`. Prefers the user's regional variant
/// (`$LANG`) and falls back to any `es_*`.
pub fn find_dictionary() -> Option<PathBuf> {
    let preferred = std::env::var("LANG")
        .ok()
        .and_then(|l| l.split('.').next().map(str::to_string))
        .filter(|l| l.starts_with("es"));

    let mut found: Vec<PathBuf> = Vec::new();
    for dir in CANDIDATES {
        let Ok(entries) = std::fs::read_dir(dir) else {
            continue;
        };
        for e in entries.flatten() {
            let p = e.path();
            let name = p.file_name().and_then(|n| n.to_str()).unwrap_or("");
            if name.starts_with("es") && name.ends_with(".dic") {
                found.push(p);
            }
        }
    }
    found.sort();
    if let Some(pref) = preferred {
        if let Some(hit) = found
            .iter()
            .find(|p| p.file_stem().and_then(|s| s.to_str()) == Some(pref.as_str()))
        {
            return Some(hit.clone());
        }
    }
    found.into_iter().next()
}

/// Loads the stems from the `.dic`, without the inflection flags.
pub fn load_words(path: &Path) -> Result<Vec<String>> {
    let raw = std::fs::read_to_string(path)?;
    let mut out = Vec::with_capacity(60_000);
    // The first line is the entry count, not a word.
    for line in raw.lines().skip(1) {
        let word = line.split('/').next().unwrap_or("").trim();
        // Skip entries that contain digits.
        if word.is_empty() || word.chars().any(|c| c.is_ascii_digit()) {
            continue;
        }
        out.push(word.to_string());
    }
    Ok(out)
}

/// The `limit` words closest to `query`.
pub fn closest(query: &str, words: &[String], limit: usize) -> Vec<String> {
    let needle = fold(query);
    // Past this distance the suggestions stop looking alike.
    let max = match needle.chars().count() {
        0..=4 => 1,
        5..=7 => 2,
        _ => 3,
    };

    // Someone typing in lowercase wants a common word, not a proper noun:
    // the .dic includes Abel, Argos or Perseo, which are just noise.
    let lowercase_query = !query.chars().any(char::is_uppercase);

    let mut scored: Vec<(usize, usize, bool, &String)> = Vec::new();
    for w in words {
        if lowercase_query && w.chars().next().is_some_and(char::is_uppercase) {
            continue;
        }
        let cand = fold(w);
        // A cheap filter before the expensive computation.
        if cand.len().abs_diff(needle.len()) > max + 2 {
            continue;
        }
        let d = distance(&needle, &cand, max);
        if d <= max {
            // On equal distance, the closest length wins, then the one that
            // starts with the same letter.
            let len_gap = cand.chars().count().abs_diff(needle.chars().count());
            let same_start = cand.chars().next() == needle.chars().next();
            scored.push((d, len_gap, !same_start, w));
        }
    }
    scored.sort_by(|a, b| (a.0, a.1, a.2, a.3).cmp(&(b.0, b.1, b.2, b.3)));
    scored.dedup_by(|a, b| a.3 == b.3);
    scored
        .into_iter()
        .take(limit)
        .map(|(_, _, _, w)| w.clone())
        .collect()
}

/// Lowercase and without accents, so the comparison does not penalize them.
pub fn fold(s: &str) -> String {
    s.trim()
        .to_lowercase()
        .chars()
        .map(|c| match c {
            'á' => 'a',
            'é' => 'e',
            'í' => 'i',
            'ó' => 'o',
            'ú' | 'ü' => 'u',
            other => other,
        })
        .collect()
}

/// Levenshtein distance, giving up as soon as it exceeds `max`.
fn distance(a: &str, b: &str, max: usize) -> usize {
    let a: Vec<char> = a.chars().collect();
    let b: Vec<char> = b.chars().collect();
    if a.is_empty() {
        return b.len();
    }
    if b.is_empty() {
        return a.len();
    }
    let mut prev: Vec<usize> = (0..=b.len()).collect();
    let mut cur = vec![0usize; b.len() + 1];
    for i in 1..=a.len() {
        cur[0] = i;
        let mut row_min = cur[0];
        for j in 1..=b.len() {
            let cost = usize::from(a[i - 1] != b[j - 1]);
            cur[j] = (prev[j] + 1).min(cur[j - 1] + 1).min(prev[j - 1] + cost);
            row_min = row_min.min(cur[j]);
        }
        if row_min > max {
            return max + 1;
        }
        std::mem::swap(&mut prev, &mut cur);
    }
    prev[b.len()]
}

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

    fn words(list: &[&str]) -> Vec<String> {
        list.iter().map(|s| s.to_string()).collect()
    }

    #[test]
    fn folds_accents_and_case() {
        assert_eq!(fold("Café"), "cafe");
        assert_eq!(fold("ÁRBOL"), "arbol");
        // ñ is a letter of its own, not an accented n: it is kept.
        assert_eq!(fold("Año"), "año");
    }

    #[test]
    fn suggests_the_accented_spelling() {
        let w = words(&["café", "cafetera", "calle", "casa"]);
        assert_eq!(
            closest("cafe", &w, 3).first().map(String::as_str),
            Some("café")
        );
    }

    #[test]
    fn suggests_for_a_typo() {
        let w = words(&["perro", "pero", "puerro", "casa"]);
        let s = closest("perrro", &w, 3);
        assert!(s.contains(&"perro".to_string()), "suggestions: {s:?}");
    }

    #[test]
    fn skips_proper_nouns_for_a_lowercase_query() {
        let w = words(&["Perseo", "perro", "perrero"]);
        let s = closest("perrro", &w, 5);
        assert!(!s.contains(&"Perseo".to_string()), "suggestions: {s:?}");
        assert!(s.contains(&"perro".to_string()));
    }

    #[test]
    fn keeps_proper_nouns_when_the_query_is_capitalised() {
        let w = words(&["Perseo", "perro"]);
        assert!(closest("Perseu", &w, 5).contains(&"Perseo".to_string()));
    }

    #[test]
    fn returns_nothing_for_gibberish() {
        let w = words(&["casa", "perro", "hablar"]);
        assert!(closest("xkcdqwzy", &w, 5).is_empty());
    }

    #[test]
    fn respects_the_limit() {
        let w = words(&["casa", "cara", "caza", "cosa", "capa", "cava"]);
        assert!(closest("casa", &w, 3).len() <= 3);
    }

    #[test]
    fn distance_gives_up_past_the_ceiling() {
        // Very different: returns max+1 without computing the exact value.
        assert_eq!(distance("abcdefgh", "zzzzzzzz", 2), 3);
        assert_eq!(distance("casa", "cosa", 2), 1);
        assert_eq!(distance("casa", "casa", 2), 0);
    }
}