Skip to main content

uqa_analysis/
tokenizer.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Tokenizers for the analysis pipeline. An [`Analyzer`] owns exactly one
8//! tokenizer.
9//!
10//! [`Analyzer`]: crate::analyzer::Analyzer
11
12use serde::{Deserialize, Serialize};
13use std::sync::OnceLock;
14use uqa_core::memory::{Budgeted, MemoryBudget};
15
16use crate::character_class::class;
17use crate::error::{AnalysisError, AnalysisResult};
18
19mod compiled;
20mod stream;
21pub(crate) use compiled::PreparedTokenizer;
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24#[serde(tag = "type", rename_all = "snake_case")]
25pub enum Tokenizer {
26    Whitespace,
27    Standard,
28    Letter,
29    NGram {
30        min_gram: usize,
31        max_gram: usize,
32    },
33    Pattern {
34        pattern: String,
35    },
36    Keyword,
37    #[cfg(feature = "nori")]
38    #[serde(rename = "nori_tokenizer")]
39    Nori(crate::nori::NoriTokenizerConfig),
40    #[cfg(feature = "kuromoji")]
41    #[serde(rename = "kuromoji_tokenizer")]
42    Kuromoji(crate::kuromoji::KuromojiTokenizerConfig),
43}
44
45impl Tokenizer {
46    /// Validate configuration without tokenizing input. This is used when an
47    /// analyzer is registered, while [`Self::tokenize`] repeats the checks so
48    /// deserialized legacy values can never bypass them.
49    pub fn validate(&self) -> AnalysisResult<()> {
50        match self {
51            #[cfg(feature = "nori")]
52            Tokenizer::Nori(_) => self.prepare().map(|_| ()),
53            #[cfg(feature = "kuromoji")]
54            Tokenizer::Kuromoji(_) => self.prepare().map(|_| ()),
55            Tokenizer::NGram { .. } | Tokenizer::Pattern { .. } => self.prepare().map(|_| ()),
56            _ => Ok(()),
57        }
58    }
59
60    pub fn tokenize(&self, text: &str) -> AnalysisResult<Vec<String>> {
61        self.tokenize_with_offsets(text)?.into_terms()
62    }
63
64    /// Tokenize source text with explicit offsets, positions, and final source coordinates.
65    pub fn tokenize_with_offsets(&self, text: &str) -> AnalysisResult<crate::AnalyzedText> {
66        self.tokenize_mapped(&crate::FilteredText::new(text))
67    }
68
69    /// Tokenize with reservations for owned buffers and retained source provenance.
70    ///
71    /// Immutable configuration resources remain separately managed. Built-in word tokenizers and every prepared pattern search poll while scanning; regex fallback workspace reserves its buffers from the caller's allowance. Emission and projection also poll. Cloning the underlying analyzed value creates separate, unreserved token buffers.
72    pub fn tokenize_with_offsets_budgeted(
73        &self,
74        text: &str,
75        budget: &MemoryBudget,
76        poll: impl FnMut() -> AnalysisResult<()>,
77    ) -> AnalysisResult<Budgeted<crate::AnalyzedText>> {
78        self.tokenize_mapped_budgeted(&crate::FilteredText::new(text), budget, poll)
79    }
80
81    /// Tokenize a retained character-filter result under a caller allowance.
82    ///
83    /// Previously prepared source allocations keep their existing owner. Newly built coordinate indexes, token buffers, terms, and retained source copies use `budget`; failed calls leave the input's coordinate caches unchanged.
84    ///
85    /// ```
86    /// use uqa_analysis::{CharFilter, Tokenizer};
87    /// use uqa_core::memory::MemoryBudget;
88    /// let budget = MemoryBudget::new(16 * 1024);
89    /// let filtered = CharFilter::HTMLStrip.filter_with_offsets_budgeted(
90    ///     "<b>韓🙂</b>", &budget, &mut || Ok(()),
91    /// )?;
92    /// let tokens = Tokenizer::Whitespace.tokenize_mapped_budgeted(
93    ///     &filtered, &budget, || Ok(()),
94    /// )?;
95    /// assert_eq!(tokens.tokens()[0].term(), "韓🙂");
96    /// assert_eq!(tokens.tokens()[0].offsets().unwrap().utf8, 3..10);
97    /// drop(filtered);
98    /// assert!(budget.used() > 0);
99    /// drop(tokens);
100    /// assert_eq!(budget.used(), 0);
101    /// # Ok::<(), uqa_analysis::AnalysisError>(())
102    /// ```
103    pub fn tokenize_mapped_budgeted(
104        &self,
105        text: &crate::FilteredText<'_>,
106        budget: &MemoryBudget,
107        mut poll: impl FnMut() -> AnalysisResult<()>,
108    ) -> AnalysisResult<Budgeted<crate::AnalyzedText>> {
109        poll()?;
110        self.prepare()?
111            .tokenize_mapped_budgeted(text, budget, &mut poll)
112    }
113
114    pub(crate) fn tokenize_mapped(
115        &self,
116        text: &crate::FilteredText<'_>,
117    ) -> AnalysisResult<crate::AnalyzedText> {
118        self.prepare()?.tokenize_mapped(text)
119    }
120}
121
122fn validate_gram_bounds(
123    component: &'static str,
124    min_gram: usize,
125    max_gram: usize,
126) -> AnalysisResult<()> {
127    if min_gram == 0 || max_gram < min_gram {
128        return Err(AnalysisError::InvalidGramBounds {
129            component,
130            min_gram,
131            max_gram,
132        });
133    }
134    Ok(())
135}
136
137pub(super) fn standard_word_class() -> AnalysisResult<&'static regex_syntax::hir::ClassUnicode> {
138    static CLASS: OnceLock<Result<regex_syntax::hir::ClassUnicode, String>> = OnceLock::new();
139    CLASS
140        .get_or_init(|| class(r"\w"))
141        .as_ref()
142        .map_err(|message| AnalysisError::BuiltInRegex {
143            component: "standard tokenizer",
144            message: message.clone(),
145        })
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn whitespace_splits_on_whitespace() {
154        let t = Tokenizer::Whitespace;
155        assert_eq!(
156            t.tokenize("hello  world\n  rust").unwrap(),
157            vec!["hello", "world", "rust"]
158        );
159    }
160
161    #[test]
162    fn standard_extracts_unicode_words() {
163        let t = Tokenizer::Standard;
164        assert_eq!(
165            t.tokenize("Rust 2024! Carácter.").unwrap(),
166            vec!["Rust", "2024", "Carácter"]
167        );
168    }
169
170    #[test]
171    fn letter_extracts_ascii_letters_only() {
172        let t = Tokenizer::Letter;
173        assert_eq!(t.tokenize("abc123 xyz").unwrap(), vec!["abc", "xyz"]);
174    }
175
176    #[test]
177    fn ngram_emits_substrings_per_word() {
178        let t = Tokenizer::NGram {
179            min_gram: 2,
180            max_gram: 3,
181        };
182        // "ab" word: 2-grams [ab]
183        // "abc" word: 2-grams [ab, bc], 3-grams [abc]
184        assert_eq!(t.tokenize("ab abc").unwrap(), vec!["ab", "ab", "bc", "abc"]);
185    }
186
187    #[test]
188    fn pattern_splits_on_regex() {
189        let t = Tokenizer::Pattern {
190            pattern: r"\W+".to_string(),
191        };
192        assert_eq!(t.tokenize("hello, world!").unwrap(), vec!["hello", "world"]);
193    }
194
195    #[test]
196    fn keyword_emits_whole_input() {
197        let t = Tokenizer::Keyword;
198        assert_eq!(t.tokenize("a b c").unwrap(), vec!["a b c"]);
199        assert!(t.tokenize("").unwrap().is_empty());
200    }
201
202    #[test]
203    fn round_trips_via_serde_json() {
204        let t = Tokenizer::NGram {
205            min_gram: 2,
206            max_gram: 4,
207        };
208        let s = serde_json::to_string(&t).unwrap();
209        let back: Tokenizer = serde_json::from_str(&s).unwrap();
210        assert_eq!(
211            back.tokenize("foobar").unwrap(),
212            t.tokenize("foobar").unwrap()
213        );
214    }
215}