rpdfium-doc 7676.6.4

// Derived from PDFium's cpdf_structtree.h/cpp
// Original: Copyright 2014 The PDFium Authors
// Licensed under BSD-3-Clause / Apache-2.0
// See pdfium-upstream/LICENSE for the original license.

//! Tagged PDF structure tree — `CPDF_StructTree`.
//!
//! Provides parsing of the PDF logical structure tree (ISO 32000-2 section 14.7),
//! which encodes accessibility (tagged PDF) information. The structure tree maps
//! marked content sequences in page content streams to semantic tags such as
//! paragraphs, headings, tables, and figures.
//!
//! All tree traversal is **iterative** (explicit `Vec` stack) for WASM safety.
//! Structure element types are in [`crate::struct_element`] (corresponds to
//! `CPDF_StructElement`).

use std::collections::HashMap;

use rpdfium_core::{Name, PdfSource};
use rpdfium_parser::{Object, ObjectId, ObjectStore};

use crate::error::{DocError, DocResult};
use crate::struct_element::{StructElement, parse_struct_element};

pub use crate::struct_element::{AttributeValue, StructAttribute};

/// Maximum number of structure elements to parse (safety limit).
const MAX_ELEMENTS: usize = 100_000;

/// Maximum depth for structure tree traversal.
const MAX_TREE_DEPTH: usize = 64;

/// Standard PDF structure type constants (ISO 32000-2 Table 368 ff.).
pub mod structure_types {
    pub const DOCUMENT: &str = "Document";
    pub const PART: &str = "Part";
    pub const SECT: &str = "Sect";
    pub const P: &str = "P";
    pub const H: &str = "H";
    pub const H1: &str = "H1";
    pub const H2: &str = "H2";
    pub const H3: &str = "H3";
    pub const H4: &str = "H4";
    pub const H5: &str = "H5";
    pub const H6: &str = "H6";
    pub const TABLE: &str = "Table";
    pub const TR: &str = "TR";
    pub const TD: &str = "TD";
    pub const TH: &str = "TH";
    pub const FIGURE: &str = "Figure";
    pub const SPAN: &str = "Span";
    pub const LINK: &str = "Link";
    pub const LIST: &str = "L";
    pub const LIST_ITEM: &str = "LI";
    pub const LABEL: &str = "Lbl";
    pub const LIST_BODY: &str = "LBody";
}

/// The parsed structure tree of a tagged PDF document.
#[derive(Debug, Clone)]
pub struct StructTree {
    /// Root structure elements.
    pub root_elements: Vec<StructElement>,
    /// Role mapping: custom tag name to standard tag name.
    pub role_map: HashMap<String, String>,
}

/// Returns indices of structure elements in the slice that reference the given MCID.
///
/// Performs an iterative depth-first search across all elements and their `/K` kids.
/// An element "references" an MCID if it contains that value in its `mcids` list.
///
/// The returned indices are positions in the top-level `elements` slice.  If a
/// descendant of element *i* holds the MCID, index *i* is included in the result.
///
/// # Examples
///
/// ```ignore
/// let result = find_elements_for_mcid(&tree.root_elements, 42);
/// // result contains indices of root elements whose subtrees include MCID 42
/// ```
pub fn find_elements_for_mcid(elements: &[StructElement], mcid: i32) -> Vec<usize> {
    let target = mcid as i64;
    let mut result = Vec::new();
    // Stack holds (root_index, element_ref) pairs.
    // Children are pushed with their ancestor's root_index so the caller
    // can identify which top-level element's subtree contains the MCID.
    let mut stack: Vec<(usize, &StructElement)> = elements.iter().enumerate().rev().collect();

    while let Some((idx, elem)) = stack.pop() {
        if elem.mcids.contains(&target) {
            result.push(idx);
        }
        // Push children with the same root index for DFS.
        for child in elem.children.iter().rev() {
            stack.push((idx, child));
        }
    }
    result
}

impl StructTree {
    /// Parse the structure tree from the document catalog dictionary.
    ///
    /// Returns `Ok(None)` if the catalog has no `/StructTreeRoot` entry.
    /// All traversal is iterative using an explicit stack.
    pub fn from_catalog<S: PdfSource>(
        catalog_dict: &HashMap<Name, Object>,
        store: &ObjectStore<S>,
    ) -> DocResult<Option<Self>> {
        // Look up /StructTreeRoot
        let root_obj = match catalog_dict.get(&Name::struct_tree_root()) {
            Some(obj) => store
                .deep_resolve(obj)
                .map_err(|e| DocError::Parser(e.to_string()))?,
            None => return Ok(None),
        };

        let root_dict = match root_obj.as_dict() {
            Some(d) => d,
            None => return Ok(None),
        };

        // Parse /RoleMap
        let role_map = parse_role_map(root_dict, store);

        // Parse /K (kids) — root content elements
        let root_elements = match root_dict.get(&Name::k()) {
            Some(k_obj) => parse_k_children(k_obj, store)?,
            None => Vec::new(),
        };

        Ok(Some(StructTree {
            root_elements,
            role_map,
        }))
    }

    /// Resolve a custom structure type to its standard type via the role map.
    ///
    /// Returns the mapped standard type if a mapping exists, or the original
    /// type name if no mapping is defined.
    ///
    /// Corresponds to `CPDF_StructTree::GetRoleMapNameFor()` in PDFium.
    pub fn role_map_name_for<'a>(&'a self, struct_type: &'a str) -> &'a str {
        self.role_map
            .get(struct_type)
            .map(|s| s.as_str())
            .unwrap_or(struct_type)
    }

    /// Upstream-aligned alias for [`role_map_name_for()`](StructTree::role_map_name_for).
    ///
    /// Corresponds to `CPDF_StructTree::GetRoleMapNameFor()` in PDFium.
    #[inline]
    pub fn get_role_map_name_for<'a>(&'a self, struct_type: &'a str) -> &'a str {
        self.role_map_name_for(struct_type)
    }

    /// Returns the number of top-level children in the structure tree.
    ///
    /// Corresponds to upstream `FPDF_StructTree_CountChildren`.
    pub fn child_count(&self) -> usize {
        self.root_elements.len()
    }

    /// ADR-019 Tier 2 alias for [`child_count()`](StructTree::child_count).
    ///
    /// Corresponds to upstream `FPDF_StructTree_CountChildren`.
    #[inline]
    pub fn struct_tree_count_children(&self) -> usize {
        self.child_count()
    }

    /// Deprecated short alias — use [`struct_tree_count_children()`](StructTree::struct_tree_count_children)
    /// or [`child_count()`](StructTree::child_count) instead.
    #[deprecated(
        since = "0.1.0",
        note = "use `struct_tree_count_children()` — matches upstream `FPDF_StructTree_CountChildren`"
    )]
    #[inline]
    pub fn count_children(&self) -> usize {
        self.child_count()
    }

    /// Returns the top-level child at the given zero-based index, or `None`.
    ///
    /// Corresponds to upstream `FPDF_StructTree_GetChildAtIndex`.
    pub fn child_at_index(&self, index: usize) -> Option<&StructElement> {
        self.root_elements.get(index)
    }

    /// ADR-019 Tier 2 alias for [`child_at_index()`](StructTree::child_at_index).
    ///
    /// Corresponds to upstream `FPDF_StructTree_GetChildAtIndex`.
    #[inline]
    pub fn struct_tree_get_child_at_index(&self, index: usize) -> Option<&StructElement> {
        self.child_at_index(index)
    }

    /// Deprecated short alias — use [`struct_tree_get_child_at_index()`](StructTree::struct_tree_get_child_at_index)
    /// or [`child_at_index()`](StructTree::child_at_index) instead.
    #[deprecated(
        since = "0.1.0",
        note = "use `struct_tree_get_child_at_index()` — matches upstream `FPDF_StructTree_GetChildAtIndex`"
    )]
    #[inline]
    pub fn get_child_at_index(&self, index: usize) -> Option<&StructElement> {
        self.child_at_index(index)
    }

    /// Iterate over all structure elements that reference the given page object ID.
    ///
    /// Equivalent to upstream `CPDF_StructTree::LoadPage`, which filters the
    /// structure tree to elements that belong to a specific page.  This variant
    /// takes the page's indirect-object ID (the `/Pg` reference value stored in
    /// each structure element) and performs an iterative depth-first walk of the
    /// entire tree, yielding every element whose `page_ref` matches.
    ///
    /// Use [`PageStructure::for_page`] if you need an owned collection instead.
    pub fn elements_for_page_ref(&self, page_ref: ObjectId) -> ElementsForPage<'_> {
        ElementsForPage {
            stack: self.root_elements.iter().rev().collect(),
            page_ref,
        }
    }

    /// Returns the indices (into `root_elements`) of elements whose subtrees reference
    /// the given marked content ID.
    ///
    /// Equivalent to calling [`find_elements_for_mcid`] with `&self.root_elements`.
    pub fn elements_for_mcid(&self, mcid: i32) -> Vec<usize> {
        find_elements_for_mcid(&self.root_elements, mcid)
    }

    /// Non-upstream alias — use [`elements_for_mcid()`](Self::elements_for_mcid).
    #[deprecated(
        note = "use `elements_for_mcid()` — no public `FPDF_StructTree_GetElementsForMcid` API"
    )]
    #[inline]
    pub fn get_elements_for_mcid(&self, mcid: i32) -> Vec<usize> {
        self.elements_for_mcid(mcid)
    }

    /// Iterate over all structure elements associated with the page at the given
    /// zero-based index, resolved against the provided ordered list of page IDs.
    ///
    /// `page_ids` must be ordered so that `page_ids[page_index]` is the indirect
    /// object ID of the page at that index — the same ordering as returned by
    /// `collect_page_ids` in `rpdfium-page`.
    ///
    /// Returns an empty iterator if `page_index` is out of range.
    pub fn elements_for_page<'a>(
        &'a self,
        page_index: usize,
        page_ids: &[ObjectId],
    ) -> ElementsForPage<'a> {
        match page_ids.get(page_index) {
            Some(&page_ref) => ElementsForPage {
                stack: self.root_elements.iter().rev().collect(),
                page_ref,
            },
            None => ElementsForPage {
                stack: Vec::new(),
                page_ref: ObjectId::new(0, 0),
            },
        }
    }
}

/// An iterator over structure elements that reference a specific page.
///
/// Created by [`StructTree::elements_for_page_ref`] and [`StructTree::elements_for_page`].
/// Traverses the structure tree iteratively (no recursion) in depth-first order,
/// yielding every [`StructElement`] whose `page_ref` field matches the target page.
pub struct ElementsForPage<'a> {
    /// DFS traversal stack holding remaining unvisited elements.
    stack: Vec<&'a StructElement>,
    /// The page object ID we are filtering for.
    page_ref: ObjectId,
}

impl<'a> Iterator for ElementsForPage<'a> {
    type Item = &'a StructElement;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let elem = self.stack.pop()?;
            // Push children in reverse order for correct left-to-right visitation.
            for child in elem.children.iter().rev() {
                self.stack.push(child);
            }
            if elem.page_ref == Some(self.page_ref) {
                return Some(elem);
            }
        }
    }
}

/// Parse the `/RoleMap` dictionary: maps custom tag names to standard tag names.
fn parse_role_map<S: PdfSource>(
    root_dict: &HashMap<Name, Object>,
    store: &ObjectStore<S>,
) -> HashMap<String, String> {
    let mut map = HashMap::new();
    let role_map_obj = match root_dict.get(&Name::role_map()) {
        Some(obj) => match store.deep_resolve(obj) {
            Ok(resolved) => resolved,
            Err(_) => return map,
        },
        None => return map,
    };

    if let Some(dict) = role_map_obj.as_dict() {
        for (key, value) in dict {
            if let Some(target) = value.as_name() {
                map.insert(key.as_str().into_owned(), target.as_str().into_owned());
            }
        }
    }
    map
}

/// Parse `/K` children, which can be a single element or an array.
/// Returns a Vec of root-level StructElements built via iterative traversal.
fn parse_k_children<S: PdfSource>(
    k_obj: &Object,
    store: &ObjectStore<S>,
) -> DocResult<Vec<StructElement>> {
    let resolved = store
        .deep_resolve(k_obj)
        .map_err(|e| DocError::Parser(e.to_string()))?;

    // Collect top-level items from /K
    let top_items: Vec<&Object> = match resolved {
        Object::Array(arr) => arr.iter().collect(),
        _ => vec![resolved],
    };

    // Each top-level item is either a struct element dict or an MCID integer.
    // We parse struct element dicts iteratively using a stack.
    // Result: flat list of (depth, StructElement) that we reconstruct into a tree.
    let mut flat: Vec<(usize, StructElement)> = Vec::new();

    // Stack entry for iterative DFS
    struct StackEntry<'a> {
        obj: &'a Object,
        depth: usize,
    }

    // Push top-level items in reverse order for correct left-to-right processing
    let mut stack: Vec<StackEntry<'_>> = Vec::new();
    for item in top_items.iter().rev() {
        stack.push(StackEntry {
            obj: item,
            depth: 0,
        });
    }

    while let Some(entry) = stack.pop() {
        if flat.len() >= MAX_ELEMENTS {
            break;
        }
        if entry.depth > MAX_TREE_DEPTH {
            return Err(DocError::DepthExceeded);
        }

        let resolved = match store.deep_resolve(entry.obj) {
            Ok(r) => r,
            Err(_) => continue,
        };

        // If it's an integer, it's a bare MCID at the top level — skip (no struct element)
        if resolved.as_i64().is_some() {
            continue;
        }

        let dict = match resolved.as_dict() {
            Some(d) => d,
            None => continue,
        };

        // Parse the structure element from this dictionary
        let mut elem = parse_struct_element(dict, store);

        // Process /K children: push child dicts onto stack, collect MCIDs into elem
        if let Some(k_val) = dict.get(&Name::k()) {
            let k_resolved = match store.deep_resolve(k_val) {
                Ok(r) => r,
                Err(_) => {
                    flat.push((entry.depth, elem));
                    continue;
                }
            };

            match k_resolved {
                Object::Integer(n) => {
                    elem.mcids.push(*n);
                }
                Object::Dictionary(child_dict) => {
                    // Could be a struct element dict or an MCID reference dict
                    if let Some(mcid) = extract_mcid_from_dict(child_dict) {
                        elem.mcids.push(mcid);
                    } else {
                        // It's a child struct element — push onto stack
                        stack.push(StackEntry {
                            obj: k_val,
                            depth: entry.depth + 1,
                        });
                    }
                }
                Object::Array(arr) => {
                    // Mixed array of MCIDs, MCID dicts, and child struct element dicts
                    // Process in reverse order for correct left-to-right stack processing
                    for child in arr.iter().rev() {
                        let child_resolved = match store.deep_resolve(child) {
                            Ok(r) => r,
                            Err(_) => continue,
                        };
                        match child_resolved {
                            Object::Integer(n) => {
                                elem.mcids.push(*n);
                            }
                            Object::Dictionary(child_dict) => {
                                if let Some(mcid) = extract_mcid_from_dict(child_dict) {
                                    elem.mcids.push(mcid);
                                } else {
                                    stack.push(StackEntry {
                                        obj: child,
                                        depth: entry.depth + 1,
                                    });
                                }
                            }
                            _ => {}
                        }
                    }
                    // MCIDs were pushed in reverse; restore original order
                    elem.mcids.reverse();
                }
                _ => {}
            }
        }

        flat.push((entry.depth, elem));
    }

    build_tree_from_flat(flat)
}

/// Extract MCID from a marked content reference dictionary (has `/Type /MCR` or just `/MCID`).
fn extract_mcid_from_dict(dict: &HashMap<Name, Object>) -> Option<i64> {
    dict.get(&Name::mcid()).and_then(|obj| obj.as_i64())
}

/// Reconstruct a tree from a flat depth-tagged list of elements.
/// Uses the same index-path approach as the bookmark tree builder.
fn build_tree_from_flat(flat: Vec<(usize, StructElement)>) -> DocResult<Vec<StructElement>> {
    if flat.is_empty() {
        return Ok(Vec::new());
    }

    let mut root: Vec<StructElement> = Vec::new();
    let mut path: Vec<usize> = Vec::new();

    for (depth, mut elem) in flat {
        path.truncate(depth);
        let container = get_children_at_path(&mut root, &path);
        let idx = container.len();
        // Root elements (depth == 0) have no parent; children record their position
        // within their parent's `children` vec so callers can implement GetParent.
        if depth > 0 {
            elem.parent_index = Some(idx);
        }
        container.push(elem);
        if path.len() <= depth {
            path.push(idx);
        }
    }

    Ok(root)
}

/// Iteratively navigate to the children vec at the given index path.
fn get_children_at_path<'a>(
    root: &'a mut Vec<StructElement>,
    path: &[usize],
) -> &'a mut Vec<StructElement> {
    let mut current = root;
    for &idx in path {
        current = &mut current[idx].children;
    }
    current
}

/// Mapping from (page ObjectId, MCID) to structure element information.
///
/// Built by walking the structure tree and collecting all MCID associations.
#[derive(Debug, Clone)]
pub struct McidMapping {
    /// Map from (page_object_id, mcid) to index into `elements`.
    entries: HashMap<(ObjectId, i64), usize>,
    /// Flat list of elements referenced by the mapping.
    elements: Vec<StructElement>,
}

impl McidMapping {
    /// Build an MCID mapping by iteratively walking the structure tree.
    pub fn from_struct_tree(tree: &StructTree) -> Self {
        let mut entries = HashMap::new();
        let mut elements = Vec::new();

        // Iterative walk of the tree
        let mut stack: Vec<&StructElement> = tree.root_elements.iter().rev().collect();

        while let Some(elem) = stack.pop() {
            if !elem.mcids.is_empty() {
                if let Some(page_id) = elem.page_ref {
                    let idx = elements.len();
                    elements.push(elem.clone());
                    for &mcid in &elem.mcids {
                        entries.insert((page_id, mcid), idx);
                    }
                }
            }
            // Push children in reverse for left-to-right processing
            for child in elem.children.iter().rev() {
                stack.push(child);
            }
        }

        McidMapping { entries, elements }
    }

    /// Look up the structure element for a given page and MCID.
    pub fn element_for_mcid(&self, page_id: ObjectId, mcid: i64) -> Option<&StructElement> {
        self.entries
            .get(&(page_id, mcid))
            .map(|&idx| &self.elements[idx])
    }
}

/// Structure elements filtered for a specific page.
#[derive(Debug, Clone)]
pub struct PageStructure {
    /// Structure elements that reference this page.
    pub elements: Vec<StructElement>,
}

impl PageStructure {
    /// Filter the structure tree to elements referencing the given page.
    pub fn for_page(tree: &StructTree, page_id: ObjectId) -> Self {
        let mut elements = Vec::new();

        // Iterative walk
        let mut stack: Vec<&StructElement> = tree.root_elements.iter().rev().collect();

        while let Some(elem) = stack.pop() {
            if elem.page_ref == Some(page_id) {
                elements.push(elem.clone());
            }
            for child in elem.children.iter().rev() {
                stack.push(child);
            }
        }

        PageStructure { elements }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rpdfium_core::PdfString;

    fn build_store() -> ObjectStore<Vec<u8>> {
        let pdf = build_minimal_pdf();
        ObjectStore::open(pdf, rpdfium_core::ParsingMode::Lenient).unwrap()
    }

    fn build_minimal_pdf() -> Vec<u8> {
        let mut pdf = Vec::new();
        pdf.extend_from_slice(b"%PDF-1.4\n");
        let obj1_offset = pdf.len();
        pdf.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");
        let obj2_offset = pdf.len();
        pdf.extend_from_slice(b"2 0 obj\n<< /Type /Pages /Kids [] /Count 0 >>\nendobj\n");
        let xref_offset = pdf.len();
        pdf.extend_from_slice(b"xref\n0 3\n");
        pdf.extend_from_slice(b"0000000000 65535 f \r\n");
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj1_offset).as_bytes());
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj2_offset).as_bytes());
        pdf.extend_from_slice(b"trailer\n<< /Size 3 /Root 1 0 R >>\n");
        pdf.extend_from_slice(format!("startxref\n{}\n%%EOF", xref_offset).as_bytes());
        pdf
    }

    fn str_obj(s: &str) -> Object {
        Object::String(PdfString::from_bytes(s.as_bytes().to_vec()))
    }

    fn name_obj(s: &str) -> Object {
        Object::Name(Name::from(s))
    }

    /// Helper: build a struct element dict with a type tag.
    fn struct_elem_dict(tag: &str) -> HashMap<Name, Object> {
        let mut d = HashMap::new();
        d.insert(Name::s(), name_obj(tag));
        d
    }

    #[test]
    fn test_no_struct_tree_root_returns_none() {
        let store = build_store();
        let catalog = HashMap::new();
        let result = StructTree::from_catalog(&catalog, &store).unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_empty_struct_tree_root() {
        let store = build_store();
        let root_dict = HashMap::new();
        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));
        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert!(tree.root_elements.is_empty());
        assert!(tree.role_map.is_empty());
    }

    #[test]
    fn test_basic_structure_tree_document_with_paragraphs() {
        let store = build_store();

        // Two paragraph children
        let p1 = struct_elem_dict("P");
        let p2 = struct_elem_dict("P");

        // Document element with /K array of children
        let mut doc = struct_elem_dict("Document");
        doc.insert(
            Name::k(),
            Object::Array(vec![Object::Dictionary(p1), Object::Dictionary(p2)]),
        );

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(doc));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(tree.root_elements.len(), 1);
        assert_eq!(tree.root_elements[0].struct_type, "Document");
        assert_eq!(tree.root_elements[0].children.len(), 2);
        assert_eq!(tree.root_elements[0].children[0].struct_type, "P");
        assert_eq!(tree.root_elements[0].children[1].struct_type, "P");
    }

    #[test]
    fn test_mcid_from_integer_in_k() {
        let store = build_store();

        // A paragraph with a single MCID integer in /K
        let mut p = struct_elem_dict("P");
        p.insert(Name::k(), Object::Integer(42));
        p.insert(Name::pg(), Object::Reference(ObjectId::new(5, 0)));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(p));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(tree.root_elements.len(), 1);
        assert_eq!(tree.root_elements[0].mcids, vec![42]);
        assert_eq!(tree.root_elements[0].page_ref, Some(ObjectId::new(5, 0)));
    }

    #[test]
    fn test_mcid_from_dict_in_k() {
        let store = build_store();

        // A marked content reference dict with /MCID
        let mut mcr = HashMap::new();
        mcr.insert(Name::mcid(), Object::Integer(7));

        let mut p = struct_elem_dict("Span");
        p.insert(Name::k(), Object::Dictionary(mcr));
        p.insert(Name::pg(), Object::Reference(ObjectId::new(3, 0)));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(p));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(tree.root_elements[0].mcids, vec![7]);
    }

    #[test]
    fn test_alt_text_extraction() {
        let store = build_store();

        let mut fig = struct_elem_dict("Figure");
        fig.insert(Name::alt(), str_obj("A photo of a cat"));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(fig));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(
            tree.root_elements[0].alt_text.as_deref(),
            Some("A photo of a cat")
        );
    }

    #[test]
    fn test_nested_structure_elements() {
        let store = build_store();

        // Span inside P inside Document
        let span = struct_elem_dict("Span");
        let mut p = struct_elem_dict("P");
        p.insert(Name::k(), Object::Dictionary(span));

        let mut doc = struct_elem_dict("Document");
        doc.insert(Name::k(), Object::Dictionary(p));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(doc));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(tree.root_elements[0].struct_type, "Document");
        assert_eq!(tree.root_elements[0].children[0].struct_type, "P");
        assert_eq!(
            tree.root_elements[0].children[0].children[0].struct_type,
            "Span"
        );
    }

    #[test]
    fn test_role_mapping() {
        let store = build_store();

        // RoleMap: "MyTag" -> "P"
        let mut role_map_dict = HashMap::new();
        role_map_dict.insert(Name::from("MyTag"), name_obj("P"));
        role_map_dict.insert(Name::from("CustomH"), name_obj("H1"));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::role_map(), Object::Dictionary(role_map_dict));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(tree.role_map.get("MyTag"), Some(&"P".to_string()));
        assert_eq!(tree.role_map.get("CustomH"), Some(&"H1".to_string()));
    }

    #[test]
    fn test_mixed_k_content_dicts_and_integers() {
        let store = build_store();

        // MCID reference dict
        let mut mcr = HashMap::new();
        mcr.insert(Name::mcid(), Object::Integer(3));

        // Child struct element
        let child = struct_elem_dict("Span");

        // Parent with mixed /K: integer, MCR dict, struct element
        let mut p = struct_elem_dict("P");
        p.insert(Name::pg(), Object::Reference(ObjectId::new(10, 0)));
        p.insert(
            Name::k(),
            Object::Array(vec![
                Object::Integer(1),
                Object::Dictionary(mcr),
                Object::Dictionary(child),
            ]),
        );

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(p));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        let elem = &tree.root_elements[0];
        assert_eq!(elem.struct_type, "P");
        assert_eq!(elem.mcids, vec![1, 3]);
        assert_eq!(elem.children.len(), 1);
        assert_eq!(elem.children[0].struct_type, "Span");
    }

    #[test]
    fn test_page_structure_filtering() {
        let store = build_store();

        let page1 = ObjectId::new(5, 0);
        let page2 = ObjectId::new(6, 0);

        let mut p1 = struct_elem_dict("P");
        p1.insert(Name::pg(), Object::Reference(page1));
        let mut p2 = struct_elem_dict("P");
        p2.insert(Name::pg(), Object::Reference(page2));
        let mut p3 = struct_elem_dict("P");
        p3.insert(Name::pg(), Object::Reference(page1));

        let mut doc = struct_elem_dict("Document");
        doc.insert(
            Name::k(),
            Object::Array(vec![
                Object::Dictionary(p1),
                Object::Dictionary(p2),
                Object::Dictionary(p3),
            ]),
        );

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(doc));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        let page1_struct = PageStructure::for_page(&tree, page1);
        assert_eq!(page1_struct.elements.len(), 2);
        for elem in &page1_struct.elements {
            assert_eq!(elem.page_ref, Some(page1));
        }

        let page2_struct = PageStructure::for_page(&tree, page2);
        assert_eq!(page2_struct.elements.len(), 1);
    }

    #[test]
    fn test_mcid_mapping_lookup() {
        let store = build_store();

        let page_id = ObjectId::new(7, 0);

        let mut p = struct_elem_dict("P");
        p.insert(Name::pg(), Object::Reference(page_id));
        p.insert(
            Name::k(),
            Object::Array(vec![Object::Integer(0), Object::Integer(1)]),
        );

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(p));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        let mapping = McidMapping::from_struct_tree(&tree);

        let elem = mapping.element_for_mcid(page_id, 0).unwrap();
        assert_eq!(elem.struct_type, "P");
        let elem1 = mapping.element_for_mcid(page_id, 1).unwrap();
        assert_eq!(elem1.struct_type, "P");
        assert!(mapping.element_for_mcid(page_id, 99).is_none());
        assert!(mapping.element_for_mcid(ObjectId::new(999, 0), 0).is_none());
    }

    #[test]
    fn test_security_limit_truncates_large_tree() {
        let store = build_store();

        // Build an array with MAX_ELEMENTS + 10 struct element dicts
        let count = MAX_ELEMENTS + 10;
        let arr: Vec<Object> = (0..count)
            .map(|_| Object::Dictionary(struct_elem_dict("P")))
            .collect();

        let mut doc = struct_elem_dict("Document");
        doc.insert(Name::k(), Object::Array(arr));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(doc));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        // The Document element itself counts as 1, and children are capped
        let total = count_elements(&tree.root_elements);
        assert!(total <= MAX_ELEMENTS + 1); // +1 for the Document root
    }

    /// Count total elements in a tree (iteratively).
    fn count_elements(roots: &[StructElement]) -> usize {
        let mut count = 0;
        let mut stack: Vec<&StructElement> = roots.iter().collect();
        while let Some(elem) = stack.pop() {
            count += 1;
            for child in &elem.children {
                stack.push(child);
            }
        }
        count
    }

    #[test]
    fn test_title_and_id_extraction() {
        let store = build_store();

        let mut elem = struct_elem_dict("Table");
        elem.insert(Name::t(), str_obj("Sales Data 2026"));
        elem.insert(Name::id(), str_obj("table-001"));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(elem));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(
            tree.root_elements[0].title.as_deref(),
            Some("Sales Data 2026")
        );
        assert_eq!(tree.root_elements[0].id.as_deref(), Some("table-001"));
    }

    #[test]
    fn test_role_map_name_for_lookup() {
        let store = build_store();

        let mut role_map_dict = HashMap::new();
        role_map_dict.insert(Name::from("MyTag"), name_obj("P"));
        role_map_dict.insert(Name::from("CustomH"), name_obj("H1"));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::role_map(), Object::Dictionary(role_map_dict));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(tree.role_map_name_for("MyTag"), "P");
        assert_eq!(tree.role_map_name_for("CustomH"), "H1");
        // Unknown tag returns itself
        assert_eq!(tree.role_map_name_for("P"), "P");
        assert_eq!(tree.role_map_name_for("UnknownTag"), "UnknownTag");
    }

    #[test]
    fn test_actual_text_and_lang() {
        let store = build_store();

        let mut span = struct_elem_dict("Span");
        span.insert(Name::actual_text(), str_obj("Hello World"));
        span.insert(Name::lang(), str_obj("en-US"));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(span));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(
            tree.root_elements[0].actual_text.as_deref(),
            Some("Hello World")
        );
        assert_eq!(tree.root_elements[0].lang.as_deref(), Some("en-US"));
    }

    #[test]
    fn test_struct_element_with_attributes() {
        let store = build_store();

        // Build an attribute dict with /O owner and some entries
        let mut attr_dict = HashMap::new();
        attr_dict.insert(Name::o(), name_obj("Layout"));
        attr_dict.insert(Name::from("WritingMode"), name_obj("LrTb"));
        attr_dict.insert(Name::from("SpaceBefore"), Object::Real(12.0));

        let mut td = struct_elem_dict("TD");
        td.insert(Name::a(), Object::Dictionary(attr_dict));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(td));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        let elem = &tree.root_elements[0];
        assert_eq!(elem.struct_type, "TD");
        assert_eq!(elem.attributes.len(), 1);
        assert_eq!(elem.attributes[0].owner, "Layout");
        assert!(elem.attributes[0].entries.len() >= 2);

        // Check that at least one entry has the expected value
        let writing_mode = elem.attributes[0]
            .entries
            .iter()
            .find(|(k, _)| k == "WritingMode");
        assert!(writing_mode.is_some());
        match &writing_mode.unwrap().1 {
            AttributeValue::Name(n) => assert_eq!(n, "LrTb"),
            _ => panic!("expected Name attribute value"),
        }
    }

    #[test]
    fn test_struct_element_with_attribute_array() {
        let store = build_store();

        let mut attr1 = HashMap::new();
        attr1.insert(Name::o(), name_obj("Layout"));
        attr1.insert(Name::from("TextAlign"), name_obj("Center"));

        let mut attr2 = HashMap::new();
        attr2.insert(Name::o(), name_obj("Table"));
        attr2.insert(Name::from("RowSpan"), Object::Integer(2));

        let mut td = struct_elem_dict("TD");
        td.insert(
            Name::a(),
            Object::Array(vec![Object::Dictionary(attr1), Object::Dictionary(attr2)]),
        );

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(td));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        let elem = &tree.root_elements[0];
        assert_eq!(elem.attributes.len(), 2);
        assert_eq!(elem.attributes[0].owner, "Layout");
        assert_eq!(elem.attributes[1].owner, "Table");
    }

    #[test]
    fn test_struct_element_no_attributes() {
        let store = build_store();

        let p = struct_elem_dict("P");

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(p));

        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert!(tree.root_elements[0].attributes.is_empty());
    }

    #[test]
    fn test_struct_element_obj_type_none_by_default() {
        let store = build_store();

        let p = struct_elem_dict("P");
        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(p));
        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert!(tree.root_elements[0].obj_type.is_none());
    }

    #[test]
    fn test_struct_element_obj_type_parsed() {
        let store = build_store();

        let mut elem = struct_elem_dict("Span");
        elem.insert(Name::obj_type(), Object::Name(Name::from("Elem")));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(elem));
        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();
        assert_eq!(tree.root_elements[0].obj_type.as_deref(), Some("Elem"));
    }

    // --- elements_for_page_ref / elements_for_page tests ---

    /// `elements_for_page_ref` returns only elements that reference the given page object ID.
    #[test]
    fn test_elements_for_page_ref_filters_correctly() {
        let store = build_store();

        let page1 = ObjectId::new(5, 0);
        let page2 = ObjectId::new(6, 0);

        // Three paragraphs: two on page1, one on page2.
        let mut p1 = struct_elem_dict("P");
        p1.insert(Name::pg(), Object::Reference(page1));
        let mut p2 = struct_elem_dict("H1");
        p2.insert(Name::pg(), Object::Reference(page2));
        let mut p3 = struct_elem_dict("Span");
        p3.insert(Name::pg(), Object::Reference(page1));

        let mut doc = struct_elem_dict("Document");
        doc.insert(
            Name::k(),
            Object::Array(vec![
                Object::Dictionary(p1),
                Object::Dictionary(p2),
                Object::Dictionary(p3),
            ]),
        );

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(doc));
        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();

        // Filter by page1 — should get "P" and "Span".
        let page1_elems: Vec<&StructElement> = tree.elements_for_page_ref(page1).collect();
        assert_eq!(page1_elems.len(), 2);
        assert!(page1_elems.iter().all(|e| e.page_ref == Some(page1)));
        let types: Vec<&str> = page1_elems.iter().map(|e| e.struct_type.as_str()).collect();
        assert!(types.contains(&"P"));
        assert!(types.contains(&"Span"));

        // Filter by page2 — should get only "H1".
        let page2_elems: Vec<&StructElement> = tree.elements_for_page_ref(page2).collect();
        assert_eq!(page2_elems.len(), 1);
        assert_eq!(page2_elems[0].struct_type, "H1");

        // Filter by a page not in the tree — should get nothing.
        let absent = ObjectId::new(99, 0);
        let absent_elems: Vec<&StructElement> = tree.elements_for_page_ref(absent).collect();
        assert!(absent_elems.is_empty());
    }

    /// `elements_for_page` with a valid page_index selects the correct page's elements.
    #[test]
    fn test_elements_for_page_with_valid_index() {
        let store = build_store();

        let page0 = ObjectId::new(10, 0);
        let page1 = ObjectId::new(11, 0);

        let mut h1 = struct_elem_dict("H1");
        h1.insert(Name::pg(), Object::Reference(page0));
        let mut p = struct_elem_dict("P");
        p.insert(Name::pg(), Object::Reference(page1));

        let mut doc = struct_elem_dict("Document");
        doc.insert(
            Name::k(),
            Object::Array(vec![Object::Dictionary(h1), Object::Dictionary(p)]),
        );

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(doc));
        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();

        // Simulate a page_ids list: page0 at index 0, page1 at index 1.
        let page_ids = vec![page0, page1];

        let idx0_elems: Vec<&StructElement> = tree.elements_for_page(0, &page_ids).collect();
        assert_eq!(idx0_elems.len(), 1);
        assert_eq!(idx0_elems[0].struct_type, "H1");

        let idx1_elems: Vec<&StructElement> = tree.elements_for_page(1, &page_ids).collect();
        assert_eq!(idx1_elems.len(), 1);
        assert_eq!(idx1_elems[0].struct_type, "P");
    }

    // --- find_elements_for_mcid tests ---

    /// `find_elements_for_mcid` on an empty slice returns empty vec.
    #[test]
    fn test_find_elements_for_mcid_empty() {
        let result = find_elements_for_mcid(&[], 42);
        assert!(result.is_empty());
    }

    /// `find_elements_for_mcid` returns the index of an element with a matching MCID.
    #[test]
    fn test_find_elements_for_mcid_found() {
        // Build two root elements: index 0 has MCID 5, index 1 has MCID 10.
        let elem0 = StructElement {
            struct_type: "P".to_string(),
            obj_type: None,
            alt_text: None,
            actual_text: None,
            lang: None,
            title: None,
            id: None,
            page_ref: None,
            mcids: vec![5],
            children: Vec::new(),
            attributes: Vec::new(),
            parent_index: None,
        };
        let elem1 = StructElement {
            struct_type: "Span".to_string(),
            obj_type: None,
            alt_text: None,
            actual_text: None,
            lang: None,
            title: None,
            id: None,
            page_ref: None,
            mcids: vec![10],
            children: Vec::new(),
            attributes: Vec::new(),
            parent_index: None,
        };
        let elements = vec![elem0, elem1];

        // MCID 5 is in element at index 0.
        let result = find_elements_for_mcid(&elements, 5);
        assert_eq!(result, vec![0]);

        // MCID 10 is in element at index 1.
        let result = find_elements_for_mcid(&elements, 10);
        assert_eq!(result, vec![1]);

        // MCID from a child element returns the root index.
        let child = StructElement {
            struct_type: "Span".to_string(),
            obj_type: None,
            alt_text: None,
            actual_text: None,
            lang: None,
            title: None,
            id: None,
            page_ref: None,
            mcids: vec![99],
            children: Vec::new(),
            attributes: Vec::new(),
            parent_index: None,
        };
        let parent = StructElement {
            struct_type: "P".to_string(),
            obj_type: None,
            alt_text: None,
            actual_text: None,
            lang: None,
            title: None,
            id: None,
            page_ref: None,
            mcids: Vec::new(),
            children: vec![child],
            attributes: Vec::new(),
            parent_index: None,
        };
        let result = find_elements_for_mcid(&[parent], 99);
        assert_eq!(result, vec![0]);
    }

    /// `find_elements_for_mcid` returns empty vec when the MCID is not present.
    #[test]
    fn test_find_elements_for_mcid_not_found() {
        let elem = StructElement {
            struct_type: "P".to_string(),
            obj_type: None,
            alt_text: None,
            actual_text: None,
            lang: None,
            title: None,
            id: None,
            page_ref: None,
            mcids: vec![1, 2, 3],
            children: Vec::new(),
            attributes: Vec::new(),
            parent_index: None,
        };
        let result = find_elements_for_mcid(&[elem], 999);
        assert!(result.is_empty());
    }

    /// `elements_for_page` with an out-of-range index returns an empty iterator.
    #[test]
    fn test_elements_for_page_out_of_range_returns_empty() {
        let store = build_store();

        let page0 = ObjectId::new(20, 0);

        let mut p = struct_elem_dict("P");
        p.insert(Name::pg(), Object::Reference(page0));

        let mut root_dict = HashMap::new();
        root_dict.insert(Name::k(), Object::Dictionary(p));
        let mut catalog = HashMap::new();
        catalog.insert(Name::struct_tree_root(), Object::Dictionary(root_dict));

        let tree = StructTree::from_catalog(&catalog, &store).unwrap().unwrap();

        let page_ids = vec![page0];
        // Index 5 is out of range for a single-page list.
        let elems: Vec<&StructElement> = tree.elements_for_page(5, &page_ids).collect();
        assert!(elems.is_empty());
    }
}