Skip to main content

mandarin_to_pinyin/
lib.rs

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// The `Encode` and `Decode` traits are for bincode's native, high-performance serialization.
78// The `Serialize` and `Deserialize` traits are for serde-based formats like JSON.
79// We keep the serde traits for two reasons:
80// 1. Future-proofing: It allows easy serialization to JSON for debugging or other purposes.
81// 2. Ecosystem compatibility: It's standard practice for data structures to be serde-compatible.
82#[derive(Encode, Decode, Debug)]
83pub struct UnicodeMapping {
84    pub mappings: HashMap<u32, Vec<String>>,
85}
86
87/// Deserializes a `UnicodeMapping` from a byte slice.
88/// This is the primary function for loading mapping data.
89pub 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/// Loads the default, embedded unicode-to-pinyin mapping.
96/// This function is only available when the `default-data` feature is enabled.
97#[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    // Pre-allocate string with estimated capacity
247    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            // Check for 2-char pinyin (e.g., "yu1")
272            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(); // consume second letter
279                    chars.pop(); // consume first letter
280                    continue;
281                }
282            }
283
284            // Check for 1-char pinyin (e.g., "a1")
285            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(); // consume the letter
290                    continue;
291                }
292            }
293        }
294        // If not a convertible pinyin, just add the char
295        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        // Test case 1: Basic valid pinyin
358        let input1 = vec!["ni3", "ha3o"];
359        let expected1: Vec<&'static str> = vec!["nǐ", "hǎo"];
360        assert_eq!(test_fn(&input1), expected1);
361
362        // Test case 2: Pinyin with all tones and 'v' for 'ü'
363        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        // Test case 3: Invalid pinyin (number not at the end, or not a valid tone)
368        // The function should pass them through unchanged.
369        let input3 = vec!["pin1yin1", "test5", "a", "b0"];
370        let expected3 = vec!["pin1yin1", "test5", "a", "b0"];
371        assert_eq!(test_fn(&input3), expected3);
372
373        // Test case 4: Mix of valid and invalid pinyin
374        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        // Test case 5: Empty input
379        let input5: Vec<&str> = vec![];
380        let expected5: Vec<String> = vec![];
381        assert_eq!(test_fn(&input5), expected5);
382
383        // Test case 6: Empty strings in the input
384        let input6 = vec!["", "hao3", ""];
385        let expected6 = vec!["", "haǒ", ""];
386        assert_eq!(test_fn(&input6), expected6);
387        
388        // Test case 7: Pinyin where the letter before the number does not form a valid diacritic
389        let input7 = vec!["x1", "z4"];
390        let expected7 = vec!["x1", "z4"];
391        assert_eq!(test_fn(&input7), expected7);
392
393        // Test case 8: Test "yu" mapping
394        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        // Test case 1: Normal sentence
408        let input1 = "你好世界";
409        let expected1 = "nǐ hǎo shì jiè";
410        assert_eq!(to_pinyin_string(input1, " ").unwrap(), expected1);
411
412        // Test case 2: Sentence with non-Chinese characters
413        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        // Test case 3: Empty string
418        let input3 = "";
419        let expected3 = "";
420        assert_eq!(to_pinyin_string(input3, " ").unwrap(), expected3);
421
422        // Test case 4: Different separator
423        let input4 = "你好";
424        let expected4 = "nǐ-hǎo";
425        assert_eq!(to_pinyin_string(input4, "-").unwrap(), expected4);
426    }
427}