lang_check/

dictionary.rs

use anyhow::Result;
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use tracing::{debug, warn};

use crate::morphology::inflection;

/// Manages custom dictionaries for the language checker.
/// Words in the dictionary are excluded from spelling diagnostics.
///
/// Internally, user-added words and bundled words are kept in separate sets.
/// Only user words are persisted to `dictionary.txt`.
pub struct Dictionary {
    user_words: HashSet<String>,
    bundled_words: HashSet<String>,
    /// Regular inflections generated from the other two sets by
    /// [`crate::morphology::inflection`].
    ///
    /// Third set rather than folded into `bundled_words` so that [`Self::persist`] is
    /// correct by construction: it reads `user_words` alone, and no derived form can
    /// ever reach the file that records what the user actually typed.
    derived_words: HashSet<String>,
    workspace_path: Option<PathBuf>,
}

impl Default for Dictionary {
    fn default() -> Self {
        Self::new()
    }
}

impl Dictionary {
    #[must_use]
    pub fn new() -> Self {
        Self {
            user_words: HashSet::new(),
            bundled_words: HashSet::new(),
            derived_words: HashSet::new(),
            workspace_path: None,
        }
    }

    /// Load dictionaries from a workspace root.
    /// Reads from .languagecheck/dictionary.txt (one word per line).
    pub fn load(workspace_root: &Path) -> Result<Self> {
        let mut dict = Self::new();
        let dict_path = workspace_root.join(".languagecheck").join("dictionary.txt");
        dict.workspace_path = Some(dict_path.clone());

        if dict_path.exists() {
            let content = std::fs::read_to_string(&dict_path)?;
            for line in content.lines() {
                let word = line.trim();
                if !word.is_empty() && !word.starts_with('#') {
                    dict.user_words.insert(word.to_lowercase());
                }
            }
        }

        Ok(dict)
    }

    /// Load the bundled dictionaries that ship with the extension.
    /// These contain domain-specific technical terms from open-source wordlists.
    /// Bundled words are kept separate and never persisted to the user's dictionary file.
    pub fn load_bundled(&mut self) {
        self.load_bundled_except(&[]);
    }

    /// Load every bundled dictionary whose name is not listed in `disabled`.
    ///
    /// Names match those in `bundled::ALL` and the section headings of
    /// `dictionaries/THIRD_PARTY_NOTICES.md`, case-insensitively. An unrecognised
    /// name is warned about rather than rejected: a config that silently does
    /// nothing is harder to diagnose than one that says so.
    pub fn load_bundled_except(&mut self, disabled: &[String]) {
        for name in disabled {
            if !bundled::ALL
                .iter()
                .any(|(known, _)| known.eq_ignore_ascii_case(name))
            {
                warn!(
                    name,
                    known = ?bundled::NAMES,
                    "Unknown bundled dictionary in dictionaries.disabled; ignoring"
                );
            }
        }

        for (name, words_str) in bundled::ALL {
            if disabled.iter().any(|d| d.eq_ignore_ascii_case(name)) {
                debug!(name, "Skipping bundled dictionary");
                continue;
            }
            parse_wordlist_into(words_str, &mut self.bundled_words);
        }
    }

    /// Load additional words from a file path. The file is expected to contain
    /// one word per line; lines starting with `#` and blank lines are skipped.
    ///
    /// Paths are resolved relative to `base` if they are not absolute.
    pub fn load_wordlist_file(&mut self, path: &Path, base: &Path) -> Result<()> {
        let resolved = if path.is_absolute() {
            path.to_path_buf()
        } else {
            base.join(path)
        };

        let resolved = resolved.canonicalize().map_err(|e| {
            anyhow::anyhow!("Cannot resolve wordlist path {}: {e}", resolved.display())
        })?;

        // Security: refuse to read files outside the workspace or common config dirs.
        // Canonicalize base too — on macOS /var → /private/var, on Windows UNC prefixes differ.
        let canonical_base = base.canonicalize().unwrap_or_else(|_| base.to_path_buf());
        if !resolved.starts_with(&canonical_base)
            && !resolved.starts_with(dirs::config_dir().unwrap_or_default())
            && !resolved.starts_with(dirs::home_dir().unwrap_or_default().join(".config"))
        {
            anyhow::bail!(
                "Wordlist path {} is outside the workspace and known config directories",
                resolved.display()
            );
        }

        let content = std::fs::read_to_string(&resolved)
            .map_err(|e| anyhow::anyhow!("Cannot read wordlist {}: {e}", resolved.display()))?;
        parse_wordlist_into(&content, &mut self.bundled_words);
        Ok(())
    }

    /// Add a word to the user dictionary and persist to disk.
    ///
    /// The word's regular inflections are derived at the same time, so adding `functor`
    /// also accepts `functors` — that is the whole point of the feature, and a user who
    /// has to add both has not been helped. Only the word itself is written to disk.
    pub fn add_word(&mut self, word: &str) -> Result<()> {
        let lower = word.to_lowercase();
        if self.user_words.insert(lower.clone()) {
            // A failed expansion must not lose the word the user just added, so this is
            // reported and stepped over rather than propagated.
            match inflection::expand([lower.as_str()]) {
                Ok(forms) => self.derived_words.extend(forms),
                Err(error) => {
                    warn!(word = %lower, %error, "Could not inflect added word; the exact form is still accepted");
                }
            }
            self.persist()?;
        }
        Ok(())
    }

    /// Generate the regular inflections of every word loaded so far.
    ///
    /// Call once, after every wordlist is in place: the expansion is a snapshot, and
    /// words added later are inflected by [`Self::add_word`] instead.
    pub fn derive_inflections(&mut self) {
        let lemmas: Vec<&str> = self
            .user_words
            .iter()
            .chain(self.bundled_words.iter())
            .map(String::as_str)
            .collect();
        match inflection::expand(lemmas) {
            Ok(forms) => {
                debug!(
                    lemmas = self.user_words.len() + self.bundled_words.len(),
                    derived = forms.len(),
                    "Generated dictionary inflections"
                );
                self.derived_words = forms;
            }
            Err(error) => warn!(%error, "Could not inflect the dictionary; exact matching only"),
        }
    }

    /// Check if a word is in the dictionary (case-insensitive).
    /// Checks both user words and bundled/external wordlists.
    ///
    /// A hyphenated compound also matches when every one of its parts is known
    /// on its own, so `Chern-Simons` resolves from `chern` and `simons`. The
    /// word handed to us is whatever span the engine reported, and the engines
    /// disagree about whether a compound is one token or two — matching both
    /// shapes here is the only place that can be correct for either.
    #[must_use]
    pub fn contains(&self, word: &str) -> bool {
        let lower = word.to_lowercase();
        if self.contains_exact(&lower) {
            return true;
        }
        self.is_known_compound(&lower)
    }

    /// Look a single already-lowercased token up in all three sets.
    fn contains_exact(&self, lower: &str) -> bool {
        self.user_words.contains(lower)
            || self.bundled_words.contains(lower)
            || self.derived_words.contains(lower)
    }

    /// Whether `lower` is a hyphenated compound whose every part is known.
    ///
    /// A present hyphen always yields at least two parts, so the only shapes to
    /// reject are the ones with an empty part — a bare `-`, a trailing `chern-`,
    /// or a doubled `--` — which would otherwise decide "every part is known"
    /// over a single token, or over nothing at all.
    fn is_known_compound(&self, lower: &str) -> bool {
        const HYPHENS: [char; 3] = ['-', '\u{2010}', '\u{2011}'];
        lower.contains(HYPHENS)
            && lower
                .split(HYPHENS)
                .all(|part| !part.is_empty() && self.contains_exact(part))
    }

    /// Return all words the user configured (user + bundled), excluding derived forms.
    pub fn words(&self) -> impl Iterator<Item = &String> {
        self.user_words.iter().chain(self.bundled_words.iter())
    }

    /// Return the number of words loaded (user + bundled), excluding derived forms.
    #[must_use]
    pub fn len(&self) -> usize {
        self.user_words.len() + self.bundled_words.len()
    }

    /// How many inflected forms were generated from those words.
    #[must_use]
    pub fn derived_len(&self) -> usize {
        self.derived_words.len()
    }

    /// Whether the dictionary is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.user_words.is_empty() && self.bundled_words.is_empty()
    }

    /// Persist only user-added words to the workspace dictionary file.
    fn persist(&self) -> Result<()> {
        let Some(path) = &self.workspace_path else {
            return Ok(());
        };

        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let mut words: Vec<&str> = self.user_words.iter().map(String::as_str).collect();
        words.sort_unstable();
        let content = words.join("\n");
        std::fs::write(path, content + "\n")?;
        Ok(())
    }
}

/// Parse a wordlist string (one word per line) into a set.
fn parse_wordlist_into(content: &str, set: &mut HashSet<String>) {
    for line in content.lines() {
        let word = line.trim();
        if !word.is_empty() && !word.starts_with('#') {
            set.insert(word.to_lowercase());
        }
    }
}

/// Bundled dictionary data embedded at compile time.
/// See `dictionaries/THIRD_PARTY_NOTICES.md` for attribution and licensing.
pub mod bundled {
    /// Software development terms, tools, acronyms, and compound words.
    /// Sources: cspell-dicts (software-terms, cpp). License: MIT.
    pub const SOFTWARE_TERMS: &str = include_str!("../dictionaries/bundled/software-terms.txt");

    /// TypeScript and JavaScript keywords, builtins, and API terms.
    /// Source: cspell-dicts (typescript). License: MIT.
    pub const TYPESCRIPT: &str = include_str!("../dictionaries/bundled/typescript.txt");

    /// Well-known company and brand names.
    /// Source: cspell-dicts (companies). License: MIT.
    pub const COMPANIES: &str = include_str!("../dictionaries/bundled/companies.txt");

    /// Computing jargon, hardware terms, and domain-specific vocabulary.
    /// Sources: hunspell-jargon (MIT), `SpellCheckDic` (MIT),
    ///          autoware-spell-check-dict (Apache-2.0).
    pub const JARGON: &str = include_str!("../dictionaries/bundled/jargon.txt");

    /// Mathematics, category theory, type theory, and mathematical physics terms.
    ///
    /// Source: nLab page titles, harvested by `scripts/build-nlab-dictionary.py`.
    /// License: none formally stated; free use granted in exchange for attribution.
    pub const MATHEMATICS: &str = include_str!("../dictionaries/bundled/mathematics.txt");

    /// All bundled wordlists, keyed by the name users write in
    /// `dictionaries.disabled` to turn one off.
    pub const ALL: &[(&str, &str)] = &[
        ("software-terms", SOFTWARE_TERMS),
        ("typescript", TYPESCRIPT),
        ("companies", COMPANIES),
        ("jargon", JARGON),
        ("mathematics", MATHEMATICS),
    ];

    /// Just the names from [`ALL`], for diagnostics.
    pub const NAMES: &[&str] = &[
        "software-terms",
        "typescript",
        "companies",
        "jargon",
        "mathematics",
    ];
}

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

    #[test]
    fn new_dictionary_is_empty() {
        let dict = Dictionary::new();
        assert!(!dict.contains("anything"));
    }

    #[test]
    fn add_and_contains() {
        let mut dict = Dictionary::new();
        dict.user_words.insert("hello".to_string());
        assert!(dict.contains("hello"));
        assert!(dict.contains("Hello")); // case-insensitive
        assert!(dict.contains("HELLO"));
    }

    #[test]
    fn persistence_roundtrip() {
        let dir = std::env::temp_dir().join("lang_check_test_dict");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        // Write
        {
            let mut dict = Dictionary::load(&dir).unwrap();
            dict.add_word("kubernetes").unwrap();
            dict.add_word("terraform").unwrap();
        }

        // Read back
        {
            let dict = Dictionary::load(&dir).unwrap();
            assert!(dict.contains("kubernetes"));
            assert!(dict.contains("Kubernetes")); // case-insensitive
            assert!(dict.contains("terraform"));
            assert!(!dict.contains("nonexistent"));
        }

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn skips_comments_and_blank_lines() {
        let dir = std::env::temp_dir().join("lang_check_test_dict_comments");
        let _ = std::fs::remove_dir_all(&dir);
        let dict_dir = dir.join(".languagecheck");
        std::fs::create_dir_all(&dict_dir).unwrap();
        std::fs::write(
            dict_dir.join("dictionary.txt"),
            "# This is a comment\n\nkubernetes\n  \n# Another comment\nterraform\n",
        )
        .unwrap();

        let dict = Dictionary::load(&dir).unwrap();
        assert!(dict.contains("kubernetes"));
        assert!(dict.contains("terraform"));
        assert_eq!(dict.words().count(), 2);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn add_duplicate_word_is_idempotent() {
        let mut dict = Dictionary::new();
        dict.user_words.insert("test".to_string());
        let initial_count = dict.words().count();
        dict.user_words.insert("test".to_string());
        assert_eq!(dict.words().count(), initial_count);
    }

    #[test]
    fn words_iterator() {
        let mut dict = Dictionary::new();
        dict.user_words.insert("alpha".to_string());
        dict.user_words.insert("beta".to_string());
        assert_eq!(dict.words().count(), 2);
    }

    #[test]
    fn bundled_dictionaries_load() {
        let mut dict = Dictionary::new();
        dict.load_bundled();

        // Should have thousands of words from bundled sources
        assert!(
            dict.len() > 5000,
            "Expected > 5000 bundled words, got {}",
            dict.len()
        );

        // Spot-check some well-known terms from each category
        assert!(
            dict.contains("kubernetes"),
            "software-terms should include kubernetes"
        );
        assert!(
            dict.contains("webpack"),
            "software-terms should include webpack"
        );
        assert!(
            dict.contains("instanceof"),
            "typescript should include instanceof"
        );
        assert!(dict.contains("stdout"), "jargon should include stdout");
    }

    #[test]
    fn mathematics_dictionary_loads() {
        let mut dict = Dictionary::new();
        dict.load_bundled();

        for term in [
            "monoidal",
            "presheaf",
            "colimit",
            "endofunctor",
            "cobordism",
        ] {
            assert!(dict.contains(term), "mathematics should include {term}");
        }
        // Harvested with diacritics intact, and matched case-insensitively.
        assert!(dict.contains("étale"), "mathematics should include étale");
        assert!(dict.contains("Grothendieck"), "lookup is case-insensitive");
    }

    #[test]
    fn disabling_a_bundled_set_drops_only_that_set() {
        let mut dict = Dictionary::new();
        dict.load_bundled_except(&["mathematics".to_string()]);

        assert!(!dict.contains("presheaf"), "mathematics should be skipped");
        assert!(dict.contains("kubernetes"), "software-terms should remain");
        assert!(dict.contains("instanceof"), "typescript should remain");
    }

    #[test]
    fn disabled_set_names_are_case_insensitive() {
        let mut dict = Dictionary::new();
        dict.load_bundled_except(&["Mathematics".to_string()]);

        assert!(!dict.contains("presheaf"));
    }

    #[test]
    fn unknown_disabled_set_name_is_tolerated() {
        // A typo'd entry must not take the other dictionaries down with it.
        let mut dict = Dictionary::new();
        dict.load_bundled_except(&["mathmatics".to_string()]);

        assert!(
            dict.contains("presheaf"),
            "nothing should have been skipped"
        );
        assert!(dict.contains("kubernetes"));
    }

    #[test]
    fn every_bundled_set_has_a_name() {
        assert_eq!(bundled::ALL.len(), bundled::NAMES.len());
        for ((name, _), listed) in bundled::ALL.iter().zip(bundled::NAMES) {
            assert_eq!(name, listed);
        }
    }

    #[test]
    fn derived_inflections_are_accepted() {
        let mut dict = Dictionary::new();
        dict.user_words.insert("functor".to_string());
        assert!(!dict.contains("functors"));
        dict.derive_inflections();
        assert!(dict.contains("functors"));
        assert!(dict.contains("Functors"), "and case-insensitively");
    }

    #[test]
    fn derived_inflections_are_never_persisted() {
        let dir = std::env::temp_dir().join("lang_check_test_derived_persist");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        let mut dict = Dictionary::load(&dir).unwrap();
        dict.add_word("functor").unwrap();
        assert!(dict.contains("functors"), "the plural is accepted");

        let written = std::fs::read_to_string(dir.join(".languagecheck/dictionary.txt")).unwrap();
        assert_eq!(
            written.trim(),
            "functor",
            "but only the typed word is recorded"
        );

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn a_bundled_word_inflects_too() {
        let mut dict = Dictionary::new();
        dict.load_bundled();
        dict.derive_inflections();
        // `mathematics.txt` carries `subalgebras` but not `subalgebra`; the reverse gap
        // is what generation closes.
        assert!(dict.contains("preorders"));
        assert!(dict.derived_len() > 1000);
    }

    #[test]
    fn hyphenated_compound_matches_when_all_parts_known() {
        let mut dict = Dictionary::new();
        dict.load_bundled();

        assert!(dict.contains("Chern-Simons"));
        assert!(dict.contains("Yang-Mills"));
        assert!(dict.contains("Seiberg-Witten"));
        // Unicode hyphen (U+2010) and non-breaking hyphen (U+2011) too.
        assert!(dict.contains("Chern\u{2010}Simons"));
        assert!(dict.contains("Chern\u{2011}Simons"));
    }

    #[test]
    fn hyphenated_compound_rejected_when_a_part_is_unknown() {
        let mut dict = Dictionary::new();
        dict.user_words.insert("chern".to_string());

        assert!(!dict.contains("chern-simmmons"));
        assert!(!dict.contains("cherm-chern"));
    }

    #[test]
    fn hyphen_split_rejects_empty_parts() {
        let mut dict = Dictionary::new();
        dict.user_words.insert("chern".to_string());

        // A trailing, leading, or doubled hyphen leaves an empty part, which must
        // not collapse into "every part is known".
        for input in ["chern-", "-chern", "chern--chern", "-", "--"] {
            assert!(!dict.contains(input), "{input} must not match");
        }
    }

    #[test]
    fn hyphen_split_only_accepts_words_the_lists_already_carry() {
        // The compound rule widens what the word lists cover; it never invents
        // vocabulary. Both halves must be present independently, so an ordinary
        // hyphenated typo stays flagged.
        let mut dict = Dictionary::new();
        dict.user_words.insert("chern".to_string());
        dict.user_words.insert("simons".to_string());

        assert!(dict.contains("chern-simons"));
        assert!(!dict.contains("well-known"));
    }

    #[test]
    fn mathematics_dictionary_excludes_nlab_misspellings() {
        let mut dict = Dictionary::new();
        dict.load_bundled();

        // These appear in nLab `[[!redirects]]` aliases, which exist precisely so
        // that misspelled links keep resolving. Harvesting them would suppress
        // real typos, so the build script reads page titles only.
        for typo in [
            "alebraic",
            "cohomlogy",
            "basises",
            "automorpism",
            "geoemtric",
        ] {
            assert!(
                !dict.contains(typo),
                "{typo} must not be an accepted spelling"
            );
        }
    }

    #[test]
    fn bundled_plus_user_words() {
        let mut dict = Dictionary::new();
        dict.load_bundled();
        let bundled_count = dict.len();

        dict.user_words.insert("myprojectword".to_string());
        assert_eq!(dict.len(), bundled_count + 1);
        assert!(dict.contains("myprojectword"));
        // Bundled words still present
        assert!(dict.contains("kubernetes"));
    }

    #[test]
    fn load_wordlist_file_works() {
        let dir = std::env::temp_dir().join("lang_check_test_wordlist");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        let wordlist = dir.join("custom.txt");
        std::fs::write(&wordlist, "# My custom words\nfoobar\nbazqux\n").unwrap();

        let mut dict = Dictionary::new();
        dict.load_wordlist_file(&wordlist, &dir).unwrap();

        assert!(dict.contains("foobar"));
        assert!(dict.contains("bazqux"));
        assert_eq!(dict.len(), 2);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn persistence_excludes_bundled_words() {
        let dir = std::env::temp_dir().join("lang_check_test_dict_bundled_persist");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        // Add a user word alongside bundled words
        {
            let mut dict = Dictionary::load(&dir).unwrap();
            dict.load_bundled();
            dict.add_word("myuserword").unwrap();
        }

        // Read the persisted file directly — should only contain user words
        let dict_path = dir.join(".languagecheck").join("dictionary.txt");
        let content = std::fs::read_to_string(&dict_path).unwrap();
        assert!(
            content.contains("myuserword"),
            "User word should be persisted"
        );
        assert!(
            !content.contains("kubernetes"),
            "Bundled words should NOT be persisted"
        );

        // Reload and verify everything still works
        {
            let mut dict = Dictionary::load(&dir).unwrap();
            dict.load_bundled();
            assert!(dict.contains("myuserword"));
            assert!(dict.contains("kubernetes"));
        }

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn load_wordlist_file_relative_path() {
        let dir = std::env::temp_dir().join("lang_check_test_wordlist_rel");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        std::fs::write(dir.join("terms.txt"), "myterm\n").unwrap();

        let mut dict = Dictionary::new();
        dict.load_wordlist_file(Path::new("terms.txt"), &dir)
            .unwrap();

        assert!(dict.contains("myterm"));

        let _ = std::fs::remove_dir_all(&dir);
    }
}