use std::collections::{HashMap, HashSet};
use zpdf_core::{ObjectId, PdfDict, PdfObject};
use zpdf_parser::PdfFile;
use crate::obj_util::{catalog_dict, resolve_dict, text};
use crate::Catalog;
const MAX_STRUCT_DEPTH: usize = 64;
const MAX_STRUCT_ELEMENTS: usize = 500_000;
const MAX_ROLE_MAP_DEPTH: usize = 32;
const MAX_ROLE_MAP_ENTRIES: usize = 65_536;
const MAX_TEXT_CHARS: usize = 64 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StructRole {
Document,
Part,
Art,
Sect,
Div,
BlockQuote,
Caption,
Toc,
Toci,
Index,
NonStruct,
Private,
P,
H,
H1,
H2,
H3,
H4,
H5,
H6,
L,
Li,
Lbl,
LBody,
Table,
Tr,
Th,
Td,
THead,
TBody,
TFoot,
Span,
Quote,
Note,
Reference,
BibEntry,
Code,
Link,
Annot,
Ruby,
Rb,
Rt,
Rp,
Warichu,
Wt,
Wp,
Figure,
Formula,
Form,
Other(String),
}
impl StructRole {
fn from_name(name: &str) -> Self {
use StructRole::*;
match name {
"Document" => Document,
"Part" => Part,
"Art" => Art,
"Sect" => Sect,
"Div" => Div,
"BlockQuote" => BlockQuote,
"Caption" => Caption,
"TOC" => Toc,
"TOCI" => Toci,
"Index" => Index,
"NonStruct" => NonStruct,
"Private" => Private,
"P" => P,
"H" => H,
"H1" => H1,
"H2" => H2,
"H3" => H3,
"H4" => H4,
"H5" => H5,
"H6" => H6,
"L" => L,
"LI" => Li,
"Lbl" => Lbl,
"LBody" => LBody,
"Table" => Table,
"TR" => Tr,
"TH" => Th,
"TD" => Td,
"THead" => THead,
"TBody" => TBody,
"TFoot" => TFoot,
"Span" => Span,
"Quote" => Quote,
"Note" => Note,
"Reference" => Reference,
"BibEntry" => BibEntry,
"Code" => Code,
"Link" => Link,
"Annot" => Annot,
"Ruby" => Ruby,
"RB" => Rb,
"RT" => Rt,
"RP" => Rp,
"Warichu" => Warichu,
"WT" => Wt,
"WP" => Wp,
"Figure" => Figure,
"Formula" => Formula,
"Form" => Form,
other => Other(other.to_string()),
}
}
pub fn as_str(&self) -> &str {
use StructRole::*;
match self {
Document => "Document",
Part => "Part",
Art => "Art",
Sect => "Sect",
Div => "Div",
BlockQuote => "BlockQuote",
Caption => "Caption",
Toc => "TOC",
Toci => "TOCI",
Index => "Index",
NonStruct => "NonStruct",
Private => "Private",
P => "P",
H => "H",
H1 => "H1",
H2 => "H2",
H3 => "H3",
H4 => "H4",
H5 => "H5",
H6 => "H6",
L => "L",
Li => "LI",
Lbl => "Lbl",
LBody => "LBody",
Table => "Table",
Tr => "TR",
Th => "TH",
Td => "TD",
THead => "THead",
TBody => "TBody",
TFoot => "TFoot",
Span => "Span",
Quote => "Quote",
Note => "Note",
Reference => "Reference",
BibEntry => "BibEntry",
Code => "Code",
Link => "Link",
Annot => "Annot",
Ruby => "Ruby",
Rb => "RB",
Rt => "RT",
Rp => "RP",
Warichu => "Warichu",
Wt => "WT",
Wp => "WP",
Figure => "Figure",
Formula => "Formula",
Form => "Form",
Other(s) => s,
}
}
pub fn is_standard(&self) -> bool {
!matches!(self, StructRole::Other(_))
}
pub fn is_heading(&self) -> bool {
use StructRole::*;
matches!(self, H | H1 | H2 | H3 | H4 | H5 | H6)
}
pub fn is_block_level(&self) -> bool {
use StructRole::*;
matches!(
self,
Document
| Part
| Art
| Sect
| Div
| BlockQuote
| Caption
| Toc
| Toci
| Index
| P
| H
| H1
| H2
| H3
| H4
| H5
| H6
| L
| Li
| Table
| Tr
| Note
| Figure
| Formula
)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StructKid {
Element(StructElem),
MarkedContent {
page: Option<usize>,
mcid: i64,
},
Object {
page: Option<usize>,
obj: ObjectId,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StructElem {
pub role: StructRole,
pub raw_type: String,
pub title: Option<String>,
pub lang: Option<String>,
pub alt: Option<String>,
pub actual_text: Option<String>,
pub expansion: Option<String>,
pub page: Option<usize>,
pub kids: Vec<StructKid>,
}
impl StructElem {
pub fn accessible_text(&self) -> Option<&str> {
self.actual_text.as_deref().or(self.alt.as_deref())
}
pub fn child_elements(&self) -> impl Iterator<Item = &StructElem> {
self.kids.iter().filter_map(|k| match k {
StructKid::Element(e) => Some(e),
_ => None,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StructTree {
pub children: Vec<StructElem>,
pub marked: bool,
}
impl StructTree {
pub fn element_count(&self) -> usize {
fn count(e: &StructElem) -> usize {
1 + e.child_elements().map(count).sum::<usize>()
}
self.children.iter().map(count).sum()
}
}
pub fn is_tagged(file: &PdfFile) -> bool {
let Some(root) = catalog_dict(file) else {
return false;
};
let Some(mark_info) = resolve_dict(file, root.get("MarkInfo")) else {
return false;
};
matches!(mark_info.get("Marked"), Some(PdfObject::Bool(true)))
}
pub fn parse_struct_tree(file: &PdfFile, catalog: &Catalog) -> Option<StructTree> {
let root = catalog_dict(file)?;
let tree_root = resolve_dict(file, root.get("StructTreeRoot"))?;
let mut visited = HashSet::new();
if let Some(PdfObject::Ref(id)) = root.get("StructTreeRoot") {
visited.insert(*id);
}
let mut walk = StructWalk {
file,
catalog,
role_map: read_role_map(file, &tree_root),
visited,
budget: MAX_STRUCT_ELEMENTS,
};
let root_page = walk.page_of(&tree_root);
let mut children = Vec::new();
for kid in normalize_kids(file, &tree_root) {
if let Some(StructKid::Element(e)) = walk.parse_kid(&kid, root_page, 0) {
children.push(e);
}
}
Some(StructTree {
children,
marked: is_tagged(file),
})
}
struct StructWalk<'a> {
file: &'a PdfFile,
catalog: &'a Catalog,
role_map: HashMap<String, String>,
visited: HashSet<ObjectId>,
budget: usize,
}
impl StructWalk<'_> {
fn parse_kid(
&mut self,
obj: &PdfObject,
parent_page: Option<usize>,
depth: usize,
) -> Option<StructKid> {
if self.budget == 0 {
return None;
}
self.budget -= 1;
match obj {
PdfObject::Integer(mcid) => Some(StructKid::MarkedContent {
page: parent_page,
mcid: *mcid,
}),
PdfObject::Ref(id) => {
let resolved = self.file.resolve(*id).ok()?;
let dict = resolved.as_dict().ok()?;
match kid_dict_kind(dict) {
KidKind::Mcr => self.marked_content(dict, parent_page),
KidKind::Objr => self.object_ref(dict, parent_page),
KidKind::Element => {
if !self.visited.insert(*id) {
return None;
}
self.element(dict, parent_page, depth)
.map(StructKid::Element)
}
}
}
PdfObject::Dict(dict) => match kid_dict_kind(dict) {
KidKind::Mcr => self.marked_content(dict, parent_page),
KidKind::Objr => self.object_ref(dict, parent_page),
KidKind::Element => self
.element(dict, parent_page, depth)
.map(StructKid::Element),
},
_ => None,
}
}
fn element(
&mut self,
dict: &PdfDict,
parent_page: Option<usize>,
depth: usize,
) -> Option<StructElem> {
if depth > MAX_STRUCT_DEPTH {
return None;
}
let page = self.page_of(dict).or(parent_page);
let raw_type = self.file_name(dict, "S").unwrap_or_default();
let role = StructRole::from_name(&self.resolve_role(&raw_type));
let kids = normalize_kids(self.file, dict)
.iter()
.filter_map(|k| self.parse_kid(k, page, depth + 1))
.collect();
Some(StructElem {
role,
raw_type,
title: capped_text(self.file, dict, "T"),
lang: capped_text(self.file, dict, "Lang"),
alt: capped_text(self.file, dict, "Alt"),
actual_text: capped_text(self.file, dict, "ActualText"),
expansion: capped_text(self.file, dict, "E"),
page,
kids,
})
}
fn marked_content(&self, dict: &PdfDict, parent_page: Option<usize>) -> Option<StructKid> {
let mcid = int_value(dict.get("MCID"))?;
let page = self.page_of(dict).or(parent_page);
Some(StructKid::MarkedContent { page, mcid })
}
fn object_ref(&self, dict: &PdfDict, parent_page: Option<usize>) -> Option<StructKid> {
let obj = dict.get_ref("Obj").ok()?;
let page = self.page_of(dict).or(parent_page);
Some(StructKid::Object { page, obj })
}
fn page_of(&self, dict: &PdfDict) -> Option<usize> {
let pg = dict.get_ref("Pg").ok()?;
self.catalog.page_index_of(pg)
}
fn file_name(&self, dict: &PdfDict, key: &str) -> Option<String> {
crate::obj_util::name_value(self.file, dict, key)
}
fn resolve_role(&self, raw: &str) -> String {
let mut current = raw.to_string();
let mut seen = HashSet::new();
for _ in 0..MAX_ROLE_MAP_DEPTH {
if !seen.insert(current.clone()) {
break;
}
match self.role_map.get(¤t) {
Some(next) if next != ¤t => current = next.clone(),
_ => break,
}
}
current
}
}
enum KidKind {
Mcr,
Objr,
Element,
}
fn kid_dict_kind(dict: &PdfDict) -> KidKind {
match dict.get_name("Type") {
Ok("MCR") => return KidKind::Mcr,
Ok("OBJR") => return KidKind::Objr,
Ok("StructElem") => return KidKind::Element,
_ => {}
}
if dict.get("S").is_none() {
if dict.get("MCID").is_some() {
return KidKind::Mcr;
}
if dict.get("Obj").is_some() {
return KidKind::Objr;
}
}
KidKind::Element
}
fn normalize_kids(file: &PdfFile, dict: &PdfDict) -> Vec<PdfObject> {
match dict.get("K") {
Some(PdfObject::Array(a)) => a.clone(),
Some(PdfObject::Ref(r)) => match file.resolve(*r) {
Ok(PdfObject::Array(a)) => a,
Ok(_) => vec![PdfObject::Ref(*r)],
Err(_) => Vec::new(),
},
Some(other) => vec![other.clone()],
None => Vec::new(),
}
}
fn read_role_map(file: &PdfFile, tree_root: &PdfDict) -> HashMap<String, String> {
let mut map = HashMap::new();
let Some(rm) = resolve_dict(file, tree_root.get("RoleMap")) else {
return map;
};
for (key, value) in rm.0.iter() {
if map.len() >= MAX_ROLE_MAP_ENTRIES {
break;
}
if let PdfObject::Name(n) = value {
map.insert(key.as_str().to_string(), n.as_str().to_string());
}
}
map
}
fn int_value(obj: Option<&PdfObject>) -> Option<i64> {
match obj? {
PdfObject::Integer(n) => Some(*n),
PdfObject::Real(f) if f.is_finite() && f.fract() == 0.0 => Some(*f as i64),
_ => None,
}
}
fn capped_text(file: &PdfFile, dict: &PdfDict, key: &str) -> Option<String> {
match text(file, dict, key) {
Some(s) if s.chars().count() > MAX_TEXT_CHARS => {
Some(s.chars().take(MAX_TEXT_CHARS).collect())
}
other => other,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_util::build_pdf;
use crate::PdfDocument;
const PAGES: &str = "<< /Type /Pages /Kids [3 0 R] /Count 1 >>";
const PAGE: &str = "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 100 100] >>";
fn open(objects: &[&str]) -> PdfDocument {
PdfDocument::open(build_pdf(objects)).expect("open pdf")
}
fn doc(catalog: &str, extra: &[&str]) -> PdfDocument {
let mut objs = vec![catalog, PAGES, PAGE];
objs.extend_from_slice(extra);
open(&objs)
}
#[test]
fn no_struct_tree_is_none() {
let d = doc("<< /Type /Catalog /Pages 2 0 R >>", &[]);
assert!(d.struct_tree().is_none());
assert!(!d.is_tagged());
}
#[test]
fn mark_info_marks_tagged() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /MarkInfo << /Marked true >> >>",
&[],
);
assert!(d.is_tagged());
assert!(d.struct_tree().is_none());
}
#[test]
fn simple_document_paragraph_with_mcids() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R \
/MarkInfo << /Marked true >> >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Document /P 4 0 R /K 6 0 R >>",
"<< /Type /StructElem /S /P /P 5 0 R /Pg 3 0 R /K [0 1] >>",
],
);
let tree = d.struct_tree().expect("tree");
assert!(tree.marked);
assert_eq!(tree.children.len(), 1);
let document = &tree.children[0];
assert_eq!(document.role, StructRole::Document);
assert_eq!(document.kids.len(), 1);
let para = document.child_elements().next().unwrap();
assert_eq!(para.role, StructRole::P);
assert_eq!(para.page, Some(0));
assert_eq!(
para.kids,
vec![
StructKid::MarkedContent {
page: Some(0),
mcid: 0
},
StructKid::MarkedContent {
page: Some(0),
mcid: 1
},
]
);
assert_eq!(tree.element_count(), 2);
}
#[test]
fn role_map_resolves_custom_type() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R /RoleMap << /Heading1 /H1 >> >>",
"<< /Type /StructElem /S /Heading1 /P 4 0 R >>",
],
);
let tree = d.struct_tree().expect("tree");
let h = &tree.children[0];
assert_eq!(h.role, StructRole::H1);
assert!(h.role.is_heading());
assert_eq!(h.raw_type, "Heading1"); }
#[test]
fn unmapped_custom_type_is_other() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /MyWidget /P 4 0 R >>",
],
);
let role = &d.struct_tree().unwrap().children[0].role;
assert_eq!(role, &StructRole::Other("MyWidget".to_string()));
assert!(!role.is_standard());
assert_eq!(role.as_str(), "MyWidget");
}
#[test]
fn figure_alt_text_is_accessible() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Figure /P 4 0 R /Alt (A bar chart) >>",
],
);
let fig = &d.struct_tree().unwrap().children[0];
assert_eq!(fig.role, StructRole::Figure);
assert_eq!(fig.alt.as_deref(), Some("A bar chart"));
assert_eq!(fig.accessible_text(), Some("A bar chart"));
}
#[test]
fn actual_text_preferred_over_alt() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Span /P 4 0 R /Alt (alt) /ActualText (exact) >>",
],
);
let span = &d.struct_tree().unwrap().children[0];
assert_eq!(span.accessible_text(), Some("exact"));
}
#[test]
fn objr_kid_resolves_object_and_page() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Link /P 4 0 R \
/K << /Type /OBJR /Obj 6 0 R /Pg 3 0 R >> >>",
"<< /Type /Annot /Subtype /Link >>",
],
);
let link = &d.struct_tree().unwrap().children[0];
assert_eq!(link.role, StructRole::Link);
assert_eq!(link.kids.len(), 1);
match &link.kids[0] {
StructKid::Object { page, obj } => {
assert_eq!(*page, Some(0));
assert_eq!(obj.0, 6); }
other => panic!("expected OBJR kid, got {other:?}"),
}
}
#[test]
fn mcr_dict_kid_with_explicit_page() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /P /P 4 0 R \
/K << /Type /MCR /Pg 3 0 R /MCID 7 >> >>",
],
);
let para = &d.struct_tree().unwrap().children[0];
assert_eq!(
para.kids[0],
StructKid::MarkedContent {
page: Some(0),
mcid: 7
}
);
}
#[test]
fn page_inherited_from_ancestor() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Sect /P 4 0 R /Pg 3 0 R /K 6 0 R >>",
"<< /Type /StructElem /S /Span /P 5 0 R /K [9] >>",
],
);
let span = d.struct_tree().unwrap().children[0]
.child_elements()
.next()
.unwrap()
.clone();
assert_eq!(span.page, Some(0), "inherited /Pg");
assert_eq!(
span.kids[0],
StructKid::MarkedContent {
page: Some(0),
mcid: 9
}
);
}
#[test]
fn single_ref_k_not_array() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Document /K 6 0 R >>",
"<< /Type /StructElem /S /P /P 5 0 R >>",
],
);
let document = &d.struct_tree().unwrap().children[0];
assert_eq!(document.child_elements().count(), 1);
assert_eq!(
document.child_elements().next().unwrap().role,
StructRole::P
);
}
#[test]
fn cyclic_kids_terminate() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Document /K 6 0 R >>",
"<< /Type /StructElem /S /Sect /K 5 0 R >>",
],
);
let tree = d.struct_tree().expect("tree (no hang)");
assert_eq!(tree.children.len(), 1);
let sect = tree.children[0].child_elements().next().unwrap();
assert_eq!(sect.role, StructRole::Sect);
assert_eq!(sect.child_elements().count(), 0);
}
#[test]
fn root_back_edge_makes_no_spurious_element() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
"<< /Type /StructElem /S /Document /K 4 0 R >>",
],
);
let document = &d.struct_tree().unwrap().children[0];
assert_eq!(document.role, StructRole::Document);
assert_eq!(document.child_elements().count(), 0, "root back-edge cut");
}
#[test]
fn role_map_cycle_terminates() {
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R /RoleMap << /Foo /Bar /Bar /Foo >> >>",
"<< /Type /StructElem /S /Foo /P 4 0 R >>",
],
);
let role = &d.struct_tree().expect("tree (no hang)").children[0].role;
assert!(!role.is_standard());
}
#[test]
fn deeply_nested_tree_terminates() {
let depth = MAX_STRUCT_DEPTH + 50;
let mut objs = vec![
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>".to_string(),
PAGES.to_string(),
PAGE.to_string(),
"<< /Type /StructTreeRoot /K 5 0 R >>".to_string(),
];
for i in 0..depth {
let obj_num = 5 + i;
if i + 1 < depth {
objs.push(format!(
"<< /Type /StructElem /S /Div /K {} 0 R >>",
obj_num + 1
));
} else {
objs.push("<< /Type /StructElem /S /Div >>".to_string());
}
}
let refs: Vec<&str> = objs.iter().map(|s| s.as_str()).collect();
let d = open(&refs);
assert!(d.struct_tree().is_some());
}
#[test]
fn huge_alt_text_is_capped() {
let big = "A".repeat(MAX_TEXT_CHARS + 1000);
let d = doc(
"<< /Type /Catalog /Pages 2 0 R /StructTreeRoot 4 0 R >>",
&[
"<< /Type /StructTreeRoot /K 5 0 R >>",
&format!("<< /Type /StructElem /S /Figure /Alt ({big}) >>"),
],
);
let alt = d.struct_tree().unwrap().children[0].alt.clone().unwrap();
assert_eq!(alt.chars().count(), MAX_TEXT_CHARS);
}
#[test]
fn role_name_round_trip() {
for name in [
"Document", "TOC", "TOCI", "P", "H1", "H6", "L", "LI", "Lbl", "LBody", "Table", "TR",
"TH", "TD", "THead", "TBody", "TFoot", "Span", "BibEntry", "Link", "RB", "WP",
"Figure", "Formula", "Form",
] {
let role = StructRole::from_name(name);
assert!(role.is_standard(), "{name} should be standard");
assert_eq!(role.as_str(), name, "round-trip for {name}");
}
}
}