nltk_porter/
lib.rs

1//! NLTK's Porter stemmer, ported to Rust.
2//!
3//! A faithful port of `nltk.stem.porter.PorterStemmer` (Porter 1980, plus the
4//! extensions NLTK ships). The word is processed by Unicode scalar value (`char`),
5//! matching Python's `str` indexing.
6//!
7//! Three modes, selectable like NLTK:
8//! - [`Mode::Nltk`] (default): NLTK contributors' improvements.
9//! - [`Mode::Martin`]: only the extensions on Martin Porter's website.
10//! - [`Mode::Original`]: faithful to the 1980 paper (Porter deprecates this).
11//!
12//! Ported from NLTK (Apache-2.0).
13//!
14//! ```
15//! use nltk_porter::{PorterStemmer, Mode};
16//! let p = PorterStemmer::new(Mode::Nltk);
17//! assert_eq!(p.stem("caresses"), "caress");
18//! assert_eq!(p.stem("happy"), "happi");
19//! ```
20
21use std::collections::HashMap;
22
23/// Which variant of the algorithm to run.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum Mode {
26    /// NLTK contributors' extensions (NLTK's default).
27    Nltk,
28    /// Martin Porter's website extensions.
29    Martin,
30    /// The original 1980 paper.
31    Original,
32}
33
34/// Condition attached to a suffix-removal rule. Evaluated against the candidate
35/// `stem` (and sometimes the whole `word`), with access to the stemmer for `measure`.
36enum Cond {
37    /// Unconditional.
38    Always,
39    /// `measure(stem) > 0`.
40    Pos,
41    /// `measure(stem) > 1`.
42    Gt1,
43    /// `measure(word[..len-n]) > 0` (step 2 `logi`, NLTK).
44    PosOfWord(usize),
45    /// `measure(word[..len-n]) > 1` (step 5b `ll`).
46    Gt1OfWord(usize),
47    /// `measure(stem) > 1 && stem ends in s or t` (step 4 `ion`).
48    IonM1,
49    /// `measure(stem) == 1 && ends_cvc(stem)` (step 1b final rule).
50    M1AndCvc,
51    /// Step 1c `y`: NLTK vs original variants.
52    Y,
53    /// A captured char is not `l`, `s`, or `z` (step 1b `*d`).
54    NotLsz(char),
55}
56
57struct Rule {
58    suffix: &'static str,
59    repl: Vec<char>,
60    cond: Cond,
61}
62
63fn rule(suffix: &'static str, repl: &str, cond: Cond) -> Rule {
64    Rule {
65        suffix,
66        repl: repl.chars().collect(),
67        cond,
68    }
69}
70
71fn is_vowel(c: char) -> bool {
72    matches!(c, 'a' | 'e' | 'i' | 'o' | 'u')
73}
74
75/// Whether `word[i]` is a consonant (a `y` is a consonant iff preceded by a vowel).
76fn is_consonant(word: &[char], i: usize) -> bool {
77    let c = word[i];
78    if is_vowel(c) {
79        return false;
80    }
81    if c == 'y' {
82        let mut negate = false;
83        let mut j = i;
84        while j > 0 && word[j] == 'y' {
85            negate = !negate;
86            j -= 1;
87        }
88        return is_vowel(word[j]) == negate;
89    }
90    true
91}
92
93fn contains_vowel(stem: &[char]) -> bool {
94    (0..stem.len()).any(|i| !is_consonant(stem, i))
95}
96
97fn ends_double_consonant(word: &[char]) -> bool {
98    word.len() >= 2 && word[word.len() - 1] == word[word.len() - 2] && is_consonant(word, word.len() - 1)
99}
100
101fn ends_with(word: &[char], suffix: &str) -> bool {
102    let s: Vec<char> = suffix.chars().collect();
103    word.len() >= s.len() && word[word.len() - s.len()..] == s[..]
104}
105
106fn concat(a: &[char], b: &[char]) -> Vec<char> {
107    let mut v = Vec::with_capacity(a.len() + b.len());
108    v.extend_from_slice(a);
109    v.extend_from_slice(b);
110    v
111}
112
113/// A Porter stemmer in a fixed [`Mode`].
114pub struct PorterStemmer {
115    mode: Mode,
116    pool: HashMap<String, String>,
117}
118
119impl PorterStemmer {
120    /// Construct a stemmer in the given mode.
121    pub fn new(mode: Mode) -> Self {
122        let mut pool = HashMap::new();
123        if mode == Mode::Nltk {
124            // Table of irregular forms (NLTK extension): val -> key.
125            let irregular: &[(&str, &[&str])] = &[
126                ("sky", &["sky", "skies"]),
127                ("die", &["dying"]),
128                ("lie", &["lying"]),
129                ("tie", &["tying"]),
130                ("news", &["news"]),
131                ("inning", &["innings", "inning"]),
132                ("outing", &["outings", "outing"]),
133                ("canning", &["cannings", "canning"]),
134                ("howe", &["howe"]),
135                ("proceed", &["proceed"]),
136                ("exceed", &["exceed"]),
137                ("succeed", &["succeed"]),
138            ];
139            for (key, vals) in irregular {
140                for v in *vals {
141                    pool.insert((*v).to_string(), (*key).to_string());
142                }
143            }
144        }
145        PorterStemmer { mode, pool }
146    }
147
148    fn measure(&self, stem: &[char]) -> usize {
149        let mut cv = String::with_capacity(stem.len());
150        for i in 0..stem.len() {
151            cv.push(if is_consonant(stem, i) { 'c' } else { 'v' });
152        }
153        cv.matches("vc").count()
154    }
155
156    fn has_positive_measure(&self, stem: &[char]) -> bool {
157        self.measure(stem) > 0
158    }
159
160    /// Condition `*o`: stem ends `cvc` where the last `c` is not `w/x/y`.
161    fn ends_cvc(&self, word: &[char]) -> bool {
162        let n = word.len();
163        (n >= 3
164            && is_consonant(word, n - 3)
165            && !is_consonant(word, n - 2)
166            && is_consonant(word, n - 1)
167            && !matches!(word[n - 1], 'w' | 'x' | 'y'))
168            || (self.mode == Mode::Nltk
169                && n == 2
170                && !is_consonant(word, 0)
171                && is_consonant(word, 1))
172    }
173
174    fn eval(&self, cond: &Cond, word: &[char], stem: &[char]) -> bool {
175        match cond {
176            Cond::Always => true,
177            Cond::Pos => self.measure(stem) > 0,
178            Cond::Gt1 => self.measure(stem) > 1,
179            Cond::PosOfWord(n) => self.measure(&word[..word.len() - n]) > 0,
180            Cond::Gt1OfWord(n) => self.measure(&word[..word.len() - n]) > 1,
181            Cond::IonM1 => self.measure(stem) > 1 && matches!(stem.last(), Some('s') | Some('t')),
182            Cond::M1AndCvc => self.measure(stem) == 1 && self.ends_cvc(stem),
183            Cond::Y => {
184                if self.mode == Mode::Nltk {
185                    stem.len() > 1 && is_consonant(stem, stem.len() - 1)
186                } else {
187                    contains_vowel(stem)
188                }
189            }
190            Cond::NotLsz(c) => !matches!(c, 'l' | 's' | 'z'),
191        }
192    }
193
194    /// Apply the first applicable rule (NLTK's `_apply_rule_list`).
195    fn apply_rules(&self, word: &[char], rules: &[Rule]) -> Vec<char> {
196        for r in rules {
197            if r.suffix == "*d" {
198                if ends_double_consonant(word) {
199                    let stem = &word[..word.len() - 2];
200                    return if self.eval(&r.cond, word, stem) {
201                        concat(stem, &r.repl)
202                    } else {
203                        word.to_vec()
204                    };
205                }
206                continue;
207            }
208            if ends_with(word, r.suffix) {
209                let stem = &word[..word.len() - r.suffix.chars().count()];
210                return if self.eval(&r.cond, word, stem) {
211                    concat(stem, &r.repl)
212                } else {
213                    word.to_vec()
214                };
215            }
216        }
217        word.to_vec()
218    }
219
220    fn step1a(&self, word: &[char]) -> Vec<char> {
221        if self.mode == Mode::Nltk && ends_with(word, "ies") && word.len() == 4 {
222            return concat(&word[..word.len() - 3], &['i', 'e']);
223        }
224        self.apply_rules(
225            word,
226            &[
227                rule("sses", "ss", Cond::Always),
228                rule("ies", "i", Cond::Always),
229                rule("ss", "ss", Cond::Always),
230                rule("s", "", Cond::Always),
231            ],
232        )
233    }
234
235    fn step1b(&self, word: &[char]) -> Vec<char> {
236        if self.mode == Mode::Nltk && ends_with(word, "ied") {
237            return if word.len() == 4 {
238                concat(&word[..word.len() - 3], &['i', 'e'])
239            } else {
240                concat(&word[..word.len() - 3], &['i'])
241            };
242        }
243
244        if ends_with(word, "eed") {
245            let stem = &word[..word.len() - 3];
246            return if self.measure(stem) > 0 {
247                concat(stem, &['e', 'e'])
248            } else {
249                word.to_vec()
250            };
251        }
252
253        let mut intermediate: Option<Vec<char>> = None;
254        for suffix in ["ed", "ing"] {
255            if ends_with(word, suffix) {
256                let stem = word[..word.len() - suffix.len()].to_vec();
257                if contains_vowel(&stem) {
258                    intermediate = Some(stem);
259                    break;
260                }
261            }
262        }
263        let inter = match intermediate {
264            Some(s) => s,
265            None => return word.to_vec(),
266        };
267
268        let last = *inter.last().unwrap();
269        let rules = [
270            rule("at", "ate", Cond::Always),
271            rule("bl", "ble", Cond::Always),
272            rule("iz", "ize", Cond::Always),
273            Rule {
274                suffix: "*d",
275                repl: vec![last],
276                cond: Cond::NotLsz(last),
277            },
278            rule("", "e", Cond::M1AndCvc),
279        ];
280        self.apply_rules(&inter, &rules)
281    }
282
283    fn step1c(&self, word: &[char]) -> Vec<char> {
284        self.apply_rules(word, &[rule("y", "i", Cond::Y)])
285    }
286
287    fn step2(&self, word: &[char]) -> Vec<char> {
288        if self.mode == Mode::Nltk
289            && ends_with(word, "alli")
290            && self.has_positive_measure(&word[..word.len() - 4])
291        {
292            let reduced = concat(&word[..word.len() - 4], &['a', 'l']);
293            return self.step2(&reduced);
294        }
295
296        let mut rules = vec![
297            rule("ational", "ate", Cond::Pos),
298            rule("tional", "tion", Cond::Pos),
299            rule("enci", "ence", Cond::Pos),
300            rule("anci", "ance", Cond::Pos),
301            rule("izer", "ize", Cond::Pos),
302            if self.mode == Mode::Original {
303                rule("abli", "able", Cond::Pos)
304            } else {
305                rule("bli", "ble", Cond::Pos)
306            },
307            rule("alli", "al", Cond::Pos),
308            rule("entli", "ent", Cond::Pos),
309            rule("eli", "e", Cond::Pos),
310            rule("ousli", "ous", Cond::Pos),
311            rule("ization", "ize", Cond::Pos),
312            rule("ation", "ate", Cond::Pos),
313            rule("ator", "ate", Cond::Pos),
314            rule("alism", "al", Cond::Pos),
315            rule("iveness", "ive", Cond::Pos),
316            rule("fulness", "ful", Cond::Pos),
317            rule("ousness", "ous", Cond::Pos),
318            rule("aliti", "al", Cond::Pos),
319            rule("iviti", "ive", Cond::Pos),
320            rule("biliti", "ble", Cond::Pos),
321        ];
322        if self.mode == Mode::Nltk {
323            rules.push(rule("fulli", "ful", Cond::Pos));
324            rules.push(rule("logi", "log", Cond::PosOfWord(3)));
325        }
326        if self.mode == Mode::Martin {
327            rules.push(rule("logi", "log", Cond::Pos));
328        }
329        self.apply_rules(word, &rules)
330    }
331
332    fn step3(&self, word: &[char]) -> Vec<char> {
333        self.apply_rules(
334            word,
335            &[
336                rule("icate", "ic", Cond::Pos),
337                rule("ative", "", Cond::Pos),
338                rule("alize", "al", Cond::Pos),
339                rule("iciti", "ic", Cond::Pos),
340                rule("ical", "ic", Cond::Pos),
341                rule("ful", "", Cond::Pos),
342                rule("ness", "", Cond::Pos),
343            ],
344        )
345    }
346
347    fn step4(&self, word: &[char]) -> Vec<char> {
348        self.apply_rules(
349            word,
350            &[
351                rule("al", "", Cond::Gt1),
352                rule("ance", "", Cond::Gt1),
353                rule("ence", "", Cond::Gt1),
354                rule("er", "", Cond::Gt1),
355                rule("ic", "", Cond::Gt1),
356                rule("able", "", Cond::Gt1),
357                rule("ible", "", Cond::Gt1),
358                rule("ant", "", Cond::Gt1),
359                rule("ement", "", Cond::Gt1),
360                rule("ment", "", Cond::Gt1),
361                rule("ent", "", Cond::Gt1),
362                rule("ion", "", Cond::IonM1),
363                rule("ou", "", Cond::Gt1),
364                rule("ism", "", Cond::Gt1),
365                rule("ate", "", Cond::Gt1),
366                rule("iti", "", Cond::Gt1),
367                rule("ous", "", Cond::Gt1),
368                rule("ive", "", Cond::Gt1),
369                rule("ize", "", Cond::Gt1),
370            ],
371        )
372    }
373
374    fn step5a(&self, word: &[char]) -> Vec<char> {
375        if ends_with(word, "e") {
376            let stem = &word[..word.len() - 1];
377            let m = self.measure(stem);
378            if m > 1 || (m == 1 && !self.ends_cvc(stem)) {
379                return stem.to_vec();
380            }
381        }
382        word.to_vec()
383    }
384
385    fn step5b(&self, word: &[char]) -> Vec<char> {
386        self.apply_rules(word, &[rule("ll", "l", Cond::Gt1OfWord(1))])
387    }
388
389    /// Stem `word` (lowercasing it first, as NLTK does by default).
390    pub fn stem(&self, word: &str) -> String {
391        let lowered = word.to_lowercase();
392
393        if self.mode == Mode::Nltk {
394            if let Some(s) = self.pool.get(&lowered) {
395                return s.clone();
396            }
397        }
398
399        // NLTK gates on the ORIGINAL word's length (`len(word)`), not the lowered one.
400        if self.mode != Mode::Original && word.chars().count() <= 2 {
401            return lowered;
402        }
403
404        let mut chars: Vec<char> = lowered.chars().collect();
405        chars = self.step1a(&chars);
406        chars = self.step1b(&chars);
407        chars = self.step1c(&chars);
408        chars = self.step2(&chars);
409        chars = self.step3(&chars);
410        chars = self.step4(&chars);
411        chars = self.step5a(&chars);
412        chars = self.step5b(&chars);
413
414        chars.into_iter().collect()
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use super::*;
421
422    #[test]
423    fn paper_examples() {
424        let p = PorterStemmer::new(Mode::Nltk);
425        assert_eq!(p.stem("caresses"), "caress");
426        assert_eq!(p.stem("ponies"), "poni");
427        assert_eq!(p.stem("cats"), "cat");
428        assert_eq!(p.stem("agreed"), "agre");
429        assert_eq!(p.stem("motoring"), "motor");
430        assert_eq!(p.stem("happy"), "happi");
431        assert_eq!(p.stem("relational"), "relat");
432        assert_eq!(p.stem("revival"), "reviv");
433        assert_eq!(p.stem("controll"), "control");
434    }
435
436    #[test]
437    fn nltk_pool_and_short_words() {
438        let p = PorterStemmer::new(Mode::Nltk);
439        assert_eq!(p.stem("skies"), "sky"); // irregular pool
440        assert_eq!(p.stem("dying"), "die");
441        assert_eq!(p.stem("by"), "by"); // length <= 2
442    }
443
444    #[test]
445    fn modes_differ() {
446        // 'happy' -> 'happi' in all; but y-handling differs for e.g. 'enjoy'.
447        assert_eq!(PorterStemmer::new(Mode::Nltk).stem("enjoy"), "enjoy");
448        assert_eq!(PorterStemmer::new(Mode::Original).stem("enjoy"), "enjoi");
449    }
450}
nltk_porter/lib.rs

nltk_porter/
lib.rs
