1use std::collections::HashMap;
36
37use docling_core::Node;
38
39use crate::outline::OutlineItem;
40
41#[derive(Clone, Debug)]
44pub struct HeadingHierarchyOptions {
45 pub enabled: bool,
48 pub use_bookmarks: bool,
50 pub use_numbering: bool,
52 pub use_style: bool,
54 pub use_font_style: bool,
57 pub style_size_tolerance: f32,
60 pub max_level: u8,
62 pub bookmark_match_threshold: f32,
65 pub numbering_schemes: Option<Vec<String>>,
69}
70
71impl Default for HeadingHierarchyOptions {
72 fn default() -> Self {
73 Self {
74 enabled: false,
75 use_bookmarks: true,
76 use_numbering: true,
77 use_style: true,
78 use_font_style: true,
79 style_size_tolerance: 0.05,
80 max_level: 6,
81 bookmark_match_threshold: 0.8,
82 numbering_schemes: None,
83 }
84 }
85}
86
87impl HeadingHierarchyOptions {
88 pub fn enabled(on: bool) -> Self {
90 Self {
91 enabled: on,
92 ..Self::default()
93 }
94 }
95}
96
97#[derive(Clone, Copy, Debug)]
101pub(crate) struct GlyphStyle {
102 pub l: f32,
103 pub t: f32,
104 pub r: f32,
105 pub b: f32,
106 pub height: f32,
109 pub weight_cls: u8,
111 pub italic: bool,
112 pub styled: bool,
114}
115
116const DEFAULT_FAMILY_ORDER: [&str; 8] = [
119 "part", "chapter", "article", "roman_u", "arabic", "alpha_u", "alpha_l", "roman_l", ];
128
129#[derive(Clone, Debug, PartialEq)]
133struct Marker {
134 family: &'static str,
135 depth: usize,
137 token: Option<String>,
139 ambiguous: bool,
141}
142
143impl Marker {
144 fn family(family: &'static str) -> Self {
145 Marker {
146 family,
147 depth: 1,
148 token: None,
149 ambiguous: false,
150 }
151 }
152}
153
154fn is_roman(token: &str) -> bool {
158 if token.is_empty() || !token.is_ascii() {
159 return false;
160 }
161 let s: Vec<u8> = token.bytes().map(|b| b.to_ascii_uppercase()).collect();
162 let mut i = 0;
163 let mut m = 0;
165 while i < s.len() && s[i] == b'M' && m < 4 {
166 i += 1;
167 m += 1;
168 }
169 if s[i..].starts_with(b"CM") || s[i..].starts_with(b"CD") {
171 i += 2;
172 } else {
173 if i < s.len() && s[i] == b'D' {
174 i += 1;
175 }
176 let mut c = 0;
177 while i < s.len() && s[i] == b'C' && c < 3 {
178 i += 1;
179 c += 1;
180 }
181 }
182 if s[i..].starts_with(b"XC") || s[i..].starts_with(b"XL") {
184 i += 2;
185 } else {
186 if i < s.len() && s[i] == b'L' {
187 i += 1;
188 }
189 let mut x = 0;
190 while i < s.len() && s[i] == b'X' && x < 3 {
191 i += 1;
192 x += 1;
193 }
194 }
195 if s[i..].starts_with(b"IX") || s[i..].starts_with(b"IV") {
197 i += 2;
198 } else {
199 if i < s.len() && s[i] == b'V' {
200 i += 1;
201 }
202 let mut n = 0;
203 while i < s.len() && s[i] == b'I' && n < 3 {
204 i += 1;
205 n += 1;
206 }
207 }
208 i == s.len()
209}
210
211fn starts_with_word(text: &str, word: &str) -> bool {
213 let Some(head) = text.get(..word.len()) else {
220 return false;
221 };
222 if !head.eq_ignore_ascii_case(word) {
223 return false;
224 }
225 text[word.len()..]
226 .chars()
227 .next()
228 .is_none_or(|c| !c.is_alphanumeric())
229}
230
231fn classify_letter(token: &str) -> Option<Marker> {
233 let upper = token.chars().all(|c| c.is_uppercase());
234 if token.chars().count() == 1 {
235 let is_roman_single = token
236 .chars()
237 .next()
238 .is_some_and(|c| "IVXLCDMivxlcdm".contains(c));
239 let family = match (is_roman_single, upper) {
240 (true, true) => "roman_u",
241 (true, false) => "roman_l",
242 (false, true) => "alpha_u",
243 (false, false) => "alpha_l",
244 };
245 return Some(Marker {
246 family,
247 depth: 1,
248 token: Some(token.to_string()),
249 ambiguous: is_roman_single,
250 });
251 }
252 if is_roman(token) {
255 return Some(Marker {
256 family: if upper { "roman_u" } else { "roman_l" },
257 depth: 1,
258 token: Some(token.to_string()),
259 ambiguous: false,
260 });
261 }
262 None
263}
264
265fn parse_marker(text: &str) -> Option<Marker> {
267 let s = text.trim_start();
268 if s.is_empty() {
269 return None;
270 }
271
272 for kw in ["part", "title", "book"] {
273 if starts_with_word(s, kw) {
274 return Some(Marker::family("part"));
275 }
276 }
277 if starts_with_word(s, "chapter") {
278 return Some(Marker::family("chapter"));
279 }
280 for kw in [
281 "article", "section", "clause", "schedule", "annex", "appendix", "rule",
282 ] {
283 if starts_with_word(s, kw) {
284 return Some(Marker::family("article"));
285 }
286 }
287 if s.starts_with('§') {
289 let after = s.trim_start_matches('§').trim_start();
290 if after.starts_with(|c: char| c.is_ascii_digit()) {
291 return Some(Marker::family("article"));
292 }
293 }
294
295 let body = s.strip_prefix('(').map(str::trim_start).unwrap_or(s);
300
301 if let Some((segments, rest)) = take_dotted(body) {
303 if segments >= 2
304 && (rest
305 .chars()
306 .next()
307 .is_none_or(|c| matches!(c, '.' | ')' | ']') || c.is_whitespace())
308 || separator_at(rest))
309 {
310 return Some(Marker {
311 family: "dotted",
312 depth: segments,
313 token: None,
314 ambiguous: false,
315 });
316 }
317 }
318 let digits = body.chars().take_while(|c| c.is_ascii_digit()).count();
320 if digits > 0 {
321 let rest = body[digits..].trim_start();
322 let dot_not_decimal = rest
323 .strip_prefix('.')
324 .is_some_and(|r| !r.starts_with(|c: char| c.is_ascii_digit()));
325 if dot_not_decimal || separator_at(rest) {
326 return Some(Marker::family("arabic"));
327 }
328 }
329
330 let letters: String = body
333 .chars()
334 .take_while(|c| c.is_ascii_alphabetic())
335 .collect();
336 if !letters.is_empty() {
337 let rest = body[letters.len()..].trim_start();
338 if rest.starts_with('.') || separator_at(rest) {
339 return classify_letter(&letters);
340 }
341 }
342 None
343}
344
345fn separator_at(rest: &str) -> bool {
348 let mut chars = rest.chars();
349 match chars.next() {
350 Some(')' | ']') => true,
351 Some(':' | '-' | '\u{2013}' | '\u{2014}') => chars.next().is_none_or(char::is_whitespace),
352 _ => false,
353 }
354}
355
356fn bare_arabic(text: &str) -> Option<u64> {
359 let s = text.trim();
360 let digits = s.chars().take_while(|c| c.is_ascii_digit()).count();
361 if digits == 0 {
362 return None;
363 }
364 let rest = &s[digits..];
365 let ws = rest.len() - rest.trim_start().len();
366 (ws > 0 && !rest[ws..].is_empty())
367 .then(|| s[..digits].parse().ok())
368 .flatten()
369}
370
371fn resolve_bare_arabic(heading_texts: &[&str], markers: &mut [Option<Marker>]) {
380 let mut sequence: Vec<usize> = Vec::new();
381 for (i, text) in heading_texts.iter().enumerate() {
382 let explicit = match &markers[i] {
383 Some(m) if m.family != "arabic" => continue,
384 Some(_) => true,
385 None => false,
386 };
387 let number = if explicit {
388 let body = text
389 .trim()
390 .strip_prefix('(')
391 .map(str::trim_start)
392 .unwrap_or(text.trim());
393 let digits = body.chars().take_while(|c| c.is_ascii_digit()).count();
394 body[..digits].parse::<u64>().ok()
395 } else {
396 bare_arabic(text)
397 };
398 let Some(number) = number else {
399 continue;
400 };
401 if number == 1 {
402 sequence = vec![i];
403 } else if !sequence.is_empty() && number == sequence.len() as u64 + 1 {
404 sequence.push(i);
405 let n = sequence.len();
408 for &index in &sequence[n.saturating_sub(2)..] {
409 if markers[index].is_none() {
410 markers[index] = Some(Marker::family("arabic"));
411 }
412 }
413 } else {
414 sequence.clear();
415 }
416 }
417}
418
419fn take_dotted(s: &str) -> Option<(usize, &str)> {
422 let mut rest = s;
423 let mut segments = 0;
424 loop {
425 let digits = rest.chars().take_while(|c| c.is_ascii_digit()).count();
426 if digits == 0 {
427 break;
428 }
429 segments += 1;
430 rest = &rest[digits..];
431 match rest.strip_prefix('.') {
432 Some(r) if r.starts_with(|c: char| c.is_ascii_digit()) => rest = r,
435 _ => break,
436 }
437 }
438 (segments >= 2).then_some((segments, rest))
439}
440
441fn resolve_ambiguous(markers: &mut [Option<Marker>]) {
447 let has = |family: &str, ms: &[Option<Marker>]| {
448 ms.iter()
449 .flatten()
450 .any(|m| !m.ambiguous && m.family == family)
451 };
452 let upper_roman = has("roman_u", markers);
453 let upper_alpha = has("alpha_u", markers);
454 let lower_roman = has("roman_l", markers);
455 let lower_alpha = has("alpha_l", markers);
456
457 for m in markers.iter_mut().flatten() {
458 if !m.ambiguous {
459 continue;
460 }
461 let Some(token) = m.token.as_deref() else {
462 continue;
463 };
464 let upper = token.chars().all(|c| c.is_uppercase());
465 let (has_roman, has_alpha) = if upper {
466 (upper_roman, upper_alpha)
467 } else {
468 (lower_roman, lower_alpha)
469 };
470 let roman = if has_roman && !has_alpha {
471 true
472 } else if has_alpha && !has_roman {
473 false
474 } else {
475 token == "I" || token == "i"
476 };
477 m.family = match (roman, upper) {
478 (true, true) => "roman_u",
479 (true, false) => "roman_l",
480 (false, true) => "alpha_u",
481 (false, false) => "alpha_l",
482 };
483 m.ambiguous = false;
484 }
485}
486
487fn family_rank(family: &str, order: &[String]) -> usize {
488 let key = if family == "dotted" { "arabic" } else { family };
489 order.iter().position(|f| f == key).unwrap_or(order.len()) }
491
492fn infer_from_numbering(
494 heading_texts: &[&str],
495 options: &HeadingHierarchyOptions,
496) -> HashMap<usize, usize> {
497 let order: Vec<String> = options
498 .numbering_schemes
499 .clone()
500 .unwrap_or_else(|| DEFAULT_FAMILY_ORDER.iter().map(|s| s.to_string()).collect());
501 let mut markers: Vec<Option<Marker>> = heading_texts.iter().map(|t| parse_marker(t)).collect();
502 resolve_ambiguous(&mut markers);
503 resolve_bare_arabic(heading_texts, &mut markers);
504
505 let mut keys: HashMap<usize, (usize, usize)> = HashMap::new();
506 for (i, m) in markers.iter().enumerate() {
507 if let Some(m) = m {
508 keys.insert(i, (family_rank(m.family, &order), m.depth));
509 }
510 }
511 compress_keys(keys)
512}
513
514fn compress_keys<K: Ord + Clone + std::hash::Hash>(
517 keys: HashMap<usize, K>,
518) -> HashMap<usize, usize> {
519 let mut distinct: Vec<K> = keys.values().cloned().collect();
520 distinct.sort();
521 distinct.dedup();
522 let level_of: HashMap<K, usize> = distinct
523 .into_iter()
524 .enumerate()
525 .map(|(i, k)| (k, i + 1))
526 .collect();
527 keys.into_iter().map(|(i, k)| (i, level_of[&k])).collect()
528}
529
530const ITALIC_RATIO: f32 = 0.6;
534
535fn is_all_caps(text: &str) -> bool {
537 let letters: Vec<char> = text.chars().filter(|c| c.is_alphabetic()).collect();
538 letters.len() >= 4 && letters.iter().all(|c| c.is_uppercase())
539}
540
541#[derive(Clone, Copy, Debug)]
543struct HeadingStyle {
544 size: f32,
545 weight_cls: u8,
546 italic: bool,
547 caps: bool,
548}
549
550fn heading_style(
555 bbox: [f32; 4],
556 text: &str,
557 glyphs: &[GlyphStyle],
558 options: &HeadingHierarchyOptions,
559) -> Option<HeadingStyle> {
560 let [hl, ht, hr, hb] = bbox;
561 let mut heights: Vec<f32> = Vec::new();
562 let mut weights = [0usize; 3];
563 let mut styled_chars = 0usize;
564 let mut italic_chars = 0usize;
565 for g in glyphs {
566 if g.l < hr && g.r > hl && g.t < hb && g.b > ht {
567 heights.push(g.height);
568 if options.use_font_style && g.styled {
569 weights[g.weight_cls.min(2) as usize] += 1;
570 styled_chars += 1;
571 if g.italic {
572 italic_chars += 1;
573 }
574 }
575 }
576 }
577 if heights.is_empty() {
578 return None;
579 }
580 heights.sort_by(f32::total_cmp);
581 let size = if heights.len() % 2 == 1 {
582 heights[heights.len() / 2]
583 } else {
584 (heights[heights.len() / 2 - 1] + heights[heights.len() / 2]) / 2.0
585 };
586 if !options.use_font_style {
587 return Some(HeadingStyle {
588 size,
589 weight_cls: 0,
590 italic: false,
591 caps: false,
592 });
593 }
594 let weight_cls = (0u8..3)
596 .max_by_key(|&cls| (weights[cls as usize], cls))
597 .unwrap_or(0);
598 Some(HeadingStyle {
599 size,
600 weight_cls,
601 italic: styled_chars > 0 && italic_chars as f32 / styled_chars as f32 >= ITALIC_RATIO,
602 caps: is_all_caps(text),
603 })
604}
605
606fn cluster_sizes(mut sizes: Vec<f32>, tolerance: f32) -> Vec<(f32, usize)> {
610 sizes.sort_by(|a, b| b.total_cmp(a));
611 sizes.dedup();
612 let mut clusters = Vec::with_capacity(sizes.len());
613 let mut index = 0usize;
614 let mut previous: Option<f32> = None;
615 for size in sizes {
616 if let Some(prev) = previous {
617 if (prev - size) > tolerance * prev {
618 index += 1;
619 }
620 }
621 clusters.push((size, index));
622 previous = Some(size);
623 }
624 clusters
625}
626
627fn infer_from_style(
629 headings: &[HeadingRef],
630 glyph_styles: &HashMap<usize, Vec<GlyphStyle>>,
631 options: &HeadingHierarchyOptions,
632) -> HashMap<usize, usize> {
633 if glyph_styles.is_empty() {
634 return HashMap::new();
635 }
636 let mut styles: HashMap<usize, HeadingStyle> = HashMap::new();
637 for (i, h) in headings.iter().enumerate() {
638 let Some(glyphs) = glyph_styles.get(&h.page_no) else {
639 continue;
640 };
641 let Some(bbox) = h.bbox_points else { continue };
642 if let Some(style) = heading_style(bbox, &h.text, glyphs, options) {
643 styles.insert(i, style);
644 }
645 }
646 if styles.is_empty() {
647 return HashMap::new();
648 }
649 let clusters = cluster_sizes(
650 styles.values().map(|s| s.size).collect(),
651 options.style_size_tolerance,
652 );
653 let cluster_of = |size: f32| -> usize {
654 clusters
655 .iter()
656 .find(|(s, _)| *s == size)
657 .map(|(_, c)| *c)
658 .unwrap_or(0)
659 };
660 let keys: HashMap<usize, (usize, i8, bool, bool)> = styles
664 .into_iter()
665 .map(|(i, s)| {
666 (
667 i,
668 (cluster_of(s.size), -(s.weight_cls as i8), s.italic, !s.caps),
669 )
670 })
671 .collect();
672 compress_keys(keys)
673}
674
675fn norm(text: &str) -> String {
679 let collapsed = text
680 .split_whitespace()
681 .collect::<Vec<_>>()
682 .join(" ")
683 .to_lowercase();
684 collapsed
685 .trim_matches(|c: char| !c.is_alphanumeric())
686 .to_string()
687}
688
689fn strip_marker(text: &str) -> String {
693 let s = text.trim_start();
694 let matched_len = leading_marker_len(s);
695 match matched_len {
696 Some(n) => {
697 let rest = &s[n..];
698 let trimmed = rest.trim_start_matches(|c: char| {
699 c.is_whitespace() || matches!(c, '.' | ':' | ')' | '-' | '\u{2013}' | '\u{2014}')
700 });
701 trimmed.to_string()
702 }
703 None => text.to_string(),
704 }
705}
706
707fn leading_marker_len(s: &str) -> Option<usize> {
709 for kw in [
711 "chapter", "article", "section", "clause", "schedule", "annex", "appendix", "rule", "part",
712 "title", "book",
713 ] {
714 if starts_with_word(s, kw) {
715 let mut i = kw.len();
716 let bytes = s.as_bytes();
717 while i < bytes.len()
718 && (bytes[i].is_ascii_whitespace() || bytes[i] == b'.' || bytes[i] == b':')
719 {
720 i += 1;
721 }
722 while i < bytes.len()
723 && (bytes[i].is_ascii_digit() || b"ivxlcdmIVXLCDM".contains(&bytes[i]))
724 {
725 i += 1;
726 }
727 return Some(i);
728 }
729 }
730 if s.starts_with('§') {
732 let rest = s.trim_start_matches('§');
733 let ws = rest.len() - rest.trim_start().len();
734 let rest2 = rest.trim_start();
735 let num = rest2
736 .bytes()
737 .take_while(|b| b.is_ascii_digit() || *b == b'.')
738 .count();
739 if num > 0 {
740 return Some(s.len() - rest.len() + ws + num);
741 }
742 }
743 let (paren, body) = match s.strip_prefix('(') {
745 Some(r) => (1, r),
746 None => (0, s),
747 };
748 let digits = body.bytes().take_while(|b| b.is_ascii_digit()).count();
749 if digits > 0 {
750 let mut i = digits;
751 let b = body.as_bytes();
752 while i < b.len() && b[i] == b'.' {
753 let d = body[i + 1..]
754 .bytes()
755 .take_while(|x| x.is_ascii_digit())
756 .count();
757 if d == 0 {
758 break;
759 }
760 i += 1 + d;
761 }
762 if i < b.len() && matches!(b[i], b')' | b'.' | b']') {
763 i += 1;
764 }
765 return Some(paren + i);
766 }
767 let letters = body.bytes().take_while(|b| b.is_ascii_alphabetic()).count();
771 if (1..=2).contains(&letters) {
772 let b = body.as_bytes();
773 if letters < b.len() && matches!(b[letters], b')' | b'.' | b']') {
774 return Some(paren + letters + 1);
775 }
776 let after = body[letters..].trim_start();
777 if after.starts_with([':', '-', '\u{2013}', '\u{2014}']) {
778 return Some(paren + letters);
779 }
780 }
781 None
782}
783
784fn similarity(a: &str, b: &str) -> f32 {
787 let a: Vec<char> = a.chars().collect();
788 let b: Vec<char> = b.chars().collect();
789 if a.is_empty() && b.is_empty() {
790 return 1.0;
791 }
792 let mut b2j: HashMap<char, Vec<usize>> = HashMap::new();
793 for (j, &c) in b.iter().enumerate() {
794 b2j.entry(c).or_default().push(j);
795 }
796 let mut matches = 0usize;
797 let mut queue = vec![(0usize, a.len(), 0usize, b.len())];
798 while let Some((alo, ahi, blo, bhi)) = queue.pop() {
799 let (mut besti, mut bestj, mut bestsize) = (alo, blo, 0usize);
801 let mut j2len: HashMap<usize, usize> = HashMap::new();
802 for (i, ch) in a.iter().enumerate().take(ahi).skip(alo) {
803 let mut newj2len: HashMap<usize, usize> = HashMap::new();
804 if let Some(js) = b2j.get(ch) {
805 for &j in js {
806 if j < blo {
807 continue;
808 }
809 if j >= bhi {
810 break;
811 }
812 let k = j
813 .checked_sub(1)
814 .and_then(|p| j2len.get(&p))
815 .copied()
816 .unwrap_or(0)
817 + 1;
818 newj2len.insert(j, k);
819 if k > bestsize {
820 besti = i + 1 - k;
821 bestj = j + 1 - k;
822 bestsize = k;
823 }
824 }
825 }
826 j2len = newj2len;
827 }
828 if bestsize == 0 {
829 continue;
830 }
831 matches += bestsize;
832 if besti > alo && bestj > blo {
833 queue.push((alo, besti, blo, bestj));
834 }
835 if besti + bestsize < ahi && bestj + bestsize < bhi {
836 queue.push((besti + bestsize, ahi, bestj + bestsize, bhi));
837 }
838 }
839 (2.0 * matches as f32) / (a.len() + b.len()) as f32
840}
841
842fn match_score(cand_text: &str, bm_title: &str) -> f32 {
847 let mut variants_a = vec![norm(cand_text), norm(&strip_marker(cand_text))];
848 let mut variants_b = vec![norm(bm_title), norm(&strip_marker(bm_title))];
849 variants_a.retain(|v| !v.is_empty());
850 variants_b.retain(|v| !v.is_empty());
851 variants_a.dedup();
852 variants_b.dedup();
853 let mut best: f32 = 0.0;
854 for a in &variants_a {
855 for b in &variants_b {
856 best = best.max(similarity(a, b));
857 if a.chars().count() >= 4
858 && b.chars().count() >= 4
859 && (a.contains(b.as_str()) || b.contains(a.as_str()))
860 {
861 best = best.max(0.92);
862 }
863 }
864 }
865 best
866}
867
868struct HeadingRef {
872 node_idx: usize,
874 text: String,
876 page_no: usize,
878 bbox_points: Option<[f32; 4]>,
880 is_list_item: bool,
882}
883
884fn collect(nodes: &[Node], with_list_items: bool) -> Vec<HeadingRef> {
886 let mut out = Vec::new();
887 let mut page_no = 0usize;
888 let mut page_w = 0f32;
889 let mut page_h = 0f32;
890 let denorm = |loc: [u16; 4], w: f32, h: f32| -> Option<[f32; 4]> {
891 (w > 0.0 && h > 0.0).then(|| {
892 [
893 loc[0] as f32 / 512.0 * w,
894 loc[1] as f32 / 512.0 * h,
895 loc[2] as f32 / 512.0 * w,
896 loc[3] as f32 / 512.0 * h,
897 ]
898 })
899 };
900 for (idx, node) in nodes.iter().enumerate() {
901 match node {
902 Node::PageInfo {
903 page_no: p,
904 width,
905 height,
906 } => {
907 page_no = *p;
908 page_w = *width;
909 page_h = *height;
910 }
911 Node::Located { location, inner } => {
912 if let Node::Heading { text, .. } = inner.as_ref() {
913 out.push(HeadingRef {
914 node_idx: idx,
915 text: text.clone(),
916 page_no,
917 bbox_points: denorm(*location, page_w, page_h),
918 is_list_item: false,
919 });
920 }
921 }
922 Node::Heading { text, .. } => out.push(HeadingRef {
923 node_idx: idx,
924 text: text.clone(),
925 page_no,
926 bbox_points: None,
927 is_list_item: false,
928 }),
929 Node::ListItem {
930 ordered,
931 number,
932 text,
933 location,
934 ..
935 } if with_list_items => {
936 let text = if *ordered {
939 format!("{number}. {text}")
940 } else {
941 text.clone()
942 };
943 out.push(HeadingRef {
944 node_idx: idx,
945 text,
946 page_no,
947 bbox_points: location.and_then(|loc| denorm(loc, page_w, page_h)),
948 is_list_item: true,
949 });
950 }
951 _ => {}
952 }
953 }
954 out
955}
956
957pub(crate) fn heading_pages(nodes: &[Node]) -> Vec<usize> {
960 let mut pages: Vec<usize> = collect(nodes, false).iter().map(|h| h.page_no).collect();
961 pages.sort_unstable();
962 pages.dedup();
963 pages.retain(|&p| p > 0);
964 pages
965}
966
967fn infer_from_bookmarks(
970 candidates: &[HeadingRef],
971 outline: &[OutlineItem],
972 options: &HeadingHierarchyOptions,
973) -> HashMap<usize, usize> {
974 let mut claimed: Vec<bool> = vec![false; candidates.len()];
975 let mut matches: Vec<(usize, usize)> = Vec::new(); for bm in outline {
978 let title = bm.title.trim();
979 if title.is_empty() {
980 continue;
981 }
982 let threshold = if bm.page_no.is_none() {
984 (options.bookmark_match_threshold + 0.1).min(1.0)
985 } else {
986 options.bookmark_match_threshold
987 };
988 let mut best: Option<(usize, f32, f32)> = None; for (idx, cand) in candidates.iter().enumerate() {
990 if claimed[idx] {
991 continue;
992 }
993 if let (Some(bp), cp) = (bm.page_no, cand.page_no) {
994 if cp != 0 && cp != bp {
995 continue;
996 }
997 }
998 let score = match_score(&cand.text, title);
999 if score < threshold {
1000 continue;
1001 }
1002 let dist = match (cand.bbox_points.map(|b| b[1]), bm.y_top) {
1003 (Some(top), Some(y)) => (top - y).abs(),
1004 _ => f32::INFINITY,
1005 };
1006 let better = match best {
1007 None => true,
1008 Some((_, bs, bd)) => score > bs + 1e-6 || ((score - bs).abs() <= 1e-6 && dist < bd),
1009 };
1010 if better {
1011 best = Some((idx, score, dist));
1012 }
1013 }
1014 if let Some((idx, _, _)) = best {
1015 claimed[idx] = true;
1016 matches.push((idx, bm.level));
1017 }
1018 }
1019 if matches.is_empty() {
1020 return HashMap::new();
1021 }
1022 compress_keys(matches.into_iter().collect())
1024}
1025
1026pub(crate) fn apply(
1032 nodes: &mut [Node],
1033 outline: &[OutlineItem],
1034 glyph_styles: &HashMap<usize, Vec<GlyphStyle>>,
1035 options: &HeadingHierarchyOptions,
1036) {
1037 if !options.enabled {
1038 return;
1039 }
1040
1041 let mut bookmark_levels: HashMap<usize, usize> = HashMap::new(); if options.use_bookmarks && !outline.is_empty() {
1045 let candidates = collect(nodes, true);
1046 let matched = infer_from_bookmarks(&candidates, outline, options);
1047 for (cand_idx, level) in matched {
1048 let cand = &candidates[cand_idx];
1049 if cand.is_list_item {
1050 promote_list_item(nodes, cand.node_idx, &cand.text);
1051 }
1052 bookmark_levels.insert(cand.node_idx, level);
1053 }
1054 }
1055
1056 let headings = collect(nodes, false);
1057 if headings.is_empty() {
1058 return;
1059 }
1060
1061 let mut levels: HashMap<usize, usize> = HashMap::new(); for (i, h) in headings.iter().enumerate() {
1064 if let Some(level) = bookmark_levels.get(&h.node_idx) {
1065 levels.insert(i, *level);
1066 }
1067 }
1068 if options.use_numbering {
1069 let texts: Vec<&str> = headings.iter().map(|h| h.text.as_str()).collect();
1070 for (i, level) in infer_from_numbering(&texts, options) {
1071 levels.entry(i).or_insert(level);
1072 }
1073 }
1074 if options.use_style && !glyph_styles.is_empty() {
1075 for (i, level) in infer_from_style(&headings, glyph_styles, options) {
1076 levels.entry(i).or_insert(level);
1077 }
1078 }
1079
1080 for (i, h) in headings.iter().enumerate() {
1081 let Some(&level) = levels.get(&i) else {
1082 continue;
1083 };
1084 let semantic = level.clamp(1, options.max_level.max(1) as usize);
1085 let rendered = (semantic + 1).min(u8::MAX as usize) as u8;
1088 set_heading_level(&mut nodes[h.node_idx], rendered);
1089 }
1090}
1091
1092fn set_heading_level(node: &mut Node, new_level: u8) {
1093 match node {
1094 Node::Heading { level, .. } => *level = new_level,
1095 Node::Located { inner, .. } => {
1096 if let Node::Heading { level, .. } = inner.as_mut() {
1097 *level = new_level;
1098 }
1099 }
1100 _ => {}
1101 }
1102}
1103
1104fn promote_list_item(nodes: &mut [Node], idx: usize, text: &str) {
1109 let Node::ListItem {
1110 first_in_list,
1111 location,
1112 ..
1113 } = &nodes[idx]
1114 else {
1115 return;
1116 };
1117 let was_first = *first_in_list;
1118 let loc = *location;
1119 let heading = Node::Heading {
1120 level: 2,
1121 text: text.to_string(),
1122 };
1123 nodes[idx] = match loc {
1124 Some(location) => Node::Located {
1125 location,
1126 inner: Box::new(heading),
1127 },
1128 None => heading,
1129 };
1130 if was_first {
1131 if let Some(Node::ListItem { first_in_list, .. }) = nodes.get_mut(idx + 1) {
1132 *first_in_list = true;
1133 }
1134 }
1135}
1136
1137#[cfg(test)]
1138mod tests {
1139 use super::*;
1140
1141 #[test]
1146 fn keyword_probe_never_slices_mid_char() {
1147 assert!(!starts_with_word("Note 1\u{a0}Overview", "chapter"));
1148 for word in ["chapter", "section", "part", "article", "appendix", "annex"] {
1149 for k in 0..12 {
1150 for ch in ['\u{a0}', '\u{e9}', '\u{3a9}', '\u{1f600}'] {
1151 let text = format!("{}{ch}Overview", "x".repeat(k));
1152 assert!(!starts_with_word(&text, word), "{text:?} vs {word}");
1153 }
1154 }
1155 }
1156 assert!(starts_with_word("Chapter\u{a0}1", "chapter"));
1158 assert!(starts_with_word("CHAPTER 2 Scope", "chapter"));
1159 assert!(starts_with_word("Section", "section"));
1160 assert!(!starts_with_word("Chapters", "chapter"));
1161 assert!(!starts_with_word("Chapt", "chapter"));
1162 }
1163
1164 fn heading(loc: [u16; 4], text: &str) -> Node {
1165 Node::Located {
1166 location: loc,
1167 inner: Box::new(Node::Heading {
1168 level: 2,
1169 text: text.to_string(),
1170 }),
1171 }
1172 }
1173
1174 fn page(no: usize) -> Node {
1175 Node::PageInfo {
1176 page_no: no,
1177 width: 512.0,
1178 height: 512.0,
1179 }
1180 }
1181
1182 fn levels(nodes: &[Node]) -> Vec<u8> {
1183 nodes
1184 .iter()
1185 .filter_map(|n| match n {
1186 Node::Located { inner, .. } => match inner.as_ref() {
1187 Node::Heading { level, .. } => Some(*level),
1188 _ => None,
1189 },
1190 Node::Heading { level, .. } => Some(*level),
1191 _ => None,
1192 })
1193 .collect()
1194 }
1195
1196 #[test]
1197 fn roman_validator_matches_difflib_regex() {
1198 for ok in ["I", "iv", "XIV", "MCMXCIX", "iii", "C"] {
1199 assert!(is_roman(ok), "{ok}");
1200 }
1201 for bad in ["", "IIII", "VX", "ABC", "Summary", "IC"] {
1202 assert!(!is_roman(bad), "{bad}");
1203 }
1204 }
1205
1206 #[test]
1207 fn markers_parse_the_docling_families() {
1208 let fam = |t: &str| parse_marker(t).map(|m| (m.family, m.depth));
1209 assert_eq!(fam("PART I — General"), Some(("part", 1)));
1210 assert_eq!(fam("Chapter 2: Scope"), Some(("chapter", 1)));
1211 assert_eq!(fam("Article 5"), Some(("article", 1)));
1212 assert_eq!(fam("§ 12 Something"), Some(("article", 1)));
1213 assert_eq!(fam("1. Introduction"), Some(("arabic", 1)));
1214 assert_eq!(fam("2) Also arabic"), Some(("arabic", 1)));
1215 assert_eq!(fam("1.1 Scope"), Some(("dotted", 2)));
1216 assert_eq!(fam("2.3.1 Deep"), Some(("dotted", 3)));
1217 assert_eq!(fam("A. Annex-ish"), Some(("alpha_u", 1)));
1218 assert_eq!(fam("(a) item"), Some(("alpha_l", 1)));
1219 assert_eq!(fam("(iv) sub"), Some(("roman_l", 1)));
1220 assert_eq!(fam("IV. Chapter"), Some(("roman_u", 1)));
1221 assert_eq!(fam("Summary."), None);
1223 assert_eq!(fam("Overview"), None);
1224 assert_eq!(fam("1: Introduction"), Some(("arabic", 1)));
1226 assert_eq!(fam("(2) Scope"), Some(("arabic", 1)));
1227 assert_eq!(fam("3] Notes"), Some(("arabic", 1)));
1228 assert_eq!(fam("4 - Methods"), Some(("arabic", 1)));
1229 assert_eq!(fam("5 — Results"), Some(("arabic", 1)));
1230 assert_eq!(fam("1.2: Detail"), Some(("dotted", 2)));
1231 assert_eq!(fam("(1.2.3) Deep"), Some(("dotted", 3)));
1232 assert_eq!(fam("A: Annex"), Some(("alpha_u", 1)));
1233 assert_eq!(fam("B - Bee"), Some(("alpha_u", 1)));
1234 assert_eq!(fam("(e) see"), Some(("alpha_l", 1)));
1235 assert_eq!(fam("1-A Something"), None);
1238 assert_eq!(fam("2 Methods"), None);
1239 assert_eq!(fam("1.5 kg"), Some(("dotted", 2)));
1240 }
1241
1242 #[test]
1246 fn bare_arabic_chapter_numbers_need_a_sequence() {
1247 let opts = HeadingHierarchyOptions::default();
1248 let texts = ["1 Intro", "1.1 Scope", "2. Methods", "3 Results", "Plain"];
1249 let map = infer_from_numbering(&texts.map(|t| t), &opts);
1250 assert_eq!(map[&0], 1, "confirmed by the explicit `2.`");
1251 assert_eq!(map[&1], 2);
1252 assert_eq!(map[&2], 1);
1253 assert_eq!(map[&3], 1, "3 follows the run 1, 2");
1254 assert!(!map.contains_key(&4));
1255 let texts = ["1 Intro", "2 Methods", "4 Results"];
1256 let map = infer_from_numbering(&texts.map(|t| t), &opts);
1257 assert_eq!(map.get(&0), Some(&1));
1258 assert_eq!(map.get(&1), Some(&1));
1259 assert!(!map.contains_key(&2), "the gap leaves 4 unrecognized");
1260 let texts = ["2024 Report", "3 kg", "Summary"];
1261 assert!(infer_from_numbering(&texts.map(|t| t), &opts).is_empty());
1262 }
1263
1264 #[test]
1265 fn ambiguous_single_letters_resolve_from_document_context() {
1266 let texts = ["I. One", "II. Two", "V. Five"];
1268 let map = infer_from_numbering(&texts.map(|t| t), &HeadingHierarchyOptions::default());
1269 assert_eq!(map[&0], map[&2]);
1271 let texts = ["B. Bee", "C. Sea", "D. Dee"];
1273 let map = infer_from_numbering(&texts.map(|t| t), &HeadingHierarchyOptions::default());
1274 assert_eq!(map[&0], map[&1]);
1275 assert_eq!(map[&1], map[&2]);
1276 }
1277
1278 #[test]
1279 fn numbering_levels_compress_to_contiguous() {
1280 let texts = ["PART I", "1.1 Scope", "1.1.1 Detail", "No marker"];
1283 let map = infer_from_numbering(&texts.map(|t| t), &HeadingHierarchyOptions::default());
1284 assert_eq!(map[&0], 1);
1285 assert_eq!(map[&1], 2);
1286 assert_eq!(map[&2], 3);
1287 assert!(!map.contains_key(&3));
1288 }
1289
1290 #[test]
1291 fn similarity_behaves_like_difflib_ratio() {
1292 assert_eq!(similarity("abc", "abc"), 1.0);
1293 assert_eq!(similarity("", ""), 1.0);
1294 assert_eq!(similarity("abc", "xyz"), 0.0);
1295 assert!((similarity("abcd", "bcde") - 0.75).abs() < 1e-6);
1297 }
1298
1299 #[test]
1300 fn bookmark_titles_match_with_and_without_markers() {
1301 assert!(match_score("1.1 Definitions", "Definitions") >= 0.9);
1302 assert!(match_score("ARTICLE 5 Payment Terms", "Payment Terms") >= 0.9);
1303 assert!(match_score("Introduction", "Conclusion") < 0.8);
1304 }
1305
1306 #[test]
1307 fn apply_assigns_numbering_levels_end_to_end() {
1308 let mut nodes = vec![
1309 page(1),
1310 heading([10, 10, 200, 20], "1. Introduction"),
1311 heading([10, 40, 200, 50], "1.1 Scope"),
1312 heading([10, 70, 200, 80], "Unnumbered"),
1313 ];
1314 apply(
1315 &mut nodes,
1316 &[],
1317 &HashMap::new(),
1318 &HeadingHierarchyOptions::enabled(true),
1319 );
1320 assert_eq!(levels(&nodes), vec![2, 3, 2]);
1323 }
1324
1325 #[test]
1329 fn apply_survives_multibyte_headings_and_bookmarks() {
1330 let mut nodes = vec![
1331 page(1),
1332 heading([10, 10, 200, 20], "Note 1\u{a0}Overview"),
1333 heading([10, 40, 200, 50], "1.\u{a0}Einf\u{fc}hrung"),
1334 heading(
1335 [10, 70, 200, 80],
1336 "1.1\u{a0}\u{dc}berblick \u{2014} Teil\u{a0}A",
1337 ),
1338 heading([10, 100, 200, 110], "Chapter\u{a0}2\u{a0}\u{3a9}mega"),
1339 heading([10, 130, 200, 140], "\u{1f600} Anhang"),
1340 ];
1341 let outline = vec![
1342 OutlineItem {
1343 title: "Note\u{a0}1 Overview".into(),
1344 level: 0,
1345 page_no: Some(1),
1346 y_top: None,
1347 },
1348 OutlineItem {
1349 title: "Einf\u{fc}hrung".into(),
1350 level: 1,
1351 page_no: Some(1),
1352 y_top: None,
1353 },
1354 ];
1355 apply(
1356 &mut nodes,
1357 &outline,
1358 &HashMap::new(),
1359 &HeadingHierarchyOptions::enabled(true),
1360 );
1361 assert_eq!(levels(&nodes).len(), 5);
1362 }
1363
1364 #[test]
1365 fn apply_is_inert_when_disabled() {
1366 let mut nodes = vec![page(1), heading([10, 10, 200, 20], "1.1.1 Deep")];
1367 apply(
1368 &mut nodes,
1369 &[],
1370 &HashMap::new(),
1371 &HeadingHierarchyOptions::default(),
1372 );
1373 assert_eq!(levels(&nodes), vec![2]);
1374 }
1375
1376 #[test]
1377 fn bookmarks_win_over_numbering_and_promote_list_items() {
1378 let outline = vec![
1379 OutlineItem {
1380 title: "1. Introduction".into(),
1381 level: 0,
1382 page_no: Some(1),
1383 y_top: None,
1384 },
1385 OutlineItem {
1386 title: "Hidden Heading".into(),
1387 level: 1,
1388 page_no: Some(1),
1389 y_top: None,
1390 },
1391 ];
1392 let mut nodes = vec![
1393 page(1),
1394 heading([10, 10, 200, 20], "1. Introduction"),
1397 Node::ListItem {
1398 ordered: false,
1399 number: 0,
1400 first_in_list: true,
1401 text: "Hidden Heading".into(),
1402 level: 0,
1403 marker: None,
1404 location: Some([10, 40, 200, 50]),
1405 dclx: None,
1406 href: None,
1407 layer: None,
1408 },
1409 Node::ListItem {
1410 ordered: false,
1411 number: 0,
1412 first_in_list: false,
1413 text: "a real item".into(),
1414 level: 0,
1415 marker: None,
1416 location: Some([10, 70, 200, 80]),
1417 dclx: None,
1418 href: None,
1419 layer: None,
1420 },
1421 ];
1422 apply(
1423 &mut nodes,
1424 &outline,
1425 &HashMap::new(),
1426 &HeadingHierarchyOptions::enabled(true),
1427 );
1428 assert_eq!(levels(&nodes), vec![2, 3]);
1431 match &nodes[3] {
1432 Node::ListItem {
1433 first_in_list,
1434 text,
1435 ..
1436 } => {
1437 assert!(*first_in_list, "sibling re-opens the list");
1438 assert_eq!(text, "a real item");
1439 }
1440 other => panic!("expected the sibling list item, got {other:?}"),
1441 }
1442 }
1443
1444 #[test]
1445 fn style_ranks_by_size_then_prominence() {
1446 let glyphs = vec![
1448 GlyphStyle {
1449 l: 10.0,
1450 t: 10.0,
1451 r: 100.0,
1452 b: 28.0,
1453 height: 18.0,
1454 weight_cls: 2,
1455 italic: false,
1456 styled: true,
1457 },
1458 GlyphStyle {
1459 l: 10.0,
1460 t: 60.0,
1461 r: 100.0,
1462 b: 72.0,
1463 height: 12.0,
1464 weight_cls: 2,
1465 italic: false,
1466 styled: true,
1467 },
1468 GlyphStyle {
1469 l: 10.0,
1470 t: 110.0,
1471 r: 100.0,
1472 b: 122.0,
1473 height: 12.0,
1474 weight_cls: 0,
1475 italic: false,
1476 styled: true,
1477 },
1478 ];
1479 let mut styles = HashMap::new();
1480 styles.insert(1usize, glyphs);
1481 let mut nodes = vec![
1482 page(1),
1483 heading([10, 10, 200, 28], "Big Title Words"),
1484 heading([10, 60, 200, 72], "Bold Twelve"),
1485 heading([10, 110, 200, 122], "Plain Twelve"),
1486 ];
1487 apply(
1489 &mut nodes,
1490 &[],
1491 &styles,
1492 &HeadingHierarchyOptions::enabled(true),
1493 );
1494 assert_eq!(levels(&nodes), vec![2, 3, 4]);
1495 }
1496
1497 #[test]
1498 fn strip_marker_removes_leading_numbering() {
1499 assert_eq!(strip_marker("1.1 Definitions"), "Definitions");
1500 assert_eq!(strip_marker("ARTICLE 5 - Payment"), "Payment");
1501 assert_eq!(strip_marker("(a) item"), "item");
1502 assert_eq!(strip_marker("No marker here"), "No marker here");
1503 }
1504}