From 5c19ff6a7e040a6a83edcaad5f7110d31705fd20 Mon Sep 17 00:00:00 2001
From: Sebastian K <s3bk@protonmail.com>
Date: Mon, 22 Jan 2024 13:50:55 +0100
Subject: [PATCH] fix pagelabels

---
 pdf/examples/names.rs   |   7 +-
 pdf/src/object/types.rs | 145 ++++++++++++++++++++++++++++++++++------
 2 files changed, 131 insertions(+), 21 deletions(-)

diff --git a/pdf/examples/names.rs b/pdf/examples/names.rs
index 82cd8b0..b273c72 100644
--- a/pdf/examples/names.rs
+++ b/pdf/examples/names.rs
@@ -119,6 +119,11 @@ fn main() {
         }
     }
 
-
     println!("{} items", count);
+
+    if let Some(ref labels) = catalog.page_labels {
+        labels.walk(&resolver, &mut |page: i32, label| {
+            println!("{page} -> {:?}", label);
+        });
+    }
 }
diff --git a/pdf/src/object/types.rs b/pdf/src/object/types.rs
index ec27c89..e93d540 100644
--- a/pdf/src/object/types.rs
+++ b/pdf/src/object/types.rs
@@ -130,7 +130,7 @@ pub struct Catalog {
     pub pages: PagesRc,
 
     #[pdf(key="PageLabels")]
-    pub page_labels: Option<NameTree<PdfString>>,
+    pub page_labels: Option<NumberTree<PageLabel>>,
 
     #[pdf(key="Names")]
     pub names: Option<MaybeRef<NameDictionary>>,
@@ -356,7 +356,7 @@ impl Page {
 impl SubType<PagesNode> for Page {}
 
 
-#[derive(Object, DataSize)]
+#[derive(Object, DataSize, Debug, ObjectWrite)]
 pub struct PageLabel {
     #[pdf(key="S")]
     pub style:  Option<Counter>,
@@ -1075,30 +1075,19 @@ impl DataSize for AppearanceStreamEntry {
     }
 }
 
-#[derive(Debug, DataSize)]
+#[derive(Debug, DataSize, Clone, Object, ObjectWrite, DeepClone)]
 pub enum Counter {
+    #[pdf(name="D")]
     Arabic,
+    #[pdf(name="r")]
     RomanUpper,
+    #[pdf(name="R")]
     RomanLower,
+    #[pdf(name="a")]
     AlphaUpper,
+    #[pdf(name="A")]
     AlphaLower
 }
-impl Object for Counter {
-    // fn serialize<W: io::Write>(&self, out: &mut W) -> Result<()> {
-    //     let style_code = match *self {
-    //         Counter::Arabic     => "D",
-    //         Counter::RomanLower => "r",
-    //         Counter::RomanUpper => "R",
-    //         Counter::AlphaLower => "a",
-    //         Counter::AlphaUpper => "A"
-    //     };
-    //     out.write_all(style_code.as_bytes())?;
-    //     Ok(())
-    // }
-    fn from_primitive(_: Primitive, _: &impl Resolve) -> Result<Self> {
-        unimplemented!();
-    }
-}
 
 #[derive(Debug, DataSize)]
 pub enum NameTreeNode<T> {
@@ -1138,7 +1127,6 @@ impl<T: Object> Object for NameTree<T> {
     fn from_primitive(p: Primitive, resolve: &impl Resolve) -> Result<Self> {
         let mut dict = t!(p.resolve(resolve)?.into_dictionary());
         
-        // Quite long function..=
         let limits = match dict.remove("Limits") {
             Some(limits) => {
                 let limits = limits.resolve(resolve)?.into_array()?;
@@ -1196,6 +1184,123 @@ impl<T: ObjectWrite> ObjectWrite for NameTree<T> {
     }
 }
 
+#[derive(DataSize, Debug)]
+pub struct NumberTree<T> {
+    pub limits: Option<(i32, i32)>,
+    pub node: NumberTreeNode<T>,
+}
+
+#[derive(DataSize, Debug)]
+pub enum NumberTreeNode<T> {
+    Leaf(Vec<(i32, T)>),
+    Intermediate(Vec<Ref<NumberTree<T>>>),
+}
+impl<T: Object> Object for NumberTree<T> {
+    fn from_primitive(p: Primitive, resolve: &impl Resolve) -> Result<Self> {
+        let mut dict = p.resolve(resolve)?.into_dictionary()?;
+
+        let limits = match dict.remove("Limits") {
+            Some(limits) => {
+                let limits = t!(limits.resolve(resolve)?.into_array());
+                if limits.len() != 2 {
+                    bail!("Error reading NameTree: 'Limits' is not of length 2");
+                }
+                let min = t!(limits[0].as_integer());
+                let max = t!(limits[1].as_integer());
+
+                Some((min, max))
+            }
+            None => None
+        };
+
+        let kids = dict.remove("Kids");
+        let nums = dict.remove("Nums");
+        match (kids, nums) {
+            (Some(kids), _) => {
+                let kids = t!(kids.resolve(resolve)?.into_array()?.iter().map(|kid|
+                    Ref::<NumberTree<T>>::from_primitive(kid.clone(), resolve)
+                ).collect::<Result<Vec<_>>>());
+                Ok(NumberTree {
+                    limits,
+                    node: NumberTreeNode::Intermediate (kids)
+                })
+            }
+            (None, Some(nums)) => {
+                let list = nums.into_array()?;
+                let mut items = Vec::with_capacity(list.len() / 2);
+                for (key, item) in list.into_iter().tuples() {
+                    let idx = t!(key.as_integer());
+                    let val = t!(T::from_primitive(item, resolve));
+                    items.push((idx, val));
+                }
+                Ok(NumberTree {
+                    limits,
+                    node: NumberTreeNode::Leaf(items)
+                })
+            }
+            (None, None) => {
+                warn!("Neither Kids nor Names present in NumberTree node.");
+                Ok(NumberTree {
+                    limits,
+                    node: NumberTreeNode::Intermediate(vec![])
+                })
+            }
+        }
+    }
+}
+impl<T: ObjectWrite> ObjectWrite for NumberTree<T> {
+    fn to_primitive(&self, update: &mut impl Updater) -> Result<Primitive> {
+        let mut dict = Dictionary::new();
+        if let Some(limits) = self.limits {
+            dict.insert("Limits", vec![limits.0.into(), limits.1.into()]);
+        }
+        match self.node {
+            NumberTreeNode::Leaf(ref items) => {
+                let mut nums = Vec::with_capacity(items.len() * 2);
+                for &(idx, ref label) in items {
+                    nums.push(idx.into());
+                    nums.push(label.to_primitive(update)?);
+                }
+                dict.insert("Nums", nums);
+            }
+            NumberTreeNode::Intermediate(ref kids) => {
+                dict.insert("Kids", kids.iter().map(|r| r.get_inner().into()).collect_vec());
+            }
+        }
+        Ok(dict.into())
+    }
+}
+impl<T: Object+DataSize> NumberTree<T> {
+    pub fn walk(&self, r: &impl Resolve, callback: &mut dyn FnMut(i32, &T)) -> Result<(), PdfError> {
+        match self.node {
+            NumberTreeNode::Leaf(ref items) => {
+                for &(idx, ref val) in items {
+                    callback(idx, val);
+                }
+            }
+            NumberTreeNode::Intermediate(ref items) => {
+                for &tree_ref in items {
+                    let tree = r.get(tree_ref)?;
+                    tree.walk(r, callback)?;
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+#[derive(Object, ObjectWrite, Clone, DeepClone, Debug)]
+pub struct LageLabel {
+    #[pdf(key="S")]
+    style: Option<Counter>,
+    
+    #[pdf(key="P")]
+    prefix: Option<PdfString>,
+
+    #[pdf(key="St")]
+    start: Option<i32>,
+}
+
 #[derive(Debug, Clone, DataSize)]
 pub enum DestView {
     // left, top, zoom