Skip to main content

docvortex_core/
text_assignment.rs

1//! 在自有文本数据上连续执行 tight-first、标点桥接及空白归属,不依赖 Python/PDFium。
2use crate::geometry::Box4;
3use std::collections::HashMap;
4
5pub const SPACE: u8 = 1;
6pub const PUNCTUATION: u8 = 2;
7pub const STOP: u8 = 4;
8pub const START: u8 = 8;
9pub const BREAK: u8 = 16;
10
11pub struct Character {
12    pub bbox: Box4,
13    pub tight: Option<Box4>,
14    pub flags: u8,
15}
16
17/// 与 Python 的有限 bbox 校验一致,空白字符额外允许零面积 advance point。
18fn valid(b: Box4, zero: bool) -> bool {
19    b.iter().all(|value| value.is_finite())
20        && if zero {
21            b[2] >= b[0] && b[3] >= b[1]
22        } else {
23            b[2] > b[0] && b[3] > b[1]
24        }
25}
26
27/// 保持原有中轴阈值及起止标点边界规则,双重标点优先按行尾规则处理。
28fn inside(b: Box4, span: Box4, flags: u8, ratio: f64) -> bool {
29    let x = (b[0] + b[2]) / 2.0;
30    let y = (b[1] + b[3]) / 2.0;
31    let middle = (span[1] + span[3]) / 2.0;
32    let height = span[3] - span[1];
33    if !(span[1] < y && y < span[3] && (y - middle).abs() < height * ratio) {
34        return false;
35    }
36    if span[0] < x && x < span[2] {
37        return true;
38    }
39    if flags & STOP != 0 {
40        span[2] - height < b[0] && b[0] < span[2] && x > span[0]
41    } else if flags & START != 0 {
42        span[0] < b[2] && b[2] < span[0] + height && x < span[2]
43    } else {
44        false
45    }
46}
47
48/// 保留归一化纵/横中心距离及原 span 顺序的稳定排名。
49fn rank(b: Box4, span: Box4, index: usize) -> (f64, f64, usize) {
50    let width = (span[2] - span[0]).max(1e-6);
51    let height = (span[3] - span[1]).max(1e-6);
52    (
53        (((b[1] + b[3]) / 2.0) - ((span[1] + span[3]) / 2.0)).abs() / height,
54        (((b[0] + b[2]) / 2.0) - ((span[0] + span[2]) / 2.0)).abs() / width,
55        index,
56    )
57}
58
59struct Grid {
60    size: f64,
61    ratio: f64,
62    ranges: Vec<(f64, f64)>,
63    cells: Option<HashMap<i64, Vec<usize>>>,
64}
65
66impl Grid {
67    /// 常见范围使用网格;巨大范围改为等价区间扫描,不截断任何候选。
68    fn new(spans: &[Box4], size: f64, ratio: f64) -> Self {
69        let ranges: Vec<_> = spans
70            .iter()
71            .map(|b| ((b[1] / size).trunc(), (b[3] / size).trunc()))
72            .collect();
73        let count: f64 = ranges.iter().map(|(a, b)| (b - a + 1.0).max(0.0)).sum();
74        let ordinary = count <= 1_000_000.0
75            && ranges
76                .iter()
77                .all(|(a, b)| a.abs() < i64::MAX as f64 && b.abs() < i64::MAX as f64);
78        let cells = ordinary.then(|| {
79            let mut cells: HashMap<i64, Vec<usize>> = HashMap::new();
80            for (index, (a, b)) in ranges.iter().enumerate() {
81                for cell in *a as i64..=*b as i64 {
82                    cells.entry(cell).or_default().push(index);
83                }
84            }
85            cells
86        });
87        Self {
88            size,
89            ratio,
90            ranges,
91            cells,
92        }
93    }
94
95    /// 候选顺序保持 span 输入顺序,空白按首个命中而非最小距离归属。
96    fn matching(&self, b: Box4, flags: u8, spans: &[Box4], first: bool) -> Option<usize> {
97        let cell = (((b[1] + b[3]) / 2.0) / self.size).trunc();
98        let x = (b[0] + b[2]) / 2.0;
99        let mut best = None;
100        let mut visit = |index: usize| {
101            let span = spans[index];
102            if flags & (STOP | START) == 0 && !(span[0] < x && x < span[2]) {
103                return false;
104            }
105            if inside(b, span, flags, self.ratio) {
106                if first {
107                    best = Some(index);
108                    return true;
109                }
110                if best.is_none_or(|old| rank(b, span, index) < rank(b, spans[old], old)) {
111                    best = Some(index);
112                }
113            }
114            false
115        };
116        if let Some(cells) = &self.cells {
117            if cell.abs() < i64::MAX as f64 {
118                for &index in cells.get(&(cell as i64)).into_iter().flatten() {
119                    if visit(index) {
120                        break;
121                    }
122                }
123            }
124        } else {
125            for (index, &(a, b)) in self.ranges.iter().enumerate() {
126                if a <= cell && cell <= b && visit(index) {
127                    break;
128                }
129            }
130        }
131        best
132    }
133}
134
135/// 一次完成可见字符匹配、断行隔离、标点补归属与空白修复,输出原字符顺序的 span 索引。
136pub fn assign(
137    chars: &[Character],
138    spans: &[Box4],
139    median_height: f64,
140    ratio: f64,
141) -> Vec<Option<usize>> {
142    let grid = Grid::new(spans, median_height.max(1.0), ratio);
143    let mut assigned = vec![None; chars.len()];
144    for (index, ch) in chars.iter().enumerate() {
145        if ch.flags & SPACE != 0 {
146            continue;
147        }
148        for bbox in [ch.tight, Some(ch.bbox)].into_iter().flatten() {
149            if valid(bbox, false) {
150                assigned[index] = grid.matching(bbox, ch.flags, spans, false);
151                if assigned[index].is_some() {
152                    break;
153                }
154            }
155        }
156    }
157    let mut previous = vec![None; chars.len()];
158    let mut next = vec![None; chars.len()];
159    let mut owner = None;
160    for (index, ch) in chars.iter().enumerate() {
161        if ch.flags & BREAK != 0 {
162            owner = None;
163            continue;
164        }
165        previous[index] = owner;
166        if ch.flags & SPACE == 0 && assigned[index].is_some() {
167            owner = assigned[index];
168        }
169    }
170    owner = None;
171    for (index, ch) in chars.iter().enumerate().rev() {
172        if ch.flags & BREAK != 0 {
173            owner = None;
174            continue;
175        }
176        next[index] = owner;
177        if ch.flags & SPACE == 0 && assigned[index].is_some() {
178            owner = assigned[index];
179        }
180    }
181    for (index, ch) in chars.iter().enumerate() {
182        if assigned[index].is_some() || ch.flags & PUNCTUATION == 0 {
183            continue;
184        }
185        if let Some(owner) = previous[index].filter(|owner| next[index] == Some(*owner)) {
186            let span = spans[owner];
187            if [ch.tight, Some(ch.bbox)].into_iter().flatten().any(|bbox| {
188                let x = (bbox[0] + bbox[2]) / 2.0;
189                let y = (bbox[1] + bbox[3]) / 2.0;
190                valid(bbox, false) && span[0] <= x && x <= span[2] && span[1] <= y && y <= span[3]
191            }) {
192                assigned[index] = Some(owner);
193            }
194        }
195    }
196    for (index, ch) in chars.iter().enumerate() {
197        if ch.flags & SPACE == 0 || !valid(ch.bbox, true) {
198            continue;
199        }
200        let same = previous[index].is_some() && previous[index] == next[index];
201        let neighbor = if same {
202            previous[index]
203        } else {
204            match (previous[index], next[index]) {
205                (owner, None) | (None, owner) => owner,
206                _ => None,
207            }
208        };
209        assigned[index] = if ch.flags & BREAK == 0
210            && neighbor.is_some_and(|owner| same || inside(ch.bbox, spans[owner], ch.flags, ratio))
211        {
212            neighbor
213        } else {
214            grid.matching(ch.bbox, ch.flags, spans, true)
215        };
216    }
217    assigned
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    /// 构造不依赖 PDFium 的真实几何输入,测试归属边界而非绑定实现。
225    fn character(bbox: Box4, flags: u8) -> Character {
226        Character {
227            bbox,
228            tight: None,
229            flags,
230        }
231    }
232
233    #[test]
234    fn tight_precedes_loose_and_ties_keep_span_order() {
235        // tight 命中第二框时不得被 loose 的首框覆盖;完全重叠时保留输入顺序。
236        let mut chars = vec![character([2.0, 2.0, 4.0, 8.0], 0)];
237        chars[0].tight = Some([12.0, 2.0, 14.0, 8.0]);
238        let spans = [[0.0, 0.0, 10.0, 10.0], [10.0, 0.0, 20.0, 10.0]];
239        assert_eq!(assign(&chars, &spans, 10.0, 0.33), [Some(1)]);
240        assert_eq!(assign(&chars, &[spans[1], spans[1]], 10.0, 0.33), [Some(0)]);
241    }
242
243    #[test]
244    fn neighbor_spaces_and_punctuation_do_not_cross_breaks() {
245        // 标点超出中轴带但中心仍在同框内时桥接;断行阻止邻居归属传播。
246        let span = [[0.0, 0.0, 10.0, 10.0]];
247        let chars = [
248            character([1.0, 2.0, 2.0, 8.0], 0),
249            character([3.0, 0.1, 4.0, 0.3], PUNCTUATION),
250            character([4.0, 0.0, 4.0, 0.0], SPACE),
251            character([5.0, 2.0, 6.0, 8.0], 0),
252            character([0.0, -10.0, 0.0, -10.0], SPACE | BREAK),
253            character([7.0, 0.0, 7.0, 0.0], SPACE),
254        ];
255        assert_eq!(
256            assign(&chars, &span, 10.0, 0.33),
257            [Some(0), Some(0), Some(0), Some(0), None, None]
258        );
259    }
260
261    #[test]
262    fn whitespace_uses_first_match_instead_of_nearest() {
263        // 无邻居空白按原候选顺序,而普通字符采用最近中心。
264        let spans = [[0.0, 0.0, 10.0, 20.0], [0.0, 5.0, 10.0, 15.0]];
265        let spaces = [character([4.0, 7.0, 5.0, 9.0], SPACE)];
266        assert_eq!(assign(&spaces, &spans, 10.0, 0.33), [Some(0)]);
267    }
268
269    #[test]
270    fn sparse_grid_preserves_large_range_candidates() {
271        // 巨大网格范围使用等价扫描,仍命中全覆盖 span,不按候选数量裁剪。
272        let spans = [[0.0, 0.0, 10.0, 100_000_000.0]];
273        let chars = [character([3.0, 49_999_999.0, 4.0, 50_000_001.0], 0)];
274        assert_eq!(assign(&chars, &spans, 1.0, 0.33), [Some(0)]);
275    }
276}