1use bincode::{Decode, Encode};
2use phf::phf_map;
3use std::{collections::HashMap, sync::OnceLock};
4
5static UNICODE_TO_PINYIN: OnceLock<HashMap<u32, Vec<String>>> = OnceLock::new();
6
7static DIACRITIC_TO_LETTER: phf::Map<char, &'static str> = phf_map! {
8 'ā' => "a1",
9 'á' => "a2",
10 'ǎ' => "a3",
11 'à' => "a4",
12 'ē' => "e1",
13 'é' => "e2",
14 'ě' => "e3",
15 'è' => "e4",
16 'ī' => "i1",
17 'í' => "i2",
18 'ǐ' => "i3",
19 'ì' => "i4",
20 'ō' => "o1",
21 'ó' => "o2",
22 'ǒ' => "o3",
23 'ò' => "o4",
24 'ǖ' => "v1",
25 'ǘ' => "v2",
26 'ǚ' => "v3",
27 'ǜ' => "v4",
28 'ū' => "u1",
29 'ú' => "u2",
30 'ǔ' => "u3",
31 'ù' => "u4",
32 'ḿ' => "m2",
33 'ń' => "n2",
34 'ň' => "n3",
35 'ǹ' => "n4"
36};
37
38static LETTER_TO_DIACRITIC: phf::Map<&'static str, char> = phf_map! {
39 "a1" => 'ā',
40 "a2" => 'á',
41 "a3" => 'ǎ',
42 "a4" => 'à',
43 "e1" => 'ē',
44 "e2" => 'é',
45 "e3" => 'ě',
46 "e4" => 'è',
47 "i1" => 'ī',
48 "i2" => 'í',
49 "i3" => 'ǐ',
50 "i4" => 'ì',
51 "o1" => 'ō',
52 "o2" => 'ó',
53 "o3" => 'ǒ',
54 "o4" => 'ò',
55 "ü1" => 'ǖ',
56 "ü2" => 'ǘ',
57 "ü3" => 'ǚ',
58 "ü4" => 'ǜ',
59 "yu1" => 'ǖ',
60 "yu2" => 'ǘ',
61 "yu3" => 'ǚ',
62 "yu4" => 'ǜ',
63 "v1" => 'ǖ',
64 "v2" => 'ǘ',
65 "v3" => 'ǚ',
66 "v4" => 'ǜ',
67 "u1" => 'ū',
68 "u2" => 'ú',
69 "u3" => 'ǔ',
70 "u4" => 'ù',
71 "m2" => 'ḿ',
72 "n2" => 'ń',
73 "n3" => 'ň',
74 "n4" => 'ǹ'
75};
76
77#[derive(Encode, Decode, Debug)]
83pub struct UnicodeMapping {
84 pub mappings: HashMap<u32, Vec<String>>,
85}
86
87pub fn load_from_bytes(bytes: &[u8]) -> Result<UnicodeMapping, Box<dyn std::error::Error>> {
90 let (decoded, _len): (UnicodeMapping, usize) =
91 bincode::decode_from_slice(bytes, bincode::config::standard())?;
92 Ok(decoded)
93}
94
95#[cfg(feature = "default-data")]
98pub fn load_default() -> Result<UnicodeMapping, Box<dyn std::error::Error>> {
99 let bytes = include_bytes!("../bincode/unicode-to-pinyin.bin");
100 load_from_bytes(bytes)
101}
102
103pub fn init_map(_bytes: Option<&[u8]>) -> Result<(), Box<dyn std::error::Error>> {
104 let unicode_mapping: UnicodeMapping;
105
106 #[cfg(feature = "default-data")]
107 {
108 unicode_mapping = load_default()?;
109 }
110
111 #[cfg(not(feature = "default-data"))]
112 {
113 let bytes = _bytes.ok_or("bytes is required but None was provided")?;
114 unicode_mapping = load_from_bytes(bytes)?;
115 }
116
117 UNICODE_TO_PINYIN
118 .set(unicode_mapping.mappings)
119 .map_err(|_| "failed to set mappings in OneLock")?;
120
121 Ok(())
122}
123
124#[derive(Debug)]
125pub struct LookupResult<K> {
126 pub map: HashMap<K, Option<Vec<String>>>,
127 pub vec: Vec<Option<Vec<String>>>,
128}
129
130pub fn lookup_unicodes_map(keys: &[u32]) -> Result<HashMap<u32, Option<Vec<String>>>, String> {
131 let map = UNICODE_TO_PINYIN
132 .get()
133 .ok_or("UNICODE_TO_PINYIN not initialized. Call init_map first.")?;
134 Ok(keys
135 .iter()
136 .map(|&key| (key, map.get(&key).cloned()))
137 .collect())
138}
139
140pub fn lookup_unicodes_vec(keys: &[u32]) -> Result<Vec<Option<Vec<String>>>, String> {
141 let map = UNICODE_TO_PINYIN
142 .get()
143 .ok_or("UNICODE_TO_PINYIN not initialized. Call init_map first.")?;
144 Ok(keys.iter().map(|&key| map.get(&key).cloned()).collect())
145}
146
147pub fn lookup_unicodes(keys: &[u32]) -> Result<LookupResult<u32>, String> {
148 let map = UNICODE_TO_PINYIN
149 .get()
150 .ok_or("UNICODE_TO_PINYIN not initialized. Call init_map first.")?;
151 let mut result_map = HashMap::new();
152 let mut values = Vec::new();
153
154 for &key in keys {
155 let value = map.get(&key).cloned();
156 result_map.insert(key, value.clone());
157 values.push(value);
158 }
159
160 Ok(LookupResult {
161 map: result_map,
162 vec: values,
163 })
164}
165
166pub fn lookup_chars_map(keys: &[char]) -> Result<HashMap<char, Option<Vec<String>>>, String> {
167 let map = UNICODE_TO_PINYIN
168 .get()
169 .ok_or("UNICODE_TO_PINYIN not initialized. Call init_map first.")?;
170 Ok(keys
171 .iter()
172 .map(|&key| (key, map.get(&(key as u32)).cloned()))
173 .collect())
174}
175
176pub fn lookup_chars_vec(keys: &[char]) -> Result<Vec<Option<Vec<String>>>, String> {
177 let map = UNICODE_TO_PINYIN
178 .get()
179 .ok_or("UNICODE_TO_PINYIN not initialized. Call init_map first.")?;
180 Ok(keys
181 .iter()
182 .map(|&key| map.get(&(key as u32)).cloned())
183 .collect())
184}
185
186pub fn lookup_chars(keys: &[char]) -> Result<LookupResult<char>, String> {
187 let map = UNICODE_TO_PINYIN
188 .get()
189 .ok_or("UNICODE_TO_PINYIN not initialized. Call init_map first.")?;
190
191 let mut result_map = HashMap::new();
192 let mut values = Vec::new();
193
194 for &key in keys {
195 let value = map.get(&(key as u32)).cloned();
196 result_map.insert(key, value.clone());
197 values.push(value);
198 }
199
200 Ok(LookupResult {
201 map: result_map,
202 vec: values,
203 })
204}
205
206pub fn lookup_chars_map_for_str(s: &str) -> Result<HashMap<char, Option<Vec<String>>>, String> {
207 let keys: Vec<char> = s.chars().collect();
208 lookup_chars_map(&keys)
209}
210
211pub fn lookup_chars_vec_for_str(s: &str) -> Result<Vec<Option<Vec<String>>>, String> {
212 let keys: Vec<char> = s.chars().collect();
213 lookup_chars_vec(&keys)
214}
215
216pub fn lookup_chars_for_str(s: &str) -> Result<LookupResult<char>, String> {
217 let keys: Vec<char> = s.chars().collect();
218 lookup_chars(&keys)
219}
220
221pub fn to_pinyin_string(text: &str, separator: &str) -> Result<String, String> {
222 let map = UNICODE_TO_PINYIN
223 .get()
224 .ok_or("UNICODE_TO_PINYIN not initialized. Call init_map first.")?;
225
226 let pinyins: Vec<String> = text
227 .chars()
228 .map(|c| {
229 map.get(&(c as u32))
230 .and_then(|p_vec| p_vec.get(0))
231 .map_or(c.to_string(), |p| p.clone())
232 })
233 .collect();
234
235 Ok(pinyins.join(separator))
236}
237
238pub fn diacritic_to_tone_plus_number(pinyins: &[&str]) -> Vec<String> {
239 pinyins
240 .iter()
241 .map(|&pinyin| replace_diacritic(pinyin))
242 .collect()
243}
244
245fn replace_diacritic(pinyin: &str) -> String {
246 let mut result = String::with_capacity(pinyin.len() * 2);
248 for c in pinyin.chars() {
249 if let Some(replacement) = DIACRITIC_TO_LETTER.get(&c) {
250 result.push_str(replacement);
251 } else {
252 result.push(c);
253 }
254 }
255 result
256}
257
258pub fn tone_plus_number_to_diacritic(pinyins: &[&str]) -> Vec<String> {
259 pinyins
260 .iter()
261 .map(|&pinyin| replace_numbered_pinyin(pinyin))
262 .collect()
263}
264
265fn replace_numbered_pinyin(pinyin: &str) -> String {
266 let mut result = String::with_capacity(pinyin.len());
267 let mut chars: Vec<char> = pinyin.chars().collect();
268
269 while let Some(current_char) = chars.pop() {
270 if current_char.is_ascii_digit() {
271 if chars.len() >= 2 {
273 let prev_char = chars[chars.len() - 1];
274 let prev_prev_char = chars[chars.len() - 2];
275 let key = format!("{}{}{}", prev_prev_char, prev_char, current_char);
276 if let Some(&diacritic) = LETTER_TO_DIACRITIC.get(key.as_str()) {
277 result.insert(0, diacritic);
278 chars.pop(); chars.pop(); continue;
281 }
282 }
283
284 if let Some(&prev_char) = chars.last() {
286 let key = format!("{}{}", prev_char, current_char);
287 if let Some(&diacritic) = LETTER_TO_DIACRITIC.get(key.as_str()) {
288 result.insert(0, diacritic);
289 chars.pop(); continue;
291 }
292 }
293 }
294 result.insert(0, current_char);
296 }
297 result
298}
299
300#[cfg(feature = "prepare-data")]
301pub fn save_to_vec(
302 pinyin_map: HashMap<u32, Vec<String>>,
303) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
304 let mappings: HashMap<u32, Vec<String>> = pinyin_map
305 .into_iter()
306 .filter_map(|(k, v)| {
307 if k <= 0x10FFFF {
308 Some((k, v))
309 } else {
310 eprintln!("Codepoint {} out of Unicode range", k);
311 None
312 }
313 })
314 .collect();
315
316 let unicode_mapping = UnicodeMapping { mappings };
317
318 let encoded = bincode::encode_to_vec(&unicode_mapping, bincode::config::standard())?;
319 Ok(encoded)
320}
321
322#[cfg(feature = "prepare-data")]
323pub fn load_pinyin_map(
324 pinyin_data_path: &str,
325) -> Result<HashMap<u32, Vec<std::string::String>>, Box<dyn std::error::Error>> {
326 use std::collections::HashMap;
327 use std::fs::File;
328 use std::io::{self, BufRead};
329 use std::path::Path;
330
331 let path = Path::new(pinyin_data_path);
332 let file = File::open(&path)?;
333 let reader = io::BufReader::new(file);
334
335 let mut pinyin_map: HashMap<u32, Vec<String>> = HashMap::new();
336
337 for line in reader.lines() {
338 let line = line?;
339 let parts: Vec<&str> = line.split('\t').collect();
340 if parts.len() == 2 {
341 if let Ok(codepoint) = u32::from_str_radix(parts[0], 16) {
342 let pinyin = parts[1].to_string();
343 let p: Vec<String> = pinyin.split_whitespace().map(|s| s.to_string()).collect();
344 pinyin_map.insert(codepoint, p);
345 }
346 }
347 }
348
349 Ok(pinyin_map)
350}
351
352#[cfg(test)]
353mod tests {
354 use super::*;
355
356 fn run_diacritic_conversion_tests(test_fn: &dyn Fn(&[&str]) -> Vec<String>) {
357 let input1 = vec!["ni3", "ha3o"];
359 let expected1: Vec<&'static str> = vec!["nǐ", "hǎo"];
360 assert_eq!(test_fn(&input1), expected1);
361
362 let input2 = vec!["ma1", "ma2", "ma3", "ma4", "nv3", "lv4"];
364 let expected2 = vec!["mā", "má", "mǎ", "mà", "nǚ", "lǜ"];
365 assert_eq!(test_fn(&input2), expected2);
366
367 let input3 = vec!["pin1yin1", "test5", "a", "b0"];
370 let expected3 = vec!["pin1yin1", "test5", "a", "b0"];
371 assert_eq!(test_fn(&input3), expected3);
372
373 let input4 = vec!["wo3", "shi4", "xue2sheng", "ni2hao3"];
375 let expected4 = vec!["wǒ", "shì", "xuésheng", "níhaǒ"];
376 assert_eq!(test_fn(&input4), expected4);
377
378 let input5: Vec<&str> = vec![];
380 let expected5: Vec<String> = vec![];
381 assert_eq!(test_fn(&input5), expected5);
382
383 let input6 = vec!["", "hao3", ""];
385 let expected6 = vec!["", "haǒ", ""];
386 assert_eq!(test_fn(&input6), expected6);
387
388 let input7 = vec!["x1", "z4"];
390 let expected7 = vec!["x1", "z4"];
391 assert_eq!(test_fn(&input7), expected7);
392
393 let input8 = vec!["yu1", "yu2", "yu3", "yu4", "abcyu4o", "abcu4o"];
395 let expected8 = vec!["ǖ", "ǘ", "ǚ", "ǜ", "abcǜo", "abcùo"];
396 assert_eq!(test_fn(&input8), expected8);
397 }
398
399 #[test]
400 fn test_tone_plus_number_to_diacritic() {
401 run_diacritic_conversion_tests(&tone_plus_number_to_diacritic);
402 }
403
404 #[test]
405 fn test_to_pinyin_string() {
406 init_map(None).unwrap();
407 let input1 = "你好世界";
409 let expected1 = "nǐ hǎo shì jiè";
410 assert_eq!(to_pinyin_string(input1, " ").unwrap(), expected1);
411
412 let input2 = "Hello 你好, world";
414 let expected2 = "H e l l o nǐ hǎo , w o r l d";
415 assert_eq!(to_pinyin_string(input2, " ").unwrap(), expected2);
416
417 let input3 = "";
419 let expected3 = "";
420 assert_eq!(to_pinyin_string(input3, " ").unwrap(), expected3);
421
422 let input4 = "你好";
424 let expected4 = "nǐ-hǎo";
425 assert_eq!(to_pinyin_string(input4, "-").unwrap(), expected4);
426 }
427}