codebook 0.3.42

A code-aware spell checker library (dependency for codebook-lsp)
Documentation
use lru::LruCache;
use unicase::UniCase;

use std::{
    collections::HashSet,
    num::NonZeroUsize,
    path::PathBuf,
    sync::{Arc, RwLock},
};

pub trait Dictionary: Send + Sync {
    fn check(&self, word: &str) -> bool;
    fn suggest(&self, word: &str) -> Vec<String>;
}

/// Failure to load a Hunspell dictionary from its `.aff`/`.dic` file pair.
#[derive(Debug, thiserror::Error)]
pub enum HunspellError {
    #[error("failed to read dictionary file {path}: {source}")]
    Read {
        path: PathBuf,
        source: std::io::Error,
    },
    #[error("failed to parse dictionary [aff: {aff}, dic: {dic}]: {message}")]
    Parse {
        aff: String,
        dic: String,
        message: String,
    },
}

enum WordCase {
    AllCaps,
    AllLower,
    TitleCase,
    Unknown,
}

#[derive(Debug)]
pub struct HunspellDictionary {
    dictionary: spellbook::Dictionary,
    suggestion_cache: Arc<RwLock<LruCache<String, Vec<String>>>>,
    check_cache: Arc<RwLock<LruCache<String, bool>>>,
}

impl HunspellDictionary {
    pub fn new(aff_path: &str, dic_path: &str) -> Result<Self, HunspellError> {
        let read = |path: &str| {
            std::fs::read_to_string(path).map_err(|source| HunspellError::Read {
                path: path.into(),
                source,
            })
        };
        let aff = read(aff_path)?;
        let dic = read(dic_path)?;
        let dict = spellbook::Dictionary::new(&aff, &dic).map_err(|e| HunspellError::Parse {
            aff: aff_path.to_string(),
            dic: dic_path.to_string(),
            message: e.to_string(),
        })?;

        Ok(HunspellDictionary {
            dictionary: dict,
            suggestion_cache: Arc::new(RwLock::new(LruCache::new(
                NonZeroUsize::new(10000).unwrap(),
            ))),
            check_cache: Arc::new(RwLock::new(LruCache::new(
                NonZeroUsize::new(10000).unwrap(),
            ))),
        })
    }
    fn get_word_case(&self, word: &str) -> WordCase {
        if word.chars().all(char::is_uppercase) {
            return WordCase::AllCaps;
        }
        if word.chars().all(char::is_lowercase) {
            return WordCase::AllLower;
        }
        if word.chars().next().unwrap().is_uppercase() {
            return WordCase::TitleCase;
        }
        WordCase::Unknown
    }
}

impl Dictionary for HunspellDictionary {
    fn check(&self, word: &str) -> bool {
        {
            let mut cache = self.check_cache.write().unwrap();
            if let Some(&result) = cache.get(word) {
                return result;
            }
        }

        // If not in cache, perform the check
        let result = self.dictionary.check(word)
            || self
                .dictionary
                .checker()
                .check_lower_as_title(true)
                .check_lower_as_upper(true)
                .check(word);

        // Cache the result
        self.check_cache
            .write()
            .unwrap()
            .put(word.to_string(), result);

        result
    }

    fn suggest(&self, word: &str) -> Vec<String> {
        {
            let mut cache = self.suggestion_cache.write().unwrap();
            if let Some(suggestions) = cache.get(word) {
                return suggestions.clone();
            }
        }

        // If not in cache, generate suggestions
        let mut suggestions = Vec::new();
        self.dictionary.suggest(word, &mut suggestions);
        suggestions.truncate(5);

        if !suggestions.is_empty() {
            let word_case = self.get_word_case(word);

            // Apply case transformations
            for suggestion in &mut suggestions {
                match word_case {
                    WordCase::AllCaps => {
                        suggestion.make_ascii_uppercase();
                    }
                    WordCase::AllLower => {
                        suggestion.make_ascii_lowercase();
                    }
                    WordCase::TitleCase => {
                        // Leave it alone if it's a title case
                    }
                    WordCase::Unknown => {}
                }
            }
        }

        // Cache the result if non-empty
        if !suggestions.is_empty() {
            self.suggestion_cache
                .write()
                .unwrap()
                .put(word.to_string(), suggestions.clone());
        }

        suggestions
    }
}

#[derive(Debug)]
pub struct TextDictionary {
    words: HashSet<UniCase<String>>,
}

impl Dictionary for TextDictionary {
    fn check(&self, word: &str) -> bool {
        self.words.contains(&UniCase::new(word.to_string()))
    }
    fn suggest(&self, _word: &str) -> Vec<String> {
        vec![]
    }
}

impl TextDictionary {
    pub fn new(word_list: &str) -> Self {
        let words = word_list
            .lines()
            .filter(|s| !s.is_empty() && !s.starts_with('#'))
            .map(|s| UniCase::new(s.to_string()))
            .collect();
        Self { words }
    }
    pub fn new_from_path(path: &PathBuf) -> Self {
        let word_list = std::fs::read_to_string(path)
            .unwrap_or_else(|_| panic!("Failed to read dictionary file: {}", path.display()));
        Self::new(&word_list)
    }

    /// Get a reference to the internal HashSet for batch operations
    pub fn word_set(&self) -> &HashSet<UniCase<String>> {
        &self.words
    }
}

#[cfg(test)]
mod dictionary_tests {
    use super::*;
    fn get_dict() -> HunspellDictionary {
        HunspellDictionary::new("./tests/en_index.aff", "./tests/en_index.dic").unwrap()
    }

    #[test]
    fn test_suggest() {
        let dict = get_dict();
        let suggestions = dict.suggest("wrld");
        println!("{suggestions:?}");
        assert!(suggestions.contains(&"world".to_string()));
    }

    #[test]
    fn test_ignore_case() {
        let dict = get_dict();
        let check = dict.check("alice");
        assert!(check);
        let suggestions = dict.suggest("alice");
        println!("{suggestions:?}");
        assert!(suggestions.contains(&"alice".to_string()));
    }

    #[test]
    fn test_text_dictionary_unicode_ignore_case() {
        let dict = TextDictionary::new("Апгрейдить\nИИ\n");

        assert!(dict.check("апгрейдить"));
        assert!(dict.check("АПГРЕЙДИТЬ"));
        assert!(dict.check("ии"));
        assert!(dict.check("ИИ"));
    }

    #[test]
    fn test_text_dictionary_preserves_original_case() {
        let dict = TextDictionary::new("Straße\n");

        assert!(dict.check("straße"));
        assert!(dict.check("STRAßE"));
        // Stored entry keeps its original casing rather than being lowercased.
        assert!(
            dict.word_set()
                .contains(&UniCase::new("Straße".to_string()))
        );
    }

    #[test]
    fn test_text_dictionary_full_case_folding() {
        // unicase handles cases that to_lowercase() cannot collapse.
        let dict = TextDictionary::new("Straße\nΣίγμα\n");

        // ß <-> SS: full case folding
        assert!(dict.check("STRASSE"));
        assert!(dict.check("strasse"));
        assert!(dict.check("Straße"));

        // Greek sigma (Σ/σ/ς) all fold together
        assert!(dict.check("ΣΊΓΜΑ"));
        assert!(dict.check("σίγμα"));
    }
}