1use std::num::NonZeroUsize;
6use std::path::{Path, PathBuf};
7use std::sync::{Arc, LazyLock, Mutex};
8
9use lru::LruCache;
10use memmap2::Mmap;
11
12use super::features::{self, InkStroke};
13use super::model::Model;
14use super::order_handwriting_candidates;
15use crate::{EngineError, Result};
16
17const CANNOT_OPEN: &str = "Cannot open handwriting model";
18const INVALID_CANVAS: &str = "Invalid handwriting canvas";
19const INVALID_STROKE: &str = "Invalid handwriting stroke";
20const INVALID_POINT: &str = "Invalid handwriting point";
21const RECOGNITION_FAILED: &str = "Handwriting recognition failed";
22
23const MAX_CANVAS: f32 = 10_000.0;
24const MAX_STROKES: usize = 64;
25const MAX_STROKE_POINTS: usize = 512;
26const INK_CANVAS: u32 = 1000;
28const INK_BOX: f32 = 800.0;
29const INK_CENTRE: f32 = 500.0;
30const NBEST: usize = 12;
32const MODEL_CACHE_CAPACITY: usize = 8;
33
34static MODELS: LazyLock<Mutex<LruCache<PathBuf, Arc<Model>>>> = LazyLock::new(|| {
35 Mutex::new(LruCache::new(
36 NonZeroUsize::new(MODEL_CACHE_CAPACITY).unwrap(),
37 ))
38});
39
40pub fn handwriting_recognize(
44 model_path: &str,
45 strokes: &[Vec<(f32, f32)>],
46 width: f32,
47 height: f32,
48) -> Result<Vec<String>> {
49 let Some(used) = strokes.iter().rposition(|stroke| !stroke.is_empty()) else {
50 return Ok(Vec::new());
51 };
52 if model_path.is_empty() {
53 return Ok(Vec::new());
54 }
55 let model = load(Path::new(model_path))?;
56 recognize(&model, &strokes[..=used], width, height)
57}
58
59#[allow(unsafe_code)]
60fn load(path: &Path) -> Result<Arc<Model>> {
61 let mut models = MODELS
62 .lock()
63 .unwrap_or_else(|poisoned| poisoned.into_inner());
64 if let Some(model) = models.get(path) {
65 return Ok(Arc::clone(model));
66 }
67 let file =
68 crate::paths::open_file_no_follow(path).map_err(|_| EngineError::failed(CANNOT_OPEN))?;
69 let bytes = unsafe { Mmap::map(&file) }.map_err(|_| EngineError::failed(CANNOT_OPEN))?;
71 let model = Arc::new(Model::parse(bytes).map_err(|_| EngineError::failed(CANNOT_OPEN))?);
72 models.put(path.to_owned(), Arc::clone(&model));
73 Ok(model)
74}
75
76fn recognize(
78 model: &Model,
79 strokes: &[Vec<(f32, f32)>],
80 width: f32,
81 height: f32,
82) -> Result<Vec<String>> {
83 let Some(ink) = ink(strokes, width, height)? else {
84 return Ok(Vec::new());
85 };
86 if model.is_empty() {
87 return Err(EngineError::failed(RECOGNITION_FAILED));
88 }
89 let features = features::extract(&ink, INK_CANVAS)
90 .ok_or_else(|| EngineError::failed(RECOGNITION_FAILED))?;
91 let candidates: Vec<String> = model
92 .classify(&features, NBEST)
93 .into_iter()
94 .map(|(label, _)| label.to_owned())
95 .collect();
96 Ok(order_handwriting_candidates(&candidates))
97}
98
99fn ink(strokes: &[Vec<(f32, f32)>], width: f32, height: f32) -> Result<Option<Vec<InkStroke>>> {
101 let canvas_ok = |size: f32| size.is_finite() && size > 0.0 && size <= MAX_CANVAS;
102 if !canvas_ok(width) || !canvas_ok(height) || strokes.len() > MAX_STROKES {
103 return Err(EngineError::invalid(INVALID_CANVAS));
104 }
105 if strokes.is_empty() {
106 return Ok(None);
107 }
108 let (mut min_x, mut min_y, mut max_x, mut max_y) = (width, height, 0.0_f32, 0.0_f32);
109 for stroke in strokes {
110 if stroke.is_empty() || stroke.len() > MAX_STROKE_POINTS {
111 return Err(EngineError::invalid(INVALID_STROKE));
112 }
113 for &(x, y) in stroke {
114 let inside = x.is_finite()
115 && y.is_finite()
116 && (0.0..=width).contains(&x)
117 && (0.0..=height).contains(&y);
118 if !inside {
119 return Err(EngineError::invalid(INVALID_POINT));
120 }
121 min_x = min_x.min(x);
122 min_y = min_y.min(y);
123 max_x = max_x.max(x);
124 max_y = max_y.max(y);
125 }
126 }
127 let scale = INK_BOX / (max_x - min_x).max(max_y - min_y).max(1.0);
128 let (centre_x, centre_y) = ((min_x + max_x) / 2.0, (min_y + max_y) / 2.0);
129 let ink: Vec<InkStroke> = strokes
131 .iter()
132 .map(|stroke| {
133 stroke
134 .iter()
135 .map(|&(x, y)| {
136 (
137 (INK_CENTRE + (x - centre_x) * scale) as i32,
138 (INK_CENTRE + (y - centre_y) * scale) as i32,
139 )
140 })
141 .collect()
142 })
143 .collect();
144 Ok(Some(ink))
145}
146
147#[cfg(test)]
148mod tests {
149 use super::super::model::tests::{encode, parse};
150 use super::*;
151
152 fn zhong() -> Vec<Vec<(f32, f32)>> {
154 vec![
155 vec![(35.0, 40.0), (35.0, 105.0)],
156 vec![(35.0, 40.0), (125.0, 40.0), (125.0, 105.0)],
157 vec![(35.0, 105.0), (125.0, 105.0)],
158 vec![(80.0, 15.0), (80.0, 140.0)],
159 ]
160 }
161
162 fn message(result: Result<Vec<String>>) -> String {
163 result.unwrap_err().to_string()
164 }
165
166 fn tiny_model() -> Model {
167 parse(&encode(&[
168 ("甲", 0.0, &[(0, 1.0)]),
169 ("x", 0.0, &[(0, 2.0)]),
170 ("乙", 0.0, &[(0, 0.5)]),
171 ]))
172 .unwrap()
173 }
174
175 #[test]
176 fn validates_like_the_reference() {
177 let model = tiny_model();
178 assert!(recognize(&model, &[], 160.0, 155.0).unwrap().is_empty());
179 assert_eq!(message(recognize(&model, &[], 0.0, 155.0)), INVALID_CANVAS);
181 for (width, height) in [
182 (0.0, 155.0),
183 (160.0, -1.0),
184 (f32::NAN, 155.0),
185 (160.0, f32::INFINITY),
186 (10_001.0, 155.0),
187 ] {
188 assert_eq!(
189 message(recognize(&model, &zhong(), width, height)),
190 INVALID_CANVAS
191 );
192 }
193 assert!(recognize(&model, &zhong(), 10_000.0, 10_000.0).is_ok());
194 let many = vec![vec![(1.0, 1.0)]; MAX_STROKES + 1];
195 assert_eq!(
196 message(recognize(&model, &many, 160.0, 155.0)),
197 INVALID_CANVAS
198 );
199 assert!(recognize(&model, &many[..MAX_STROKES], 160.0, 155.0).is_ok());
200 assert_eq!(
201 message(recognize(&model, &[vec![(-1.0, 0.0)]], 160.0, 155.0)),
202 INVALID_POINT
203 );
204 assert_eq!(
205 message(recognize(&model, &[vec![(0.0, 155.5)]], 160.0, 155.0)),
206 INVALID_POINT
207 );
208 assert_eq!(
209 message(recognize(&model, &[vec![(f32::NAN, 1.0)]], 160.0, 155.0)),
210 INVALID_POINT
211 );
212 assert!(recognize(&model, &[vec![(0.0, 0.0), (160.0, 155.0)]], 160.0, 155.0).is_ok());
213 assert_eq!(
214 message(recognize(&model, &[vec![(1.0, 1.0)], vec![]], 160.0, 155.0)),
215 INVALID_STROKE
216 );
217 let long = vec![vec![(1.0, 1.0); MAX_STROKE_POINTS + 1]];
218 assert_eq!(
219 message(recognize(&model, &long, 160.0, 155.0)),
220 INVALID_STROKE
221 );
222 assert_eq!(
224 message(recognize(
225 &model,
226 &[vec![(-1.0, 0.0)], vec![]],
227 160.0,
228 155.0
229 )),
230 INVALID_POINT
231 );
232 }
233
234 #[test]
235 fn classifies_and_orders_candidates() {
236 let ordered = recognize(&tiny_model(), &zhong(), 160.0, 155.0).unwrap();
238 assert_eq!(ordered, vec!["甲", "乙", "x"]);
239 let empty = parse(&encode(&[])).unwrap();
240 assert_eq!(
241 message(recognize(&empty, &zhong(), 160.0, 155.0)),
242 RECOGNITION_FAILED
243 );
244 }
245
246 #[test]
247 fn bridge_entry_skips_empty_input_and_trims_trailing_strokes() {
248 let directory = tempfile::tempdir().unwrap();
249 let path = directory.path().join("tiny.model");
250 std::fs::write(&path, encode(&[("甲", 0.0, &[(0, 1.0)])])).unwrap();
251 let path = path.to_str().unwrap();
252
253 assert!(handwriting_recognize(path, &[], 160.0, 155.0)
254 .unwrap()
255 .is_empty());
256 assert!(handwriting_recognize(path, &[vec![], vec![]], 0.0, 155.0)
257 .unwrap()
258 .is_empty());
259 assert!(handwriting_recognize("", &zhong(), 160.0, 155.0)
260 .unwrap()
261 .is_empty());
262 let mut trailing = zhong();
263 trailing.push(Vec::new());
264 assert_eq!(
265 handwriting_recognize(path, &trailing, 160.0, 155.0).unwrap(),
266 vec!["甲"]
267 );
268 let mut inner = zhong();
269 inner.insert(1, Vec::new());
270 assert_eq!(
271 message(handwriting_recognize(path, &inner, 160.0, 155.0)),
272 INVALID_STROKE
273 );
274 }
275
276 fn parse_strokes(body: &str) -> Vec<Vec<(f32, f32)>> {
277 body.split(';')
278 .map(|stroke| {
279 stroke
280 .split(' ')
281 .map(|point| {
282 let (x, y) = point.split_once(',').unwrap();
283 (x.parse().unwrap(), y.parse().unwrap())
284 })
285 .collect()
286 })
287 .collect()
288 }
289
290 #[test]
292 fn matches_zinnia_on_the_shipped_model() {
293 let Some(path) = std::env::var_os("MSIME_HANDWRITING_MODEL") else {
294 eprintln!("skipped: set MSIME_HANDWRITING_MODEL to handwriting-zh_CN.model");
295 return;
296 };
297 let model = load(Path::new(&path)).unwrap();
298 let fixtures = include_str!("testdata/strokes.tsv").lines();
299 let expected = include_str!("testdata/zinnia-12best.tsv").lines();
300 let mut compared = 0;
301 for (fixture, expected) in fixtures.zip(expected) {
302 let fields: Vec<&str> = fixture.split('\t').collect();
303 let strokes = parse_strokes(fields[3]);
304 let ink = ink(
305 &strokes,
306 fields[1].parse().unwrap(),
307 fields[2].parse().unwrap(),
308 )
309 .unwrap()
310 .unwrap();
311 let features = features::extract(&ink, INK_CANVAS).unwrap();
312 let actual = model.classify(&features, NBEST);
313 let mut columns = expected.split('\t');
314 assert_eq!(columns.next(), Some(fields[0]));
315 let expected: Vec<(&str, f32)> = columns
316 .map(|column| {
317 let (label, score) = column.split_once(' ').unwrap();
318 (label, score.parse().unwrap())
319 })
320 .collect();
321 let labels = |items: &[(&str, f32)]| {
322 items
323 .iter()
324 .map(|item| item.0.to_owned())
325 .collect::<Vec<_>>()
326 };
327 assert_eq!(labels(&actual), labels(&expected), "{fixture}");
328 for ((_, actual), (_, expected)) in actual.iter().zip(&expected) {
329 assert!(
331 (actual - expected).abs() <= 1e-5,
332 "{fixture}: {actual} vs {expected}"
333 );
334 }
335 compared += 1;
336 }
337 assert_eq!(compared, 22);
338
339 let written =
341 handwriting_recognize(path.to_str().unwrap(), &zhong(), 160.0, 155.0).unwrap();
342 assert!(
343 written.iter().any(|candidate| candidate == "中"),
344 "{written:?}"
345 );
346 let moved: Vec<Vec<(f32, f32)>> = zhong()
347 .into_iter()
348 .map(|stroke| {
349 stroke
350 .into_iter()
351 .map(|(x, y)| (x * 0.7 + 180.0, y * 0.7 + 10.0))
352 .collect()
353 })
354 .collect();
355 let moved = handwriting_recognize(path.to_str().unwrap(), &moved, 400.0, 155.0).unwrap();
356 assert!(moved.iter().any(|candidate| candidate == "中"), "{moved:?}");
357 }
358
359 #[test]
361 fn a_model_is_mapped_once_per_path() {
362 let directory = tempfile::tempdir().unwrap();
363 let path = directory.path().join("once.model");
364 std::fs::write(&path, encode(&[("甲", 0.0, &[(0, 1.0)])])).unwrap();
365 let first = load(&path).unwrap();
366 let second = load(&path).unwrap();
367 assert!(Arc::ptr_eq(&first, &second));
368 std::fs::remove_file(&path).unwrap();
369 assert_eq!(
370 handwriting_recognize(path.to_str().unwrap(), &zhong(), 160.0, 155.0).unwrap(),
371 vec!["甲"]
372 );
373 assert_eq!(
375 message(handwriting_recognize(
376 directory.path().to_str().unwrap(),
377 &zhong(),
378 160.0,
379 155.0
380 )),
381 CANNOT_OPEN
382 );
383 }
384
385 #[test]
386 fn the_model_cache_is_bounded_across_paths() {
387 let mut directories = Vec::new();
388 for index in 0..=MODEL_CACHE_CAPACITY {
389 let directory = tempfile::tempdir().unwrap();
390 let path = directory.path().join(format!("model-{index}.model"));
391 std::fs::write(&path, encode(&[("甲", 0.0, &[(0, 1.0)])])).unwrap();
392 assert!(load(&path).is_ok());
393 directories.push(directory);
394 }
395 let cache = MODELS.lock().unwrap();
396 assert!(cache.len() <= MODEL_CACHE_CAPACITY);
397 }
398
399 #[test]
400 fn unreadable_models_fail_to_open() {
401 let directory = tempfile::tempdir().unwrap();
402 let missing = directory.path().join("missing-handwriting-model");
403 let error =
404 handwriting_recognize(missing.to_str().unwrap(), &zhong(), 160.0, 155.0).unwrap_err();
405 assert!(matches!(error, EngineError::Failed(ref text) if text == CANNOT_OPEN));
406 let broken = directory.path().join("broken.model");
407 std::fs::write(&broken, b"not a zinnia model").unwrap();
408 assert_eq!(
409 message(handwriting_recognize(
410 broken.to_str().unwrap(),
411 &zhong(),
412 160.0,
413 155.0
414 )),
415 CANNOT_OPEN
416 );
417 assert_eq!(
419 message(handwriting_recognize(
420 broken.to_str().unwrap(),
421 &zhong(),
422 0.0,
423 155.0
424 )),
425 CANNOT_OPEN
426 );
427 std::fs::write(&missing, encode(&[("甲", 0.0, &[])])).unwrap();
429 assert_eq!(
430 handwriting_recognize(missing.to_str().unwrap(), &zhong(), 160.0, 155.0).unwrap(),
431 vec!["甲"]
432 );
433 }
434}