rpdfium-doc 7676.6.2

//! PDF number tree parser (ISO 32000-2 section 7.9.7).
//!
//! A number tree maps integer keys to values. It has the same structure as
//! a name tree but uses `/Nums` instead of `/Names`.

use std::collections::HashMap;

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

use crate::error::{DocError, DocResult};

/// Maximum traversal depth for number trees.
const MAX_TREE_DEPTH: usize = 64;

/// A parsed number tree mapping integer keys to values of type `V`.
#[derive(Debug, Clone)]
pub struct NumberTree<V> {
    entries: Vec<(i64, V)>,
}

impl<V> NumberTree<V> {
    /// Parse a number tree starting from the given root object.
    ///
    /// The `convert` function transforms each value `Object` into a `V`.
    /// Traversal is fully iterative using an explicit stack.
    pub fn parse<S: PdfSource, F>(
        root: &Object,
        store: &ObjectStore<S>,
        convert: F,
    ) -> DocResult<Self>
    where
        F: Fn(&Object, &ObjectStore<S>) -> DocResult<V>,
    {
        let root_dict = resolve_dict(root, store)?;
        let mut entries = Vec::new();

        // Stack of (dict, depth) for iterative traversal
        let mut stack: Vec<(&HashMap<Name, Object>, usize)> = vec![(root_dict, 0)];

        while let Some((dict, depth)) = stack.pop() {
            if depth > MAX_TREE_DEPTH {
                return Err(DocError::DepthExceeded);
            }

            // Leaf node: has /Nums array
            if let Some(Object::Array(nums_arr)) = dict.get(&Name::nums()) {
                let mut i = 0;
                while i + 1 < nums_arr.len() {
                    let key_obj = store
                        .deep_resolve(&nums_arr[i])
                        .map_err(|e| DocError::Parser(e.to_string()))?;
                    let val_obj = store
                        .deep_resolve(&nums_arr[i + 1])
                        .map_err(|e| DocError::Parser(e.to_string()))?;

                    let key = key_obj.as_i64().ok_or(DocError::UnexpectedType)?;
                    let value = convert(val_obj, store)?;
                    entries.push((key, value));
                    i += 2;
                }
            }

            // Intermediate node: has /Kids array
            if let Some(Object::Array(kids)) = dict.get(&Name::kids()) {
                for kid in kids.iter().rev() {
                    let kid_obj = store
                        .deep_resolve(kid)
                        .map_err(|e| DocError::Parser(e.to_string()))?;
                    let kid_dict = kid_obj.as_dict().ok_or(DocError::UnexpectedType)?;
                    stack.push((kid_dict, depth + 1));
                }
            }
        }

        // Sort by key for binary search in floor
        entries.sort_by_key(|(k, _)| *k);

        Ok(Self { entries })
    }

    /// Look up a value by exact integer key.
    pub fn get(&self, num: i64) -> Option<&V> {
        self.entries
            .binary_search_by_key(&num, |(k, _)| *k)
            .ok()
            .map(|idx| &self.entries[idx].1)
    }

    /// Find the entry with the largest key less than or equal to `num`.
    ///
    /// Corresponds to `CPDF_NumberTree::GetFloor()` in PDFium.
    pub fn floor(&self, num: i64) -> Option<&V> {
        match self.entries.binary_search_by_key(&num, |(k, _)| *k) {
            Ok(idx) => Some(&self.entries[idx].1),
            Err(idx) => {
                if idx > 0 {
                    Some(&self.entries[idx - 1].1)
                } else {
                    None
                }
            }
        }
    }

    /// ADR-019 alias for [`floor()`](Self::floor).
    ///
    /// Corresponds to `CPDF_NumberTree::GetFloor()` in PDFium.
    #[inline]
    pub fn get_floor(&self, num: i64) -> Option<&V> {
        self.floor(num)
    }

    /// Return all entries as a slice of `(i64, V)` pairs, sorted by key.
    pub fn entries(&self) -> &[(i64, V)] {
        &self.entries
    }
}

/// Resolve an object to a dictionary reference.
fn resolve_dict<'a, S: PdfSource>(
    obj: &'a Object,
    store: &'a ObjectStore<S>,
) -> DocResult<&'a HashMap<Name, Object>> {
    let resolved = store
        .deep_resolve(obj)
        .map_err(|e| DocError::Parser(e.to_string()))?;
    resolved.as_dict().ok_or(DocError::UnexpectedType)
}

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

    fn int_obj(n: i64) -> Object {
        Object::Integer(n)
    }

    fn convert_int<S: PdfSource>(obj: &Object, _store: &ObjectStore<S>) -> DocResult<i64> {
        obj.as_i64().ok_or(DocError::UnexpectedType)
    }

    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
    }

    #[test]
    fn test_empty_tree() {
        let store = build_store();
        let root = Object::Dictionary(HashMap::new());
        let tree = NumberTree::parse(&root, &store, convert_int).unwrap();
        assert!(tree.entries().is_empty());
    }

    #[test]
    fn test_single_level_leaf() {
        let store = build_store();
        let mut dict = HashMap::new();
        dict.insert(
            Name::nums(),
            Object::Array(vec![int_obj(0), int_obj(100), int_obj(5), int_obj(500)]),
        );
        let root = Object::Dictionary(dict);
        let tree = NumberTree::parse(&root, &store, convert_int).unwrap();
        assert_eq!(tree.entries().len(), 2);
        assert_eq!(tree.get(0), Some(&100));
        assert_eq!(tree.get(5), Some(&500));
    }

    #[test]
    fn test_get_floor_exact() {
        let store = build_store();
        let mut dict = HashMap::new();
        dict.insert(
            Name::nums(),
            Object::Array(vec![int_obj(0), int_obj(10), int_obj(5), int_obj(50)]),
        );
        let root = Object::Dictionary(dict);
        let tree = NumberTree::parse(&root, &store, convert_int).unwrap();
        assert_eq!(tree.floor(5), Some(&50));
    }

    #[test]
    fn test_get_floor_between() {
        let store = build_store();
        let mut dict = HashMap::new();
        dict.insert(
            Name::nums(),
            Object::Array(vec![
                int_obj(0),
                int_obj(10),
                int_obj(5),
                int_obj(50),
                int_obj(10),
                int_obj(100),
            ]),
        );
        let root = Object::Dictionary(dict);
        let tree = NumberTree::parse(&root, &store, convert_int).unwrap();
        // 7 is between 5 and 10, floor should be 5 -> value 50
        assert_eq!(tree.floor(7), Some(&50));
    }

    #[test]
    fn test_get_floor_before_first() {
        let store = build_store();
        let mut dict = HashMap::new();
        dict.insert(Name::nums(), Object::Array(vec![int_obj(5), int_obj(50)]));
        let root = Object::Dictionary(dict);
        let tree = NumberTree::parse(&root, &store, convert_int).unwrap();
        assert_eq!(tree.floor(3), None);
    }

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

        let mut leaf1 = HashMap::new();
        leaf1.insert(Name::nums(), Object::Array(vec![int_obj(0), int_obj(100)]));

        let mut leaf2 = HashMap::new();
        leaf2.insert(Name::nums(), Object::Array(vec![int_obj(10), int_obj(200)]));

        let mut root_dict = HashMap::new();
        root_dict.insert(
            Name::kids(),
            Object::Array(vec![Object::Dictionary(leaf1), Object::Dictionary(leaf2)]),
        );

        let root = Object::Dictionary(root_dict);
        let tree = NumberTree::parse(&root, &store, convert_int).unwrap();
        assert_eq!(tree.entries().len(), 2);
        assert_eq!(tree.get(0), Some(&100));
        assert_eq!(tree.get(10), Some(&200));
    }

    #[test]
    fn test_lookup_miss() {
        let store = build_store();
        let mut dict = HashMap::new();
        dict.insert(Name::nums(), Object::Array(vec![int_obj(1), int_obj(11)]));
        let root = Object::Dictionary(dict);
        let tree = NumberTree::parse(&root, &store, convert_int).unwrap();
        assert_eq!(tree.get(1), Some(&11));
        assert_eq!(tree.get(99), None);
    }
}