1use crate::{median, quantile};
4
5pub type Box4 = [f64; 4];
6pub type Record = (Box4, Option<Box4>, Option<[f64; 2]>, u32, i64);
7const SPACE: u32 = 1;
8const BREAK: u32 = 2;
9const ALNUM: u32 = 4;
10const DECIMAL: u32 = 8;
11const LOCAL: u32 = 16;
12const SLASH: u32 = 32;
13const SEPARATOR: u32 = 64;
14const PROTECTED: u32 = 128;
15const VALID: u32 = 256;
16
17#[derive(Clone)]
18struct Feature {
19 loose: Box4,
20 tight: Option<Box4>,
21 origin: Option<[f64; 2]>,
22 flags: u32,
23 font: i64,
24}
25impl Feature {
26 fn has(&self, flag: u32) -> bool {
28 self.flags & flag != 0
29 }
30 fn lh(&self) -> f64 {
32 self.loose[3] - self.loose[1]
33 }
34 fn th(&self) -> f64 {
36 self.tight.map_or(0.0, |b| b[3] - b[1])
37 }
38 fn lc(&self) -> f64 {
40 (self.loose[1] + self.loose[3]) / 2.0
41 }
42 fn tc(&self) -> f64 {
44 self.tight.map_or(0.0, |b| (b[1] + b[3]) / 2.0)
45 }
46 fn y(&self) -> f64 {
48 self.origin.map_or(0.0, |p| p[1])
49 }
50}
51
52#[derive(Clone)]
53struct Cluster {
54 baseline: f64,
55 indices: Vec<usize>,
56 tight: f64,
57 loose: f64,
58}
59
60fn role(shift: f64) -> u8 {
62 if shift > 0.0 {
63 2
64 } else {
65 1
66 }
67}
68
69fn gap(a: Box4, b: Box4) -> f64 {
71 0.0_f64.max(a[0] - b[2]).max(b[0] - a[2])
72}
73
74fn components(f: &[Feature]) -> Vec<Vec<usize>> {
76 let heights: Vec<_> = f.iter().map(Feature::lh).filter(|v| *v > 0.0).collect();
77 let scale = if heights.is_empty() {
78 1.0
79 } else {
80 median(heights)
81 };
82 let mut result = Vec::new();
83 let mut current = Vec::new();
84 let mut previous: Option<usize> = None;
85 for (i, item) in f.iter().enumerate() {
86 if item.has(SPACE | BREAK) {
87 if !current.is_empty() {
88 result.push(std::mem::take(&mut current));
89 }
90 previous = None;
91 continue;
92 }
93 if item.has(VALID) {
94 if let Some(p) = previous {
95 if (f[p].loose[0] - item.loose[0] > scale * 0.5
96 || item.loose[0] - f[p].loose[2] > scale * 1.5)
97 && !current.is_empty()
98 {
99 result.push(std::mem::take(&mut current));
100 }
101 }
102 }
103 current.push(i);
104 if item.has(VALID) {
105 previous = Some(i);
106 }
107 }
108 if !current.is_empty() {
109 result.push(current);
110 }
111 result
112}
113
114fn clusters(f: &[Feature], indices: &[usize]) -> (Vec<Cluster>, f64) {
116 let mut anchors: Vec<_> = indices
117 .iter()
118 .copied()
119 .filter(|i| f[*i].has(VALID) && f[*i].has(ALNUM) && !f[*i].has(PROTECTED))
120 .collect();
121 if anchors.is_empty() {
122 return (Vec::new(), 0.0);
123 }
124 let tolerance = 0.35_f64.max(median(anchors.iter().map(|i| f[*i].lh()).collect()) * 0.04);
125 anchors.sort_by(|a, b| f[*a].y().partial_cmp(&f[*b].y()).unwrap());
126 let mut groups: Vec<Vec<usize>> = Vec::new();
127 for i in anchors {
128 let append = groups.last().is_some_and(|g| {
129 let n = g.len();
130 let mid = if n % 2 == 1 {
131 f[g[n / 2]].y()
132 } else {
133 (f[g[n / 2 - 1]].y() + f[g[n / 2]].y()) / 2.0
134 };
135 (f[i].y() - mid).abs() <= tolerance
136 });
137 if append {
138 groups.last_mut().unwrap().push(i);
139 } else {
140 groups.push(vec![i]);
141 }
142 }
143 (
144 groups.into_iter().map(|g| make_cluster(f, g)).collect(),
145 tolerance,
146 )
147}
148
149fn make_cluster(f: &[Feature], indices: Vec<usize>) -> Cluster {
151 Cluster {
152 baseline: median(indices.iter().map(|i| f[*i].y()).collect()),
153 tight: quantile(
154 indices
155 .iter()
156 .map(|i| f[*i].th())
157 .filter(|v| *v > 0.0)
158 .collect(),
159 0.9,
160 ),
161 loose: quantile(
162 indices
163 .iter()
164 .map(|i| f[*i].lh())
165 .filter(|v| *v > 0.0)
166 .collect(),
167 0.9,
168 ),
169 indices,
170 }
171}
172
173fn body(cs: &[Cluster]) -> Option<usize> {
175 let max_height = cs.iter().map(|c| c.tight).fold(0.0_f64, f64::max);
176 let mut best: Option<usize> = None;
177 for (i, c) in cs.iter().enumerate() {
178 if c.tight < max_height * 0.9 {
179 continue;
180 }
181 if best.is_none_or(|j| {
182 (c.indices.len(), c.tight, c.loose) > (cs[j].indices.len(), cs[j].tight, cs[j].loose)
183 }) {
184 best = Some(i);
185 }
186 }
187 best
188}
189
190fn center(f: &[Feature], c: &Cluster, tight: bool) -> f64 {
192 median(
193 c.indices
194 .iter()
195 .map(|i| if tight { f[*i].tc() } else { f[*i].lc() })
196 .collect(),
197 )
198}
199
200fn consistent(f: &[Feature], c: &Cluster, b: &Cluster) -> bool {
202 let shift = c.baseline - b.baseline;
203 let t = center(f, c, true) - center(f, b, true);
204 let l = center(f, c, false) - center(f, b, false);
205 let count = [shift, t, l]
206 .iter()
207 .filter(|v| if shift > 0.0 { **v > 0.05 } else { **v < -0.05 })
208 .count();
209 count >= 2
210 && (shift.abs() / b.tight.max(1e-6) >= 0.22
211 || (t.abs() / b.tight.max(1e-6) >= 0.3 && l.abs() / b.loose.max(1e-6) >= 0.15))
212}
213
214fn recheck(f: &[Feature], indices: &[usize], cs: &[Cluster], bidx: usize, roles: &mut [u8]) {
216 let b = &cs[bidx];
217 let initial = roles.to_vec();
218 let order: Vec<_> = std::iter::once(bidx)
219 .chain((0..cs.len()).filter(|i| *i != bidx))
220 .collect();
221 for &ci in &order {
222 let c = &cs[ci];
223 if initial[ci] == 0 || (c.baseline - b.baseline).abs() >= b.tight * 0.3 {
224 continue;
225 }
226 let key = f[c.indices[0]].font;
227 if key < 0 || c.indices.iter().any(|i| f[*i].font != key) {
228 continue;
229 }
230 let belongs = |i: usize| {
231 i < f.len()
232 && indices.contains(&i)
233 && f[i].has(VALID)
234 && f[i].has(LOCAL)
235 && f[i].font == key
236 };
237 let mut left = *c.indices.iter().min().unwrap();
238 let mut right = left + 1;
239 while left > 0 && belongs(left - 1) {
240 left -= 1;
241 }
242 while belongs(right) {
243 right += 1;
244 }
245 if c.indices
246 .iter()
247 .any(|i| *i < left || *i >= right || !belongs(*i))
248 {
249 continue;
250 }
251 let mut best: Option<(f64, usize, Cluster)> = None;
252 for &ri in &order {
253 if initial[ri] != 0 {
254 continue;
255 }
256 let members: Vec<_> = cs[ri]
257 .indices
258 .iter()
259 .copied()
260 .filter(|i| *i >= left && *i < right)
261 .collect();
262 if members.len() < 2 {
263 continue;
264 }
265 let distance = c
266 .indices
267 .iter()
268 .flat_map(|a| {
269 members
270 .iter()
271 .map(|i| gap(f[*a].tight.unwrap(), f[*i].tight.unwrap()))
272 })
273 .fold(f64::INFINITY, f64::min);
274 if distance > 2.5_f64.max(b.tight * 1.2) {
275 continue;
276 }
277 if best
278 .as_ref()
279 .is_none_or(|(d, n, _)| distance < *d || (distance == *d && members.len() > *n))
280 {
281 best = Some((distance, members.len(), make_cluster(f, members)));
282 }
283 }
284 if let Some((_, _, reference)) = best {
285 if reference.tight <= 0.0 {
286 continue;
287 }
288 let shift = c.baseline - reference.baseline;
289 let ratio = c.tight / reference.tight;
290 if shift.abs() < 0.5_f64.max(reference.tight * 0.12)
291 || (ratio > 0.9 && !(shift.abs() >= reference.tight * 0.3 && ratio <= 1.1))
292 {
293 roles[ci] = 0;
294 }
295 }
296 }
297}
298
299fn assign(f: &[Feature], indices: &[usize], roles: &mut [u8]) -> Option<()> {
301 let (cs, tolerance) = clusters(f, indices);
302 let Some(bidx) = body(&cs) else {
303 return Some(());
304 };
305 let b = &cs[bidx];
306 if b.tight <= 0.0 || b.loose <= 0.0 {
307 return Some(());
308 }
309 let mut cr = vec![0; cs.len()];
310 for (i, c) in cs.iter().enumerate() {
311 if i == bidx {
312 continue;
313 }
314 let shift = c.baseline - b.baseline;
315 let tr = c.tight / b.tight;
316 let lr = c.loose / b.loose;
317 if shift.abs() < 0.5_f64.max(b.tight * 0.12)
318 || (tr > 0.9 && !(shift.abs() >= b.tight * 0.3 && tr <= 1.1))
319 || (lr > 1.35 && !consistent(f, c, b))
320 {
321 continue;
322 }
323 cr[i] = role(shift);
324 }
325 recheck(f, indices, &cs, bidx, &mut cr);
326 for &i in indices {
327 if f[i].has(PROTECTED | SPACE | BREAK) || f[i].origin.is_none() {
328 continue;
329 }
330 let mut best = 0;
331 for j in 1..cs.len() {
332 if (f[i].y() - cs[j].baseline).abs() < (f[i].y() - cs[best].baseline).abs() {
333 best = j;
334 }
335 }
336 if (f[i].y() - cs[best].baseline).abs() <= tolerance {
337 roles[i] = cr[best];
338 }
339 }
340 let seeds: Vec<_> = indices
341 .iter()
342 .copied()
343 .filter(|i| roles[*i] != 0 && f[*i].has(ALNUM) && f[*i].tight.is_some())
344 .collect();
345 for &i in indices {
346 if roles[i] == 0 || f[i].has(ALNUM) || f[i].tight.is_none() {
347 continue;
348 }
349 if !seeds.iter().any(|j| {
350 roles[*j] == roles[i]
351 && gap(f[i].tight.unwrap(), f[*j].tight.unwrap()) <= 2.0_f64.max(b.tight * 0.5)
352 && i.abs_diff(*j) <= 4
353 }) {
354 roles[i] = 0;
355 }
356 }
357 for &i in indices {
358 let item = &f[i];
359 if roles[i] != 0 || b.indices.contains(&i) || !item.has(VALID) {
360 continue;
361 }
362 let mut best: Option<(f64, usize, usize)> = None;
363 let mut tied = false;
364 for &j in &b.indices {
365 let key = (gap(item.tight.unwrap(), f[j].tight.unwrap()), i.abs_diff(j));
366 match best {
367 Some((d, n, _)) if key == (d, n) => {
368 tied = true;
369 }
370 Some((d, n, _)) if key > (d, n) => {}
371 _ => {
372 best = Some((key.0, key.1, j));
373 tied = false;
374 }
375 }
376 }
377 let Some((distance, _, j)) = best else {
378 continue;
379 };
380 if tied {
381 return None;
382 }
383 let reference = &f[j];
384 if distance > 2.5_f64.max(reference.th() * 1.2)
385 || item.th() / reference.th().max(1e-6) > 0.88
386 {
387 continue;
388 }
389 let shifts = [
390 item.y() - reference.y(),
391 item.tc() - reference.tc(),
392 item.lc() - reference.lc(),
393 ];
394 let pos = shifts.iter().filter(|v| **v > 0.05).count();
395 let neg = shifts.iter().filter(|v| **v < -0.05).count();
396 if pos.max(neg) < 2 {
397 continue;
398 }
399 let r = if pos > neg { 2 } else { 1 };
400 if role(shifts[0]) != r {
401 continue;
402 }
403 let or = shifts[0].abs() / reference.th().max(1e-6);
404 let tr = shifts[1].abs() / reference.th().max(1e-6);
405 let lr = shifts[2].abs() / reference.lh().max(1e-6);
406 if !item.has(ALNUM)
407 && (!shifts
408 .iter()
409 .all(|v| if r == 2 { *v > 0.05 } else { *v < -0.05 })
410 || !(or >= 0.3 && tr >= 0.3 && lr >= 0.15))
411 {
412 continue;
413 }
414 if or >= 0.22 || (tr >= 0.3 && lr >= 0.15) {
415 roles[i] = r;
416 }
417 }
418 loop {
419 let mut changed = false;
420 for &i in indices {
421 let item = &f[i];
422 if roles[i] != 0
423 || item.has(PROTECTED | SPACE)
424 || !item.has(VALID)
425 || item.th() > b.tight * 1.1
426 {
427 continue;
428 }
429 let previous = i
430 .checked_sub(1)
431 .filter(|j| indices.contains(j))
432 .map_or(0, |j| roles[j]);
433 let next = if indices.contains(&(i + 1)) {
434 roles[i + 1]
435 } else {
436 0
437 };
438 if (previous == 0 && next == 0) || (previous != 0 && next != 0 && previous != next) {
439 continue;
440 }
441 let r = previous.max(next);
442 let shift = item.y() - b.baseline;
443 if shift.abs() < b.tight * 0.08 || role(shift) != r {
444 continue;
445 }
446 roles[i] = r;
447 changed = true;
448 }
449 if !changed {
450 break;
451 }
452 }
453 for &i in indices {
454 if f[i].has(PROTECTED) {
455 roles[i] = 0;
456 }
457 }
458 Some(())
459}
460
461fn numeric(f: &[Feature], roles: &[u8]) -> Vec<usize> {
463 let mut accepted = Vec::new();
464 let mut start = 0;
465 while start < f.len() {
466 if !f[start].has(DECIMAL) {
467 start += 1;
468 continue;
469 }
470 let begin = start;
471 let mut end = start + 1;
472 while end < f.len() && f[end].has(DECIMAL | SEPARATOR) {
473 end += 1;
474 }
475 while end > start && !f[end - 1].has(DECIMAL) {
476 end -= 1;
477 }
478 start = end;
479 if (begin > 0 && f[begin - 1].has(SLASH)) || (end < f.len() && f[end].has(SLASH)) {
480 continue;
481 }
482 let mut previous = begin;
483 while previous > 0 && f[previous - 1].has(SPACE) {
484 if f[previous - 1].has(BREAK) {
485 break;
486 }
487 previous -= 1;
488 }
489 let reference_end = previous;
490 while previous > 0
491 && f[previous - 1].has(VALID)
492 && roles[previous - 1] == 0
493 && !f[previous - 1].has(PROTECTED)
494 {
495 previous -= 1;
496 }
497 if previous == reference_end
498 || (begin..end).any(|i| !f[i].has(VALID) || f[i].has(PROTECTED) || roles[i] == 2)
499 {
500 continue;
501 }
502 let refs: Vec<_> = (previous..reference_end).collect();
503 let (cs, _) = clusters(f, &refs);
504 let Some(bidx) = body(&cs) else {
505 continue;
506 };
507 let b = &cs[bidx];
508 if b.indices.len() < 2 || b.tight <= 0.0 || b.loose <= 0.0 {
509 continue;
510 }
511 let distance = f[begin].tight.unwrap()[0] - f[reference_end - 1].tight.unwrap()[2];
512 if distance < -0.15 * b.tight || distance > 2.5_f64.max(1.2 * b.tight) {
513 continue;
514 }
515 let digits: Vec<_> = (begin..end).filter(|i| f[*i].has(DECIMAL)).collect();
516 let baseline = median(digits.iter().map(|i| f[*i].y()).collect());
517 let shift = b.baseline - baseline;
518 if shift < 0.5_f64.max(0.22 * b.tight) || shift > 0.8 * b.tight {
519 continue;
520 }
521 let tolerance = 0.35_f64.max(0.04 * b.loose);
522 if (begin..end).any(|i| (f[i].y() - baseline).abs() > tolerance) {
523 continue;
524 }
525 let tr = digits.iter().map(|i| f[*i].th()).fold(0.0_f64, f64::max) / b.tight;
526 if tr > 0.88 {
527 let lr = digits.iter().map(|i| f[*i].lh()).fold(0.0_f64, f64::max) / b.loose;
528 if !(shift >= 0.3 * b.tight && tr <= 1.1 && lr <= 0.88) {
529 continue;
530 }
531 }
532 let digit = make_cluster(f, digits);
533 if center(f, &digit, true) >= center(f, b, true) - 0.05
534 || center(f, &digit, false) >= center(f, b, false) - 0.05
535 {
536 continue;
537 }
538 accepted.extend(begin..end);
539 }
540 accepted
541}
542
543pub fn classify(records: Vec<Record>) -> Option<Vec<u8>> {
545 let f: Vec<_> = records
546 .into_iter()
547 .map(|(loose, tight, origin, flags, font)| Feature {
548 loose,
549 tight,
550 origin,
551 flags,
552 font,
553 })
554 .collect();
555 if f.iter().any(|v| {
556 v.loose
557 .iter()
558 .chain(v.tight.iter().flatten())
559 .chain(v.origin.iter().flatten())
560 .any(|x| !x.is_finite() || x.abs() > 1e150)
561 }) {
562 return None;
563 }
564 let mut roles = vec![0; f.len()];
565 for indices in components(&f) {
566 assign(&f, &indices, &mut roles)?;
567 }
568 for i in numeric(&f, &roles) {
569 roles[i] = 1;
570 }
571 Some(roles)
572}