1use icu_segmenter::{WordSegmenter, WordSegmenterBorrowed, options::WordBreakInvariantOptions};
5
6use crate::{TextAnalyzer, TextAnalyzerBuilder, Token, TokenStream, Tokenizer};
7
8#[derive(Clone)]
9pub struct IcuTokenizer {
10 segmenter: WordSegmenterBorrowed<'static>,
11 split_on_non_alphanumeric: bool,
12}
13
14impl Default for IcuTokenizer {
15 fn default() -> Self {
16 Self {
17 segmenter: WordSegmenter::new_dictionary(WordBreakInvariantOptions::default()),
18 split_on_non_alphanumeric: false,
19 }
20 }
21}
22
23impl IcuTokenizer {
24 pub fn with_simple_split(mut self) -> Self {
26 self.split_on_non_alphanumeric = true;
27 self
28 }
29
30 pub fn analyzer(self) -> TextAnalyzer {
31 TextAnalyzer::builder(self).build()
32 }
33
34 pub fn analyzer_builder(self) -> TextAnalyzerBuilder {
35 TextAnalyzer::builder(self).dynamic()
36 }
37}
38
39pub struct IcuTokenStream {
40 tokens: Vec<Token>,
41 index: usize,
42}
43
44fn push_token(tokens: &mut Vec<Token>, text: &str, offset_from: usize, offset_to: usize) {
45 if offset_from == offset_to {
46 return;
47 }
48
49 let token_text = &text[offset_from..offset_to];
50 if token_text.chars().any(char::is_alphanumeric) {
51 tokens.push(Token {
52 offset_from,
53 offset_to,
54 position: tokens.len(),
55 text: token_text.to_owned(),
56 position_length: 1,
57 });
58 }
59}
60
61fn push_tokens_split_on_non_alphanumeric(
62 tokens: &mut Vec<Token>,
63 text: &str,
64 offset_from: usize,
65 offset_to: usize,
66) {
67 let mut part_start = offset_from;
68 for (relative_offset, c) in text[offset_from..offset_to].char_indices() {
69 if !c.is_alphanumeric() {
70 let delimiter_offset = offset_from + relative_offset;
71 push_token(tokens, text, part_start, delimiter_offset);
72 part_start = delimiter_offset + c.len_utf8();
73 }
74 }
75 push_token(tokens, text, part_start, offset_to);
76}
77
78impl TokenStream for IcuTokenStream {
79 fn advance(&mut self) -> bool {
80 if self.index < self.tokens.len() {
81 self.index += 1;
82 true
83 } else {
84 false
85 }
86 }
87
88 fn token(&self) -> &Token {
89 &self.tokens[self.index - 1]
90 }
91
92 fn token_mut(&mut self) -> &mut Token {
93 &mut self.tokens[self.index - 1]
94 }
95}
96
97impl Tokenizer for IcuTokenizer {
98 type TokenStream<'a> = IcuTokenStream;
99
100 fn token_stream<'a>(&'a mut self, text: &'a str) -> Self::TokenStream<'a> {
101 let mut boundaries = self.segmenter.segment_str(text);
102 let mut tokens = Vec::new();
103 let Some(mut offset_from) = boundaries.next() else {
104 return IcuTokenStream { tokens, index: 0 };
105 };
106
107 for offset_to in boundaries {
108 if self.split_on_non_alphanumeric {
109 push_tokens_split_on_non_alphanumeric(&mut tokens, text, offset_from, offset_to);
110 } else {
111 push_token(&mut tokens, text, offset_from, offset_to);
112 }
113 offset_from = offset_to;
114 }
115
116 IcuTokenStream { tokens, index: 0 }
117 }
118}
119
120#[cfg(test)]
121mod tests {
122 use crate::{IcuTokenizer, Token, TokenStream, Tokenizer};
123
124 fn collect_tokens_with_split(text: &str, split_on_non_alphanumeric: bool) -> Vec<Token> {
125 let mut tokenizer = IcuTokenizer::default();
126 if split_on_non_alphanumeric {
127 tokenizer = tokenizer.with_simple_split();
128 }
129 let mut stream = tokenizer.token_stream(text);
130 let mut tokens = Vec::new();
131 stream.process(&mut |token| tokens.push(token.clone()));
132 tokens
133 }
134
135 fn collect_tokens(text: &str) -> Vec<Token> {
136 collect_tokens_with_split(text, false)
137 }
138
139 #[test]
140 fn test_icu_tokenizer_segments_mixed_text() {
141 let tokens = collect_tokens("Hello, こんにちは世界!");
142
143 assert_eq!(
144 tokens
145 .iter()
146 .map(|token| token.text.as_str())
147 .collect::<Vec<_>>(),
148 vec!["Hello", "こんにちは", "世界"]
149 );
150 assert_eq!(
151 tokens
152 .iter()
153 .map(|token| (token.offset_from, token.offset_to, token.position))
154 .collect::<Vec<_>>(),
155 vec![(0, 5, 0), (7, 22, 1), (22, 28, 2)]
156 );
157 }
158
159 #[test]
160 fn test_icu_tokenizer_skips_non_word_segments() {
161 let tokens = collect_tokens("Mark'd ye his words?");
162
163 assert_eq!(
164 tokens
165 .iter()
166 .map(|token| token.text.as_str())
167 .collect::<Vec<_>>(),
168 vec!["Mark'd", "ye", "his", "words"]
169 );
170 }
171
172 #[test]
173 fn test_icu_tokenizer_splits_on_non_alphanumeric_when_enabled() {
174 let tokens = collect_tokens_with_split("foo_bar__baz-alpha.beta", true);
175
176 assert_eq!(
177 tokens
178 .iter()
179 .map(|token| token.text.as_str())
180 .collect::<Vec<_>>(),
181 vec!["foo", "bar", "baz", "alpha", "beta"]
182 );
183 assert_eq!(
184 tokens
185 .iter()
186 .map(|token| (token.offset_from, token.offset_to, token.position))
187 .collect::<Vec<_>>(),
188 vec![(0, 3, 0), (4, 7, 1), (9, 12, 2), (13, 18, 3), (19, 23, 4)]
189 );
190 }
191
192 #[test]
193 fn test_icu_tokenizer_split_control_keeps_icu_segmentation() {
194 let tokens = collect_tokens_with_split("hello_world こんにちは世界", true);
195
196 assert_eq!(
197 tokens
198 .iter()
199 .map(|token| token.text.as_str())
200 .collect::<Vec<_>>(),
201 vec!["hello", "world", "こんにちは", "世界"]
202 );
203 }
204}