1use std::collections::HashSet;
4
5use crate::error::Result;
6use crate::language_dictionary::LanguageDictionary;
7
8pub const MAX_SEGMENTATIONS: usize = 16;
10
11#[derive(Debug, Clone, Default)]
13pub struct Inventory {
14 syllables: HashSet<String>,
15 prefixes: HashSet<String>,
16 longest: usize,
17}
18
19impl Inventory {
20 pub fn new<I, S>(syllables: I) -> Self
22 where
23 I: IntoIterator<Item = S>,
24 S: Into<String>,
25 {
26 let mut inventory = Self::default();
27 let mut first_prefixes = [false; 256];
28 for syllable in syllables {
29 let syllable = syllable.into();
30 if syllable.len() > 1 {
31 let first = usize::from(syllable.as_bytes()[0]);
32 let known_first = first_prefixes[first];
34 for end in (1..syllable.len()).rev() {
35 let prefix = &syllable[..end];
36 if known_first && (end == 1 || inventory.prefixes.contains(prefix)) {
38 break;
39 }
40 inventory.prefixes.insert(prefix.to_owned());
41 }
42 first_prefixes[first] = true;
43 }
44 inventory.longest = inventory.longest.max(syllable.len());
45 inventory.syllables.insert(syllable);
46 }
47 inventory
48 }
49
50 pub fn load(dictionary: &LanguageDictionary) -> Result<Self> {
52 Ok(Self::new(dictionary.syllables()?))
53 }
54
55 pub fn contains(&self, syllable: &str) -> bool {
56 self.syllables.contains(syllable)
57 }
58
59 pub fn is_prefix(&self, letters: &str) -> bool {
61 self.prefixes.contains(letters)
62 }
63
64 pub fn syllables(&self) -> impl Iterator<Item = &str> {
66 self.syllables.iter().map(String::as_str)
67 }
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub struct Syllable {
73 pub start: usize,
74 pub end: usize,
75 pub complete: bool,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Default)]
81pub struct Segmentation {
82 pub syllables: Vec<Syllable>,
83}
84
85impl Segmentation {
86 pub fn texts<'a>(&'a self, input: &'a str) -> impl Iterator<Item = &'a str> + 'a {
88 self.syllables
89 .iter()
90 .map(move |syllable| &input[syllable.start..syllable.end])
91 }
92
93 pub fn key(&self, input: &str, count: usize) -> String {
95 let syllables = self.syllables.iter().take(count);
96 let capacity = syllables
97 .clone()
98 .map(|syllable| syllable.end - syllable.start)
99 .sum::<usize>()
100 .saturating_add(self.syllables.len().min(count).saturating_sub(1));
101 let mut key = String::with_capacity(capacity);
102 for (index, syllable) in self.syllables.iter().take(count).enumerate() {
103 if index > 0 {
104 key.push(' ');
105 }
106 key.push_str(&input[syllable.start..syllable.end]);
107 }
108 key
109 }
110
111 pub fn key_into(&self, input: &str, count: usize, key: &mut String) {
112 let syllables = self.syllables.iter().take(count);
113 let capacity = syllables
114 .clone()
115 .map(|syllable| syllable.end - syllable.start)
116 .sum::<usize>()
117 .saturating_add(self.syllables.len().min(count).saturating_sub(1));
118 key.clear();
119 key.reserve(capacity);
120 for (index, syllable) in self.syllables.iter().take(count).enumerate() {
121 if index > 0 {
122 key.push(' ');
123 }
124 key.push_str(&input[syllable.start..syllable.end]);
125 }
126 }
127
128 pub fn end(&self) -> usize {
130 self.syllables.last().map_or(0, |syllable| syllable.end)
131 }
132
133 pub fn ends_in_prefix(&self) -> bool {
135 self.syllables
136 .last()
137 .is_some_and(|syllable| !syllable.complete)
138 }
139}
140
141pub fn segment(input: &str, inventory: &Inventory) -> Vec<Segmentation> {
143 segment_with(input, inventory, true)
144}
145
146fn segment_with(input: &str, inventory: &Inventory, allow_prefix: bool) -> Vec<Segmentation> {
148 let mut found = Vec::new();
149 let mut dead = vec![false; input.len() + 1];
150 let mut path = Vec::new();
151 walk(
152 input,
153 inventory,
154 allow_prefix,
155 0,
156 &mut path,
157 &mut dead,
158 &mut found,
159 );
160 found
161}
162
163fn first_segmentation(
165 input: &str,
166 inventory: &Inventory,
167 allow_prefix: bool,
168) -> Option<Segmentation> {
169 let mut dead = vec![false; input.len() + 1];
170 let mut path = Vec::new();
171 walk_first(input, inventory, allow_prefix, 0, &mut path, &mut dead)
172}
173
174pub fn best(input: &str, inventory: &Inventory) -> Segmentation {
176 if let Some(full) = first_segmentation(input, inventory, true) {
177 return full;
178 }
179 let mut reachable = vec![false; input.len() + 1];
181 reachable[0] = true;
182 let mut furthest = 0;
183 for position in 0..input.len() {
184 if !reachable[position] {
185 continue;
186 }
187 if input.as_bytes()[position] == b'\'' {
188 reachable[position + 1] = true;
189 continue;
190 }
191 let chunk_end = chunk_end(input, position);
192 for end in position + 1..=chunk_end.min(position + inventory.longest) {
193 if inventory.contains(&input[position..end]) {
194 reachable[end] = true;
195 furthest = furthest.max(end);
196 }
197 }
198 }
199 segment_with(&input[..furthest], inventory, false)
201 .into_iter()
202 .next()
203 .unwrap_or_default()
204}
205
206fn chunk_end(input: &str, position: usize) -> usize {
208 input[position..]
209 .find('\'')
210 .map_or(input.len(), |offset| position + offset)
211}
212
213fn push_path(path: &mut Vec<Syllable>, syllable: Syllable, input_length: usize) {
214 if path.capacity() == 0 {
215 path.reserve_exact(input_length);
216 }
217 path.push(syllable);
218}
219
220fn walk(
222 input: &str,
223 inventory: &Inventory,
224 allow_prefix: bool,
225 mut position: usize,
226 path: &mut Vec<Syllable>,
227 dead: &mut [bool],
228 found: &mut Vec<Segmentation>,
229) -> bool {
230 while input.as_bytes().get(position) == Some(&b'\'') {
231 position += 1;
232 }
233 if position == input.len() {
234 if found.is_empty() {
235 found.reserve_exact(MAX_SEGMENTATIONS);
236 }
237 found.push(Segmentation {
238 syllables: path.clone(),
239 });
240 return true;
241 }
242 if dead[position] {
243 return false;
244 }
245 let chunk_end = chunk_end(input, position);
246 let max_end = chunk_end.min(position + inventory.longest);
247 let has_prefix =
248 allow_prefix && chunk_end == input.len() && inventory.is_prefix(&input[position..]);
249 let mut any = false;
250 for end in (position + 1..=max_end).rev() {
252 if inventory.contains(&input[position..end]) {
253 push_path(
254 path,
255 Syllable {
256 start: position,
257 end,
258 complete: true,
259 },
260 input.len(),
261 );
262 any |= walk(input, inventory, allow_prefix, end, path, dead, found);
263 path.pop();
264 if found.len() >= MAX_SEGMENTATIONS {
265 return true;
266 }
267 }
268 if has_prefix && end == input.len() {
269 push_path(
270 path,
271 Syllable {
272 start: position,
273 end,
274 complete: false,
275 },
276 input.len(),
277 );
278 any |= walk(input, inventory, allow_prefix, end, path, dead, found);
279 path.pop();
280 if found.len() >= MAX_SEGMENTATIONS {
281 return true;
282 }
283 }
284 }
285 if !any {
286 dead[position] = true;
287 }
288 any
289}
290
291fn walk_first(
293 input: &str,
294 inventory: &Inventory,
295 allow_prefix: bool,
296 mut position: usize,
297 path: &mut Vec<Syllable>,
298 dead: &mut [bool],
299) -> Option<Segmentation> {
300 while input.as_bytes().get(position) == Some(&b'\'') {
301 position += 1;
302 }
303 if position == input.len() {
304 return Some(Segmentation {
305 syllables: path.clone(),
306 });
307 }
308 if dead[position] {
309 return None;
310 }
311 let chunk_end = chunk_end(input, position);
312 let max_end = chunk_end.min(position + inventory.longest);
313 let has_prefix =
314 allow_prefix && chunk_end == input.len() && inventory.is_prefix(&input[position..]);
315 for end in (position + 1..=max_end).rev() {
316 if inventory.contains(&input[position..end]) {
317 push_path(
318 path,
319 Syllable {
320 start: position,
321 end,
322 complete: true,
323 },
324 input.len(),
325 );
326 if let Some(found) = walk_first(input, inventory, allow_prefix, end, path, dead) {
327 return Some(found);
328 }
329 path.pop();
330 }
331 if has_prefix && end == input.len() {
332 push_path(
333 path,
334 Syllable {
335 start: position,
336 end,
337 complete: false,
338 },
339 input.len(),
340 );
341 if let Some(found) = walk_first(input, inventory, allow_prefix, end, path, dead) {
342 return Some(found);
343 }
344 path.pop();
345 }
346 }
347 dead[position] = true;
348 None
349}
350
351#[cfg(test)]
352pub(crate) mod tests {
353 use super::*;
354
355 pub(crate) const SYLLABLES: [&str; 22] = [
357 "nei", "hou", "ngo", "ngoi", "oi", "gwong", "dung", "waa", "gwo", "ngok", "dei", "di",
358 "hoeng", "gong", "jan", "jat", "aa", "m", "ng", "i", "do", "ji",
359 ];
360
361 pub(crate) fn inventory() -> Inventory {
362 Inventory::new(SYLLABLES)
363 }
364
365 fn read(input: &str, segmentation: &Segmentation) -> Vec<String> {
366 segmentation
367 .syllables
368 .iter()
369 .map(|syllable| {
370 let text = &input[syllable.start..syllable.end];
371 if syllable.complete {
372 text.to_owned()
373 } else {
374 format!("{text}*")
375 }
376 })
377 .collect()
378 }
379
380 fn first(input: &str) -> Vec<String> {
381 read(input, &segment(input, &inventory())[0])
382 }
383
384 #[test]
385 fn inventory_allocates_only_distinct_prefixes() {
386 let syllables: Vec<String> = (b'a'..=b'z')
387 .map(|last| format!("baa{}", char::from(last)))
388 .collect();
389 let (inventory, allocations) =
390 crate::ime::personal_rerank::allocations::count(|| Inventory::new(syllables));
391
392 assert_eq!(inventory.longest, 4);
393 assert_eq!(inventory.syllables.len(), 26);
394 assert_eq!(inventory.prefixes.len(), 3);
395 for prefix in ["b", "ba", "baa"] {
396 assert!(inventory.is_prefix(prefix));
397 }
398 for last in b'a'..=b'z' {
399 let syllable = format!("baa{}", char::from(last));
400 assert!(inventory.contains(&syllable));
401 assert!(!inventory.is_prefix(&syllable));
402 }
403 assert!(allocations <= 10, "共享前缀库存分配次数: {allocations}");
404 eprintln!("inventory_shared_prefix_allocations={allocations}");
405 }
406
407 #[test]
408 fn inventory_preserves_nested_prefixes_and_empty_inputs() {
409 let empty = Inventory::new(std::iter::empty::<String>());
410 assert_eq!(empty.longest, 0);
411 assert!(empty.syllables.is_empty());
412 assert!(empty.prefixes.is_empty());
413 for syllables in [
414 ["abc", "ab", "a", "abd", "abc", ""],
415 ["", "abc", "abd", "a", "ab", "abc"],
416 ["a", "ab", "abc", "abd", "abc", ""],
417 ] {
418 let inventory = Inventory::new(syllables);
419 assert_eq!(inventory.longest, 3);
420 assert_eq!(inventory.syllables.len(), 5);
421 assert_eq!(inventory.prefixes.len(), 2);
422 for syllable in ["", "a", "ab", "abc", "abd"] {
423 assert!(inventory.contains(syllable));
424 }
425 for prefix in ["a", "ab"] {
426 assert!(inventory.is_prefix(prefix));
427 }
428 for absent in ["", "abc", "abd", "b", "ac"] {
429 assert!(!inventory.is_prefix(absent));
430 }
431 assert_eq!(read("abcabd", &best("abcabd", &inventory)), ["abc", "abd"]);
432 }
433 }
434
435 #[test]
436 fn segments_the_design_cases() {
437 assert_eq!(first("neihou"), ["nei", "hou"]);
438 assert_eq!(first("gwongdungwaa"), ["gwong", "dung", "waa"]);
439 assert_eq!(first("ngo'oi"), ["ngo", "oi"]);
440 assert_eq!(first("ngoi"), ["ngoi"]);
441 let all = segment("ngoi", &inventory());
442 assert_eq!(all.capacity(), MAX_SEGMENTATIONS);
443 let all: Vec<_> = all
444 .iter()
445 .map(|segmentation| read("ngoi", segmentation))
446 .collect();
447 assert_eq!(all, [vec!["ngoi"], vec!["ngo", "i"], vec!["ng", "oi"]]);
448 }
449
450 #[test]
451 fn backtracks_when_the_longest_piece_leads_nowhere() {
452 let small = Inventory::new(["ab", "abc", "cd"]);
454 let readings = segment("abcd", &small);
455 assert_eq!(readings.len(), 1);
456 assert_eq!(read("abcd", &readings[0]), ["ab", "cd"]);
457 let empty = segment("hoengong", &inventory());
459 assert!(empty.is_empty());
460 assert_eq!(empty.capacity(), 0);
461 assert_eq!(first("hoenggong"), ["hoeng", "gong"]);
462 }
463
464 #[test]
465 fn the_apostrophe_is_a_hard_boundary() {
466 assert!(segment("ngo'i'", &inventory())
467 .iter()
468 .all(|segmentation| segmentation.syllables.len() == 2));
469 assert_eq!(first("ngo'i"), ["ngo", "i"]);
470 assert_eq!(first("ngoi"), ["ngoi"]);
472 assert!(segment("ng'oi'x", &inventory()).is_empty());
474 assert!(segment("n'ei", &inventory()).is_empty());
475 assert_eq!(first("''nei''hou'"), ["nei", "hou"]);
476 }
477
478 #[test]
479 fn only_the_trailing_letters_may_be_an_incomplete_syllable() {
480 assert_eq!(first("neih"), ["nei", "h*"]);
481 assert_eq!(first("neiho"), ["nei", "ho*"]);
482 assert_eq!(first("gwon"), ["gwon*"]);
484 assert_eq!(first("gwongd"), ["gwong", "d*"]);
485 assert!(segment("hx", &inventory()).is_empty());
486 assert!(segment("", &inventory())
487 .iter()
488 .all(|s| s.syllables.is_empty()));
489 }
490
491 #[test]
492 fn best_falls_back_to_the_longest_leading_run() {
493 let input = "neihoux";
494 let reading = best(input, &inventory());
495 assert_eq!(read(input, &reading), ["nei", "hou"]);
496 assert_eq!(reading.end(), 6);
497 let input = "xnei";
498 assert!(best(input, &inventory()).syllables.is_empty());
499 let input = "nei'qq";
500 assert_eq!(read(input, &best(input, &inventory())), ["nei"]);
501 }
502
503 #[test]
504 fn the_fallback_reads_only_complete_syllables() {
505 let small = Inventory::new(["ho", "e", "hoeng", "nei"]);
507 let input = "hoex";
508 let reading = best(input, &small);
509 assert_eq!(read(input, &reading), ["ho", "e"]);
510 assert_eq!(reading.end(), 3);
511 assert!(!reading.ends_in_prefix());
512 let input = "neihoex";
513 assert_eq!(read(input, &best(input, &small)), ["nei", "ho", "e"]);
514 }
515
516 #[test]
517 fn segmentations_are_capped() {
518 let input = "ng".repeat(12);
520 let all = segment(&input, &Inventory::new(["n", "g", "ng"]));
521 assert_eq!(all.len(), MAX_SEGMENTATIONS);
522 assert_eq!(all[0].syllables.len(), 12);
523 }
524
525 #[test]
526 fn walk_does_not_allocate_candidate_pieces() {
527 let input = "abcx";
528 let inventory = Inventory::new(["ab", "abc"]);
529 let mut path = Vec::with_capacity(input.len());
530 let mut dead = vec![false; input.len() + 1];
531 let mut found = Vec::with_capacity(MAX_SEGMENTATIONS);
532
533 let (readable, allocations) = crate::ime::personal_rerank::allocations::count(|| {
534 walk(
535 input, &inventory, false, 0, &mut path, &mut dead, &mut found,
536 )
537 });
538
539 assert!(!readable);
540 assert!(found.is_empty());
541 assert_eq!(allocations, 0);
542 }
543
544 #[test]
545 fn unreadable_segmentation_does_not_allocate_path_storage() {
546 let input = "x".repeat(200);
547 let inventory = inventory();
548 let (readings, allocations) =
549 crate::ime::personal_rerank::allocations::count(|| segment(&input, &inventory));
550
551 assert!(readings.is_empty());
552 assert_eq!(allocations, 1, "只应为死路表分配存储: {allocations}");
553 }
554
555 #[test]
556 fn best_full_reading_does_not_collect_unused_alternatives() {
557 let inventory = inventory();
558 let (reading, allocations) =
559 crate::ime::personal_rerank::allocations::count(|| best("neihou", &inventory));
560
561 assert_eq!(read("neihou", &reading), ["nei", "hou"]);
562 assert!(allocations <= 3, "best reading allocations: {allocations}");
563 }
564
565 #[test]
566 fn complete_piece_precedes_a_prefix_with_the_same_range() {
567 let input = "ab";
568 let inventory = Inventory::new(["ab", "abc"]);
569 let readings = segment(input, &inventory);
570 assert_eq!(readings.len(), 2);
571 assert_eq!(read(input, &readings[0]), ["ab"]);
572 assert_eq!(read(input, &readings[1]), ["ab*"]);
573 }
574
575 #[test]
576 fn long_unreadable_input_returns_quickly() {
577 let input = format!("{}x", "a".repeat(200));
578 assert!(segment(&input, &Inventory::new(["a", "aa"])).is_empty());
579 assert_eq!(best(&input, &Inventory::new(["a", "aa"])).end(), 200);
580 }
581
582 #[test]
583 fn segmentation_keys() {
584 let input = "gwong'dungwa";
585 let reading = &segment(input, &inventory())[0];
586 let key = reading.key(input, 3);
587 assert_eq!(key, "gwong dung wa");
588 assert_eq!(key.capacity(), key.len());
589 assert_eq!(reading.key(input, 2), "gwong dung");
590 assert!(reading.ends_in_prefix());
591 assert_eq!(reading.end(), input.len());
592 }
593}