use std::collections::HashSet;
use crate::error::Result;
use crate::language_dictionary::LanguageDictionary;
pub const MAX_SEGMENTATIONS: usize = 16;
#[derive(Debug, Clone, Default)]
pub struct Inventory {
syllables: HashSet<String>,
prefixes: HashSet<String>,
longest: usize,
}
impl Inventory {
pub fn new<I, S>(syllables: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
let mut inventory = Self::default();
let mut first_prefixes = [false; 256];
for syllable in syllables {
let syllable = syllable.into();
if syllable.len() > 1 {
let first = usize::from(syllable.as_bytes()[0]);
let known_first = first_prefixes[first];
for end in (1..syllable.len()).rev() {
let prefix = &syllable[..end];
if known_first && (end == 1 || inventory.prefixes.contains(prefix)) {
break;
}
inventory.prefixes.insert(prefix.to_owned());
}
first_prefixes[first] = true;
}
inventory.longest = inventory.longest.max(syllable.len());
inventory.syllables.insert(syllable);
}
inventory
}
pub fn load(dictionary: &LanguageDictionary) -> Result<Self> {
Ok(Self::new(dictionary.syllables()?))
}
pub fn contains(&self, syllable: &str) -> bool {
self.syllables.contains(syllable)
}
pub fn is_prefix(&self, letters: &str) -> bool {
self.prefixes.contains(letters)
}
pub fn syllables(&self) -> impl Iterator<Item = &str> {
self.syllables.iter().map(String::as_str)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Syllable {
pub start: usize,
pub end: usize,
pub complete: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct Segmentation {
pub syllables: Vec<Syllable>,
}
impl Segmentation {
pub fn texts<'a>(&'a self, input: &'a str) -> impl Iterator<Item = &'a str> + 'a {
self.syllables
.iter()
.map(move |syllable| &input[syllable.start..syllable.end])
}
pub fn key(&self, input: &str, count: usize) -> String {
let syllables = self.syllables.iter().take(count);
let capacity = syllables
.clone()
.map(|syllable| syllable.end - syllable.start)
.sum::<usize>()
.saturating_add(self.syllables.len().min(count).saturating_sub(1));
let mut key = String::with_capacity(capacity);
for (index, syllable) in self.syllables.iter().take(count).enumerate() {
if index > 0 {
key.push(' ');
}
key.push_str(&input[syllable.start..syllable.end]);
}
key
}
pub fn key_into(&self, input: &str, count: usize, key: &mut String) {
let syllables = self.syllables.iter().take(count);
let capacity = syllables
.clone()
.map(|syllable| syllable.end - syllable.start)
.sum::<usize>()
.saturating_add(self.syllables.len().min(count).saturating_sub(1));
key.clear();
key.reserve(capacity);
for (index, syllable) in self.syllables.iter().take(count).enumerate() {
if index > 0 {
key.push(' ');
}
key.push_str(&input[syllable.start..syllable.end]);
}
}
pub fn end(&self) -> usize {
self.syllables.last().map_or(0, |syllable| syllable.end)
}
pub fn ends_in_prefix(&self) -> bool {
self.syllables
.last()
.is_some_and(|syllable| !syllable.complete)
}
}
pub fn segment(input: &str, inventory: &Inventory) -> Vec<Segmentation> {
segment_with(input, inventory, true)
}
fn segment_with(input: &str, inventory: &Inventory, allow_prefix: bool) -> Vec<Segmentation> {
let mut found = Vec::new();
let mut dead = vec![false; input.len() + 1];
let mut path = Vec::new();
walk(
input,
inventory,
allow_prefix,
0,
&mut path,
&mut dead,
&mut found,
);
found
}
fn first_segmentation(
input: &str,
inventory: &Inventory,
allow_prefix: bool,
) -> Option<Segmentation> {
let mut dead = vec![false; input.len() + 1];
let mut path = Vec::new();
walk_first(input, inventory, allow_prefix, 0, &mut path, &mut dead)
}
pub fn best(input: &str, inventory: &Inventory) -> Segmentation {
if let Some(full) = first_segmentation(input, inventory, true) {
return full;
}
let mut reachable = vec![false; input.len() + 1];
reachable[0] = true;
let mut furthest = 0;
for position in 0..input.len() {
if !reachable[position] {
continue;
}
if input.as_bytes()[position] == b'\'' {
reachable[position + 1] = true;
continue;
}
let chunk_end = chunk_end(input, position);
for end in position + 1..=chunk_end.min(position + inventory.longest) {
if inventory.contains(&input[position..end]) {
reachable[end] = true;
furthest = furthest.max(end);
}
}
}
segment_with(&input[..furthest], inventory, false)
.into_iter()
.next()
.unwrap_or_default()
}
fn chunk_end(input: &str, position: usize) -> usize {
input[position..]
.find('\'')
.map_or(input.len(), |offset| position + offset)
}
fn push_path(path: &mut Vec<Syllable>, syllable: Syllable, input_length: usize) {
if path.capacity() == 0 {
path.reserve_exact(input_length);
}
path.push(syllable);
}
fn walk(
input: &str,
inventory: &Inventory,
allow_prefix: bool,
mut position: usize,
path: &mut Vec<Syllable>,
dead: &mut [bool],
found: &mut Vec<Segmentation>,
) -> bool {
while input.as_bytes().get(position) == Some(&b'\'') {
position += 1;
}
if position == input.len() {
if found.is_empty() {
found.reserve_exact(MAX_SEGMENTATIONS);
}
found.push(Segmentation {
syllables: path.clone(),
});
return true;
}
if dead[position] {
return false;
}
let chunk_end = chunk_end(input, position);
let max_end = chunk_end.min(position + inventory.longest);
let has_prefix =
allow_prefix && chunk_end == input.len() && inventory.is_prefix(&input[position..]);
let mut any = false;
for end in (position + 1..=max_end).rev() {
if inventory.contains(&input[position..end]) {
push_path(
path,
Syllable {
start: position,
end,
complete: true,
},
input.len(),
);
any |= walk(input, inventory, allow_prefix, end, path, dead, found);
path.pop();
if found.len() >= MAX_SEGMENTATIONS {
return true;
}
}
if has_prefix && end == input.len() {
push_path(
path,
Syllable {
start: position,
end,
complete: false,
},
input.len(),
);
any |= walk(input, inventory, allow_prefix, end, path, dead, found);
path.pop();
if found.len() >= MAX_SEGMENTATIONS {
return true;
}
}
}
if !any {
dead[position] = true;
}
any
}
fn walk_first(
input: &str,
inventory: &Inventory,
allow_prefix: bool,
mut position: usize,
path: &mut Vec<Syllable>,
dead: &mut [bool],
) -> Option<Segmentation> {
while input.as_bytes().get(position) == Some(&b'\'') {
position += 1;
}
if position == input.len() {
return Some(Segmentation {
syllables: path.clone(),
});
}
if dead[position] {
return None;
}
let chunk_end = chunk_end(input, position);
let max_end = chunk_end.min(position + inventory.longest);
let has_prefix =
allow_prefix && chunk_end == input.len() && inventory.is_prefix(&input[position..]);
for end in (position + 1..=max_end).rev() {
if inventory.contains(&input[position..end]) {
push_path(
path,
Syllable {
start: position,
end,
complete: true,
},
input.len(),
);
if let Some(found) = walk_first(input, inventory, allow_prefix, end, path, dead) {
return Some(found);
}
path.pop();
}
if has_prefix && end == input.len() {
push_path(
path,
Syllable {
start: position,
end,
complete: false,
},
input.len(),
);
if let Some(found) = walk_first(input, inventory, allow_prefix, end, path, dead) {
return Some(found);
}
path.pop();
}
}
dead[position] = true;
None
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
pub(crate) const SYLLABLES: [&str; 22] = [
"nei", "hou", "ngo", "ngoi", "oi", "gwong", "dung", "waa", "gwo", "ngok", "dei", "di",
"hoeng", "gong", "jan", "jat", "aa", "m", "ng", "i", "do", "ji",
];
pub(crate) fn inventory() -> Inventory {
Inventory::new(SYLLABLES)
}
fn read(input: &str, segmentation: &Segmentation) -> Vec<String> {
segmentation
.syllables
.iter()
.map(|syllable| {
let text = &input[syllable.start..syllable.end];
if syllable.complete {
text.to_owned()
} else {
format!("{text}*")
}
})
.collect()
}
fn first(input: &str) -> Vec<String> {
read(input, &segment(input, &inventory())[0])
}
#[test]
fn inventory_allocates_only_distinct_prefixes() {
let syllables: Vec<String> = (b'a'..=b'z')
.map(|last| format!("baa{}", char::from(last)))
.collect();
let (inventory, allocations) =
crate::ime::personal_rerank::allocations::count(|| Inventory::new(syllables));
assert_eq!(inventory.longest, 4);
assert_eq!(inventory.syllables.len(), 26);
assert_eq!(inventory.prefixes.len(), 3);
for prefix in ["b", "ba", "baa"] {
assert!(inventory.is_prefix(prefix));
}
for last in b'a'..=b'z' {
let syllable = format!("baa{}", char::from(last));
assert!(inventory.contains(&syllable));
assert!(!inventory.is_prefix(&syllable));
}
assert!(allocations <= 10, "共享前缀库存分配次数: {allocations}");
eprintln!("inventory_shared_prefix_allocations={allocations}");
}
#[test]
fn inventory_preserves_nested_prefixes_and_empty_inputs() {
let empty = Inventory::new(std::iter::empty::<String>());
assert_eq!(empty.longest, 0);
assert!(empty.syllables.is_empty());
assert!(empty.prefixes.is_empty());
for syllables in [
["abc", "ab", "a", "abd", "abc", ""],
["", "abc", "abd", "a", "ab", "abc"],
["a", "ab", "abc", "abd", "abc", ""],
] {
let inventory = Inventory::new(syllables);
assert_eq!(inventory.longest, 3);
assert_eq!(inventory.syllables.len(), 5);
assert_eq!(inventory.prefixes.len(), 2);
for syllable in ["", "a", "ab", "abc", "abd"] {
assert!(inventory.contains(syllable));
}
for prefix in ["a", "ab"] {
assert!(inventory.is_prefix(prefix));
}
for absent in ["", "abc", "abd", "b", "ac"] {
assert!(!inventory.is_prefix(absent));
}
assert_eq!(read("abcabd", &best("abcabd", &inventory)), ["abc", "abd"]);
}
}
#[test]
fn segments_the_design_cases() {
assert_eq!(first("neihou"), ["nei", "hou"]);
assert_eq!(first("gwongdungwaa"), ["gwong", "dung", "waa"]);
assert_eq!(first("ngo'oi"), ["ngo", "oi"]);
assert_eq!(first("ngoi"), ["ngoi"]);
let all = segment("ngoi", &inventory());
assert_eq!(all.capacity(), MAX_SEGMENTATIONS);
let all: Vec<_> = all
.iter()
.map(|segmentation| read("ngoi", segmentation))
.collect();
assert_eq!(all, [vec!["ngoi"], vec!["ngo", "i"], vec!["ng", "oi"]]);
}
#[test]
fn backtracks_when_the_longest_piece_leads_nowhere() {
let small = Inventory::new(["ab", "abc", "cd"]);
let readings = segment("abcd", &small);
assert_eq!(readings.len(), 1);
assert_eq!(read("abcd", &readings[0]), ["ab", "cd"]);
let empty = segment("hoengong", &inventory());
assert!(empty.is_empty());
assert_eq!(empty.capacity(), 0);
assert_eq!(first("hoenggong"), ["hoeng", "gong"]);
}
#[test]
fn the_apostrophe_is_a_hard_boundary() {
assert!(segment("ngo'i'", &inventory())
.iter()
.all(|segmentation| segmentation.syllables.len() == 2));
assert_eq!(first("ngo'i"), ["ngo", "i"]);
assert_eq!(first("ngoi"), ["ngoi"]);
assert!(segment("ng'oi'x", &inventory()).is_empty());
assert!(segment("n'ei", &inventory()).is_empty());
assert_eq!(first("''nei''hou'"), ["nei", "hou"]);
}
#[test]
fn only_the_trailing_letters_may_be_an_incomplete_syllable() {
assert_eq!(first("neih"), ["nei", "h*"]);
assert_eq!(first("neiho"), ["nei", "ho*"]);
assert_eq!(first("gwon"), ["gwon*"]);
assert_eq!(first("gwongd"), ["gwong", "d*"]);
assert!(segment("hx", &inventory()).is_empty());
assert!(segment("", &inventory())
.iter()
.all(|s| s.syllables.is_empty()));
}
#[test]
fn best_falls_back_to_the_longest_leading_run() {
let input = "neihoux";
let reading = best(input, &inventory());
assert_eq!(read(input, &reading), ["nei", "hou"]);
assert_eq!(reading.end(), 6);
let input = "xnei";
assert!(best(input, &inventory()).syllables.is_empty());
let input = "nei'qq";
assert_eq!(read(input, &best(input, &inventory())), ["nei"]);
}
#[test]
fn the_fallback_reads_only_complete_syllables() {
let small = Inventory::new(["ho", "e", "hoeng", "nei"]);
let input = "hoex";
let reading = best(input, &small);
assert_eq!(read(input, &reading), ["ho", "e"]);
assert_eq!(reading.end(), 3);
assert!(!reading.ends_in_prefix());
let input = "neihoex";
assert_eq!(read(input, &best(input, &small)), ["nei", "ho", "e"]);
}
#[test]
fn segmentations_are_capped() {
let input = "ng".repeat(12);
let all = segment(&input, &Inventory::new(["n", "g", "ng"]));
assert_eq!(all.len(), MAX_SEGMENTATIONS);
assert_eq!(all[0].syllables.len(), 12);
}
#[test]
fn walk_does_not_allocate_candidate_pieces() {
let input = "abcx";
let inventory = Inventory::new(["ab", "abc"]);
let mut path = Vec::with_capacity(input.len());
let mut dead = vec![false; input.len() + 1];
let mut found = Vec::with_capacity(MAX_SEGMENTATIONS);
let (readable, allocations) = crate::ime::personal_rerank::allocations::count(|| {
walk(
input, &inventory, false, 0, &mut path, &mut dead, &mut found,
)
});
assert!(!readable);
assert!(found.is_empty());
assert_eq!(allocations, 0);
}
#[test]
fn unreadable_segmentation_does_not_allocate_path_storage() {
let input = "x".repeat(200);
let inventory = inventory();
let (readings, allocations) =
crate::ime::personal_rerank::allocations::count(|| segment(&input, &inventory));
assert!(readings.is_empty());
assert_eq!(allocations, 1, "只应为死路表分配存储: {allocations}");
}
#[test]
fn best_full_reading_does_not_collect_unused_alternatives() {
let inventory = inventory();
let (reading, allocations) =
crate::ime::personal_rerank::allocations::count(|| best("neihou", &inventory));
assert_eq!(read("neihou", &reading), ["nei", "hou"]);
assert!(allocations <= 3, "best reading allocations: {allocations}");
}
#[test]
fn complete_piece_precedes_a_prefix_with_the_same_range() {
let input = "ab";
let inventory = Inventory::new(["ab", "abc"]);
let readings = segment(input, &inventory);
assert_eq!(readings.len(), 2);
assert_eq!(read(input, &readings[0]), ["ab"]);
assert_eq!(read(input, &readings[1]), ["ab*"]);
}
#[test]
fn long_unreadable_input_returns_quickly() {
let input = format!("{}x", "a".repeat(200));
assert!(segment(&input, &Inventory::new(["a", "aa"])).is_empty());
assert_eq!(best(&input, &Inventory::new(["a", "aa"])).end(), 200);
}
#[test]
fn segmentation_keys() {
let input = "gwong'dungwa";
let reading = &segment(input, &inventory())[0];
let key = reading.key(input, 3);
assert_eq!(key, "gwong dung wa");
assert_eq!(key.capacity(), key.len());
assert_eq!(reading.key(input, 2), "gwong dung");
assert!(reading.ends_in_prefix());
assert_eq!(reading.end(), input.len());
}
}