nltk-porter 0.1.0

//! NLTK's Porter stemmer, ported to Rust.
//!
//! A faithful port of `nltk.stem.porter.PorterStemmer` (Porter 1980, plus the
//! extensions NLTK ships). The word is processed by Unicode scalar value (`char`),
//! matching Python's `str` indexing.
//!
//! Three modes, selectable like NLTK:
//! - [`Mode::Nltk`] (default): NLTK contributors' improvements.
//! - [`Mode::Martin`]: only the extensions on Martin Porter's website.
//! - [`Mode::Original`]: faithful to the 1980 paper (Porter deprecates this).
//!
//! Ported from NLTK (Apache-2.0).
//!
//! ```
//! use nltk_porter::{PorterStemmer, Mode};
//! let p = PorterStemmer::new(Mode::Nltk);
//! assert_eq!(p.stem("caresses"), "caress");
//! assert_eq!(p.stem("happy"), "happi");
//! ```

use std::collections::HashMap;

/// Which variant of the algorithm to run.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Mode {
    /// NLTK contributors' extensions (NLTK's default).
    Nltk,
    /// Martin Porter's website extensions.
    Martin,
    /// The original 1980 paper.
    Original,
}

/// Condition attached to a suffix-removal rule. Evaluated against the candidate
/// `stem` (and sometimes the whole `word`), with access to the stemmer for `measure`.
enum Cond {
    /// Unconditional.
    Always,
    /// `measure(stem) > 0`.
    Pos,
    /// `measure(stem) > 1`.
    Gt1,
    /// `measure(word[..len-n]) > 0` (step 2 `logi`, NLTK).
    PosOfWord(usize),
    /// `measure(word[..len-n]) > 1` (step 5b `ll`).
    Gt1OfWord(usize),
    /// `measure(stem) > 1 && stem ends in s or t` (step 4 `ion`).
    IonM1,
    /// `measure(stem) == 1 && ends_cvc(stem)` (step 1b final rule).
    M1AndCvc,
    /// Step 1c `y`: NLTK vs original variants.
    Y,
    /// A captured char is not `l`, `s`, or `z` (step 1b `*d`).
    NotLsz(char),
}

struct Rule {
    suffix: &'static str,
    repl: Vec<char>,
    cond: Cond,
}

fn rule(suffix: &'static str, repl: &str, cond: Cond) -> Rule {
    Rule {
        suffix,
        repl: repl.chars().collect(),
        cond,
    }
}

fn is_vowel(c: char) -> bool {
    matches!(c, 'a' | 'e' | 'i' | 'o' | 'u')
}

/// Whether `word[i]` is a consonant (a `y` is a consonant iff preceded by a vowel).
fn is_consonant(word: &[char], i: usize) -> bool {
    let c = word[i];
    if is_vowel(c) {
        return false;
    }
    if c == 'y' {
        let mut negate = false;
        let mut j = i;
        while j > 0 && word[j] == 'y' {
            negate = !negate;
            j -= 1;
        }
        return is_vowel(word[j]) == negate;
    }
    true
}

fn contains_vowel(stem: &[char]) -> bool {
    (0..stem.len()).any(|i| !is_consonant(stem, i))
}

fn ends_double_consonant(word: &[char]) -> bool {
    word.len() >= 2 && word[word.len() - 1] == word[word.len() - 2] && is_consonant(word, word.len() - 1)
}

fn ends_with(word: &[char], suffix: &str) -> bool {
    let s: Vec<char> = suffix.chars().collect();
    word.len() >= s.len() && word[word.len() - s.len()..] == s[..]
}

fn concat(a: &[char], b: &[char]) -> Vec<char> {
    let mut v = Vec::with_capacity(a.len() + b.len());
    v.extend_from_slice(a);
    v.extend_from_slice(b);
    v
}

/// A Porter stemmer in a fixed [`Mode`].
pub struct PorterStemmer {
    mode: Mode,
    pool: HashMap<String, String>,
}

impl PorterStemmer {
    /// Construct a stemmer in the given mode.
    pub fn new(mode: Mode) -> Self {
        let mut pool = HashMap::new();
        if mode == Mode::Nltk {
            // Table of irregular forms (NLTK extension): val -> key.
            let irregular: &[(&str, &[&str])] = &[
                ("sky", &["sky", "skies"]),
                ("die", &["dying"]),
                ("lie", &["lying"]),
                ("tie", &["tying"]),
                ("news", &["news"]),
                ("inning", &["innings", "inning"]),
                ("outing", &["outings", "outing"]),
                ("canning", &["cannings", "canning"]),
                ("howe", &["howe"]),
                ("proceed", &["proceed"]),
                ("exceed", &["exceed"]),
                ("succeed", &["succeed"]),
            ];
            for (key, vals) in irregular {
                for v in *vals {
                    pool.insert((*v).to_string(), (*key).to_string());
                }
            }
        }
        PorterStemmer { mode, pool }
    }

    fn measure(&self, stem: &[char]) -> usize {
        let mut cv = String::with_capacity(stem.len());
        for i in 0..stem.len() {
            cv.push(if is_consonant(stem, i) { 'c' } else { 'v' });
        }
        cv.matches("vc").count()
    }

    fn has_positive_measure(&self, stem: &[char]) -> bool {
        self.measure(stem) > 0
    }

    /// Condition `*o`: stem ends `cvc` where the last `c` is not `w/x/y`.
    fn ends_cvc(&self, word: &[char]) -> bool {
        let n = word.len();
        (n >= 3
            && is_consonant(word, n - 3)
            && !is_consonant(word, n - 2)
            && is_consonant(word, n - 1)
            && !matches!(word[n - 1], 'w' | 'x' | 'y'))
            || (self.mode == Mode::Nltk
                && n == 2
                && !is_consonant(word, 0)
                && is_consonant(word, 1))
    }

    fn eval(&self, cond: &Cond, word: &[char], stem: &[char]) -> bool {
        match cond {
            Cond::Always => true,
            Cond::Pos => self.measure(stem) > 0,
            Cond::Gt1 => self.measure(stem) > 1,
            Cond::PosOfWord(n) => self.measure(&word[..word.len() - n]) > 0,
            Cond::Gt1OfWord(n) => self.measure(&word[..word.len() - n]) > 1,
            Cond::IonM1 => self.measure(stem) > 1 && matches!(stem.last(), Some('s') | Some('t')),
            Cond::M1AndCvc => self.measure(stem) == 1 && self.ends_cvc(stem),
            Cond::Y => {
                if self.mode == Mode::Nltk {
                    stem.len() > 1 && is_consonant(stem, stem.len() - 1)
                } else {
                    contains_vowel(stem)
                }
            }
            Cond::NotLsz(c) => !matches!(c, 'l' | 's' | 'z'),
        }
    }

    /// Apply the first applicable rule (NLTK's `_apply_rule_list`).
    fn apply_rules(&self, word: &[char], rules: &[Rule]) -> Vec<char> {
        for r in rules {
            if r.suffix == "*d" {
                if ends_double_consonant(word) {
                    let stem = &word[..word.len() - 2];
                    return if self.eval(&r.cond, word, stem) {
                        concat(stem, &r.repl)
                    } else {
                        word.to_vec()
                    };
                }
                continue;
            }
            if ends_with(word, r.suffix) {
                let stem = &word[..word.len() - r.suffix.chars().count()];
                return if self.eval(&r.cond, word, stem) {
                    concat(stem, &r.repl)
                } else {
                    word.to_vec()
                };
            }
        }
        word.to_vec()
    }

    fn step1a(&self, word: &[char]) -> Vec<char> {
        if self.mode == Mode::Nltk && ends_with(word, "ies") && word.len() == 4 {
            return concat(&word[..word.len() - 3], &['i', 'e']);
        }
        self.apply_rules(
            word,
            &[
                rule("sses", "ss", Cond::Always),
                rule("ies", "i", Cond::Always),
                rule("ss", "ss", Cond::Always),
                rule("s", "", Cond::Always),
            ],
        )
    }

    fn step1b(&self, word: &[char]) -> Vec<char> {
        if self.mode == Mode::Nltk && ends_with(word, "ied") {
            return if word.len() == 4 {
                concat(&word[..word.len() - 3], &['i', 'e'])
            } else {
                concat(&word[..word.len() - 3], &['i'])
            };
        }

        if ends_with(word, "eed") {
            let stem = &word[..word.len() - 3];
            return if self.measure(stem) > 0 {
                concat(stem, &['e', 'e'])
            } else {
                word.to_vec()
            };
        }

        let mut intermediate: Option<Vec<char>> = None;
        for suffix in ["ed", "ing"] {
            if ends_with(word, suffix) {
                let stem = word[..word.len() - suffix.len()].to_vec();
                if contains_vowel(&stem) {
                    intermediate = Some(stem);
                    break;
                }
            }
        }
        let inter = match intermediate {
            Some(s) => s,
            None => return word.to_vec(),
        };

        let last = *inter.last().unwrap();
        let rules = [
            rule("at", "ate", Cond::Always),
            rule("bl", "ble", Cond::Always),
            rule("iz", "ize", Cond::Always),
            Rule {
                suffix: "*d",
                repl: vec![last],
                cond: Cond::NotLsz(last),
            },
            rule("", "e", Cond::M1AndCvc),
        ];
        self.apply_rules(&inter, &rules)
    }

    fn step1c(&self, word: &[char]) -> Vec<char> {
        self.apply_rules(word, &[rule("y", "i", Cond::Y)])
    }

    fn step2(&self, word: &[char]) -> Vec<char> {
        if self.mode == Mode::Nltk
            && ends_with(word, "alli")
            && self.has_positive_measure(&word[..word.len() - 4])
        {
            let reduced = concat(&word[..word.len() - 4], &['a', 'l']);
            return self.step2(&reduced);
        }

        let mut rules = vec![
            rule("ational", "ate", Cond::Pos),
            rule("tional", "tion", Cond::Pos),
            rule("enci", "ence", Cond::Pos),
            rule("anci", "ance", Cond::Pos),
            rule("izer", "ize", Cond::Pos),
            if self.mode == Mode::Original {
                rule("abli", "able", Cond::Pos)
            } else {
                rule("bli", "ble", Cond::Pos)
            },
            rule("alli", "al", Cond::Pos),
            rule("entli", "ent", Cond::Pos),
            rule("eli", "e", Cond::Pos),
            rule("ousli", "ous", Cond::Pos),
            rule("ization", "ize", Cond::Pos),
            rule("ation", "ate", Cond::Pos),
            rule("ator", "ate", Cond::Pos),
            rule("alism", "al", Cond::Pos),
            rule("iveness", "ive", Cond::Pos),
            rule("fulness", "ful", Cond::Pos),
            rule("ousness", "ous", Cond::Pos),
            rule("aliti", "al", Cond::Pos),
            rule("iviti", "ive", Cond::Pos),
            rule("biliti", "ble", Cond::Pos),
        ];
        if self.mode == Mode::Nltk {
            rules.push(rule("fulli", "ful", Cond::Pos));
            rules.push(rule("logi", "log", Cond::PosOfWord(3)));
        }
        if self.mode == Mode::Martin {
            rules.push(rule("logi", "log", Cond::Pos));
        }
        self.apply_rules(word, &rules)
    }

    fn step3(&self, word: &[char]) -> Vec<char> {
        self.apply_rules(
            word,
            &[
                rule("icate", "ic", Cond::Pos),
                rule("ative", "", Cond::Pos),
                rule("alize", "al", Cond::Pos),
                rule("iciti", "ic", Cond::Pos),
                rule("ical", "ic", Cond::Pos),
                rule("ful", "", Cond::Pos),
                rule("ness", "", Cond::Pos),
            ],
        )
    }

    fn step4(&self, word: &[char]) -> Vec<char> {
        self.apply_rules(
            word,
            &[
                rule("al", "", Cond::Gt1),
                rule("ance", "", Cond::Gt1),
                rule("ence", "", Cond::Gt1),
                rule("er", "", Cond::Gt1),
                rule("ic", "", Cond::Gt1),
                rule("able", "", Cond::Gt1),
                rule("ible", "", Cond::Gt1),
                rule("ant", "", Cond::Gt1),
                rule("ement", "", Cond::Gt1),
                rule("ment", "", Cond::Gt1),
                rule("ent", "", Cond::Gt1),
                rule("ion", "", Cond::IonM1),
                rule("ou", "", Cond::Gt1),
                rule("ism", "", Cond::Gt1),
                rule("ate", "", Cond::Gt1),
                rule("iti", "", Cond::Gt1),
                rule("ous", "", Cond::Gt1),
                rule("ive", "", Cond::Gt1),
                rule("ize", "", Cond::Gt1),
            ],
        )
    }

    fn step5a(&self, word: &[char]) -> Vec<char> {
        if ends_with(word, "e") {
            let stem = &word[..word.len() - 1];
            let m = self.measure(stem);
            if m > 1 || (m == 1 && !self.ends_cvc(stem)) {
                return stem.to_vec();
            }
        }
        word.to_vec()
    }

    fn step5b(&self, word: &[char]) -> Vec<char> {
        self.apply_rules(word, &[rule("ll", "l", Cond::Gt1OfWord(1))])
    }

    /// Stem `word` (lowercasing it first, as NLTK does by default).
    pub fn stem(&self, word: &str) -> String {
        let lowered = word.to_lowercase();

        if self.mode == Mode::Nltk {
            if let Some(s) = self.pool.get(&lowered) {
                return s.clone();
            }
        }

        // NLTK gates on the ORIGINAL word's length (`len(word)`), not the lowered one.
        if self.mode != Mode::Original && word.chars().count() <= 2 {
            return lowered;
        }

        let mut chars: Vec<char> = lowered.chars().collect();
        chars = self.step1a(&chars);
        chars = self.step1b(&chars);
        chars = self.step1c(&chars);
        chars = self.step2(&chars);
        chars = self.step3(&chars);
        chars = self.step4(&chars);
        chars = self.step5a(&chars);
        chars = self.step5b(&chars);

        chars.into_iter().collect()
    }
}

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

    #[test]
    fn paper_examples() {
        let p = PorterStemmer::new(Mode::Nltk);
        assert_eq!(p.stem("caresses"), "caress");
        assert_eq!(p.stem("ponies"), "poni");
        assert_eq!(p.stem("cats"), "cat");
        assert_eq!(p.stem("agreed"), "agre");
        assert_eq!(p.stem("motoring"), "motor");
        assert_eq!(p.stem("happy"), "happi");
        assert_eq!(p.stem("relational"), "relat");
        assert_eq!(p.stem("revival"), "reviv");
        assert_eq!(p.stem("controll"), "control");
    }

    #[test]
    fn nltk_pool_and_short_words() {
        let p = PorterStemmer::new(Mode::Nltk);
        assert_eq!(p.stem("skies"), "sky"); // irregular pool
        assert_eq!(p.stem("dying"), "die");
        assert_eq!(p.stem("by"), "by"); // length <= 2
    }

    #[test]
    fn modes_differ() {
        // 'happy' -> 'happi' in all; but y-handling differs for e.g. 'enjoy'.
        assert_eq!(PorterStemmer::new(Mode::Nltk).stem("enjoy"), "enjoy");
        assert_eq!(PorterStemmer::new(Mode::Original).stem("enjoy"), "enjoi");
    }
}