oxidize_pdf/page_labels/
page_label_tree.rs

1//! Page label tree structure for managing page numbering
2
3use crate::objects::{Array, Dictionary, Object};
4use crate::page_labels::PageLabel;
5use std::collections::BTreeMap;
6
7/// Page label tree - manages custom page numbering for a document
8#[derive(Debug, Clone)]
9pub struct PageLabelTree {
10    /// Page label ranges, sorted by starting page
11    ranges: BTreeMap<u32, PageLabel>,
12}
13
14impl PageLabelTree {
15    /// Create a new empty page label tree
16    pub fn new() -> Self {
17        Self {
18            ranges: BTreeMap::new(),
19        }
20    }
21
22    /// Add a page label range
23    pub fn add_range(&mut self, start_page: u32, label: PageLabel) {
24        self.ranges.insert(start_page, label);
25    }
26
27    /// Get the page label for a specific page
28    pub fn get_label(&self, page_index: u32) -> Option<String> {
29        // Find the applicable range
30        let mut applicable_range = None;
31        let mut range_start = 0;
32
33        for (&start, label) in &self.ranges {
34            if start <= page_index {
35                applicable_range = Some(label);
36                range_start = start;
37            } else {
38                break;
39            }
40        }
41
42        // Format the label if found
43        applicable_range.map(|label| {
44            let offset = page_index - range_start;
45            label.format_label(offset)
46        })
47    }
48
49    /// Get all page labels for a document
50    pub fn get_all_labels(&self, total_pages: u32) -> Vec<String> {
51        (0..total_pages)
52            .map(|i| self.get_label(i).unwrap_or_else(|| (i + 1).to_string()))
53            .collect()
54    }
55
56    /// Convert to PDF number tree dictionary
57    pub fn to_dict(&self) -> Dictionary {
58        let mut dict = Dictionary::new();
59
60        // Create nums array [key1 val1 key2 val2 ...]
61        let mut nums = Array::new();
62
63        for (&start_page, label) in &self.ranges {
64            nums.push(Object::Integer(start_page as i64));
65            nums.push(Object::Dictionary(label.to_dict()));
66        }
67
68        dict.set("Nums", Object::Array(nums.into()));
69
70        dict
71    }
72
73    /// Create from PDF dictionary
74    pub fn from_dict(dict: &Dictionary) -> Option<Self> {
75        let nums_array = match dict.get("Nums")? {
76            Object::Array(arr) => arr,
77            _ => return None,
78        };
79        let mut tree = Self::new();
80
81        // Parse pairs of [page_index, label_dict]
82        let elements: Vec<&Object> = nums_array.iter().collect();
83        for i in (0..elements.len()).step_by(2) {
84            if i + 1 >= elements.len() {
85                break;
86            }
87
88            let page_index = match elements[i] {
89                Object::Integer(n) => *n as u32,
90                _ => continue,
91            };
92            let label_dict = match elements[i + 1] {
93                Object::Dictionary(d) => d,
94                _ => continue,
95            };
96
97            // Parse label from dictionary
98            let style = if let Some(Object::Name(type_name)) = label_dict.get("Type") {
99                match type_name.as_str() {
100                    "D" => PageLabelStyle::DecimalArabic,
101                    "r" => PageLabelStyle::UppercaseRoman,
102                    "R" => PageLabelStyle::LowercaseRoman,
103                    "A" => PageLabelStyle::UppercaseLetters,
104                    "a" => PageLabelStyle::LowercaseLetters,
105                    _ => PageLabelStyle::None,
106                }
107            } else {
108                PageLabelStyle::None
109            };
110
111            let mut label = PageLabel::new(style);
112
113            if let Some(Object::String(prefix)) = label_dict.get("P") {
114                label = label.with_prefix(prefix);
115            }
116
117            if let Some(Object::Integer(start)) = label_dict.get("St") {
118                label = label.starting_at(*start as u32);
119            }
120
121            tree.add_range(page_index, label);
122        }
123
124        Some(tree)
125    }
126}
127
128impl Default for PageLabelTree {
129    fn default() -> Self {
130        Self::new()
131    }
132}
133
134/// Builder for creating page label trees
135pub struct PageLabelBuilder {
136    tree: PageLabelTree,
137    current_page: u32,
138}
139
140impl Default for PageLabelBuilder {
141    fn default() -> Self {
142        Self::new()
143    }
144}
145
146impl PageLabelBuilder {
147    /// Create a new page label builder
148    pub fn new() -> Self {
149        Self {
150            tree: PageLabelTree::new(),
151            current_page: 0,
152        }
153    }
154
155    /// Add a range with a specific label
156    pub fn add_range(mut self, num_pages: u32, label: PageLabel) -> Self {
157        self.tree.add_range(self.current_page, label);
158        self.current_page += num_pages;
159        self
160    }
161
162    /// Add pages with decimal numbering
163    pub fn decimal_pages(self, num_pages: u32) -> Self {
164        self.add_range(num_pages, PageLabel::decimal())
165    }
166
167    /// Add pages with roman numbering
168    pub fn roman_pages(self, num_pages: u32, uppercase: bool) -> Self {
169        let label = if uppercase {
170            PageLabel::roman_uppercase()
171        } else {
172            PageLabel::roman_lowercase()
173        };
174        self.add_range(num_pages, label)
175    }
176
177    /// Add pages with letter numbering
178    pub fn letter_pages(self, num_pages: u32, uppercase: bool) -> Self {
179        let label = if uppercase {
180            PageLabel::letters_uppercase()
181        } else {
182            PageLabel::letters_lowercase()
183        };
184        self.add_range(num_pages, label)
185    }
186
187    /// Add pages with only a prefix
188    pub fn prefix_pages(self, num_pages: u32, prefix: impl Into<String>) -> Self {
189        self.add_range(num_pages, PageLabel::prefix_only(prefix))
190    }
191
192    /// Build the page label tree
193    pub fn build(self) -> PageLabelTree {
194        self.tree
195    }
196}
197
198// Import PageLabelStyle from the other module
199use crate::page_labels::PageLabelStyle;
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn test_page_label_tree() {
207        let mut tree = PageLabelTree::new();
208
209        // Add roman numerals for first 3 pages
210        tree.add_range(0, PageLabel::roman_lowercase());
211
212        // Add decimal starting at page 3
213        tree.add_range(3, PageLabel::decimal());
214
215        // Test labels
216        assert_eq!(tree.get_label(0), Some("i".to_string()));
217        assert_eq!(tree.get_label(1), Some("ii".to_string()));
218        assert_eq!(tree.get_label(2), Some("iii".to_string()));
219        assert_eq!(tree.get_label(3), Some("1".to_string()));
220        assert_eq!(tree.get_label(4), Some("2".to_string()));
221        assert_eq!(tree.get_label(5), Some("3".to_string()));
222    }
223
224    #[test]
225    fn test_page_label_with_prefix() {
226        let mut tree = PageLabelTree::new();
227
228        // Preface with prefix
229        tree.add_range(0, PageLabel::prefix_only("Cover"));
230        tree.add_range(1, PageLabel::roman_lowercase().with_prefix("p. "));
231        tree.add_range(4, PageLabel::decimal().with_prefix("Chapter "));
232
233        assert_eq!(tree.get_label(0), Some("Cover".to_string()));
234        assert_eq!(tree.get_label(1), Some("p. i".to_string()));
235        assert_eq!(tree.get_label(2), Some("p. ii".to_string()));
236        assert_eq!(tree.get_label(3), Some("p. iii".to_string()));
237        assert_eq!(tree.get_label(4), Some("Chapter 1".to_string()));
238        assert_eq!(tree.get_label(5), Some("Chapter 2".to_string()));
239    }
240
241    #[test]
242    fn test_page_label_with_start() {
243        let mut tree = PageLabelTree::new();
244
245        // Start numbering at 10
246        tree.add_range(0, PageLabel::decimal().starting_at(10));
247
248        assert_eq!(tree.get_label(0), Some("10".to_string()));
249        assert_eq!(tree.get_label(1), Some("11".to_string()));
250        assert_eq!(tree.get_label(2), Some("12".to_string()));
251    }
252
253    #[test]
254    fn test_get_all_labels() {
255        let mut tree = PageLabelTree::new();
256        tree.add_range(0, PageLabel::roman_lowercase());
257        tree.add_range(2, PageLabel::decimal());
258
259        let labels = tree.get_all_labels(5);
260        assert_eq!(labels, vec!["i", "ii", "1", "2", "3"]);
261    }
262
263    #[test]
264    fn test_page_label_builder() {
265        let tree = PageLabelBuilder::new()
266            .prefix_pages(1, "Cover")
267            .roman_pages(3, false)
268            .decimal_pages(10)
269            .letter_pages(3, true)
270            .build();
271
272        assert_eq!(tree.get_label(0), Some("Cover".to_string()));
273        assert_eq!(tree.get_label(1), Some("i".to_string()));
274        assert_eq!(tree.get_label(2), Some("ii".to_string()));
275        assert_eq!(tree.get_label(3), Some("iii".to_string()));
276        assert_eq!(tree.get_label(4), Some("1".to_string()));
277        assert_eq!(tree.get_label(13), Some("10".to_string()));
278        assert_eq!(tree.get_label(14), Some("A".to_string()));
279        assert_eq!(tree.get_label(15), Some("B".to_string()));
280        assert_eq!(tree.get_label(16), Some("C".to_string()));
281    }
282
283    #[test]
284    fn test_to_dict() {
285        let mut tree = PageLabelTree::new();
286        tree.add_range(0, PageLabel::roman_lowercase());
287        tree.add_range(3, PageLabel::decimal().with_prefix("Page "));
288
289        let dict = tree.to_dict();
290        assert!(dict.get("Nums").is_some());
291    }
292
293    #[test]
294    fn test_page_label_tree_default() {
295        let tree = PageLabelTree::default();
296        // Empty tree should return None for any page
297        assert!(tree.get_label(0).is_none());
298        assert!(tree.get_label(100).is_none());
299    }
300
301    #[test]
302    fn test_page_label_tree_clone() {
303        let mut tree = PageLabelTree::new();
304        tree.add_range(0, PageLabel::decimal());
305        let cloned = tree.clone();
306        assert_eq!(tree.get_label(0), cloned.get_label(0));
307    }
308
309    #[test]
310    fn test_page_label_tree_debug() {
311        let tree = PageLabelTree::new();
312        let debug_str = format!("{:?}", tree);
313        assert!(debug_str.contains("PageLabelTree"));
314    }
315
316    #[test]
317    fn test_page_label_builder_default() {
318        let builder = PageLabelBuilder::default();
319        let tree = builder.build();
320        // Empty tree
321        assert!(tree.get_label(0).is_none());
322    }
323
324    #[test]
325    fn test_page_label_builder_roman_uppercase() {
326        let tree = PageLabelBuilder::new().roman_pages(5, true).build();
327
328        assert_eq!(tree.get_label(0), Some("I".to_string()));
329        assert_eq!(tree.get_label(1), Some("II".to_string()));
330        assert_eq!(tree.get_label(2), Some("III".to_string()));
331        assert_eq!(tree.get_label(3), Some("IV".to_string()));
332        assert_eq!(tree.get_label(4), Some("V".to_string()));
333    }
334
335    #[test]
336    fn test_page_label_builder_letter_lowercase() {
337        let tree = PageLabelBuilder::new().letter_pages(3, false).build();
338
339        assert_eq!(tree.get_label(0), Some("a".to_string()));
340        assert_eq!(tree.get_label(1), Some("b".to_string()));
341        assert_eq!(tree.get_label(2), Some("c".to_string()));
342    }
343
344    #[test]
345    fn test_get_all_labels_empty_tree() {
346        let tree = PageLabelTree::new();
347        // Empty tree should return default numbering
348        let labels = tree.get_all_labels(3);
349        assert_eq!(labels, vec!["1", "2", "3"]);
350    }
351
352    #[test]
353    fn test_from_dict_empty() {
354        let mut dict = Dictionary::new();
355        dict.set("Nums", Object::Array(Array::new().into()));
356
357        let tree = PageLabelTree::from_dict(&dict);
358        assert!(tree.is_some());
359        let tree = tree.unwrap();
360        assert!(tree.get_label(0).is_none());
361    }
362
363    #[test]
364    fn test_from_dict_missing_nums() {
365        let dict = Dictionary::new();
366        let tree = PageLabelTree::from_dict(&dict);
367        assert!(tree.is_none());
368    }
369
370    #[test]
371    fn test_from_dict_invalid_nums_type() {
372        let mut dict = Dictionary::new();
373        dict.set("Nums", Object::Integer(42));
374
375        let tree = PageLabelTree::from_dict(&dict);
376        assert!(tree.is_none());
377    }
378
379    #[test]
380    fn test_from_dict_with_decimal_labels() {
381        let mut label_dict = Dictionary::new();
382        label_dict.set("Type", Object::Name("D".to_string()));
383
384        let mut nums = Array::new();
385        nums.push(Object::Integer(0));
386        nums.push(Object::Dictionary(label_dict));
387
388        let mut dict = Dictionary::new();
389        dict.set("Nums", Object::Array(nums.into()));
390
391        let tree = PageLabelTree::from_dict(&dict);
392        assert!(tree.is_some());
393    }
394
395    #[test]
396    fn test_from_dict_with_prefix_and_start() {
397        let mut label_dict = Dictionary::new();
398        label_dict.set("Type", Object::Name("D".to_string()));
399        label_dict.set("P", Object::String("Page ".to_string()));
400        label_dict.set("St", Object::Integer(10));
401
402        let mut nums = Array::new();
403        nums.push(Object::Integer(0));
404        nums.push(Object::Dictionary(label_dict));
405
406        let mut dict = Dictionary::new();
407        dict.set("Nums", Object::Array(nums.into()));
408
409        let tree = PageLabelTree::from_dict(&dict);
410        assert!(tree.is_some());
411    }
412
413    #[test]
414    fn test_from_dict_invalid_page_index() {
415        // Test with non-integer page index
416        let label_dict = Dictionary::new();
417
418        let mut nums = Array::new();
419        nums.push(Object::String("not_an_integer".to_string()));
420        nums.push(Object::Dictionary(label_dict));
421
422        let mut dict = Dictionary::new();
423        dict.set("Nums", Object::Array(nums.into()));
424
425        // Should skip invalid entries
426        let tree = PageLabelTree::from_dict(&dict);
427        assert!(tree.is_some());
428    }
429
430    #[test]
431    fn test_from_dict_invalid_label_dict() {
432        // Test with non-dictionary label
433        let mut nums = Array::new();
434        nums.push(Object::Integer(0));
435        nums.push(Object::String("not_a_dict".to_string()));
436
437        let mut dict = Dictionary::new();
438        dict.set("Nums", Object::Array(nums.into()));
439
440        // Should skip invalid entries
441        let tree = PageLabelTree::from_dict(&dict);
442        assert!(tree.is_some());
443    }
444
445    #[test]
446    fn test_from_dict_odd_length_array() {
447        // Array with odd length - last element should be ignored
448        let label_dict = Dictionary::new();
449
450        let mut nums = Array::new();
451        nums.push(Object::Integer(0));
452        nums.push(Object::Dictionary(label_dict));
453        nums.push(Object::Integer(5)); // Missing pair
454
455        let mut dict = Dictionary::new();
456        dict.set("Nums", Object::Array(nums.into()));
457
458        let tree = PageLabelTree::from_dict(&dict);
459        assert!(tree.is_some());
460    }
461
462    #[test]
463    fn test_from_dict_all_style_types() {
464        // Test r (uppercase roman)
465        let mut label_dict1 = Dictionary::new();
466        label_dict1.set("Type", Object::Name("r".to_string()));
467
468        // Test R (lowercase roman)
469        let mut label_dict2 = Dictionary::new();
470        label_dict2.set("Type", Object::Name("R".to_string()));
471
472        // Test A (uppercase letters)
473        let mut label_dict3 = Dictionary::new();
474        label_dict3.set("Type", Object::Name("A".to_string()));
475
476        // Test a (lowercase letters)
477        let mut label_dict4 = Dictionary::new();
478        label_dict4.set("Type", Object::Name("a".to_string()));
479
480        // Test unknown type
481        let mut label_dict5 = Dictionary::new();
482        label_dict5.set("Type", Object::Name("unknown".to_string()));
483
484        let mut nums = Array::new();
485        nums.push(Object::Integer(0));
486        nums.push(Object::Dictionary(label_dict1));
487        nums.push(Object::Integer(5));
488        nums.push(Object::Dictionary(label_dict2));
489        nums.push(Object::Integer(10));
490        nums.push(Object::Dictionary(label_dict3));
491        nums.push(Object::Integer(15));
492        nums.push(Object::Dictionary(label_dict4));
493        nums.push(Object::Integer(20));
494        nums.push(Object::Dictionary(label_dict5));
495
496        let mut dict = Dictionary::new();
497        dict.set("Nums", Object::Array(nums.into()));
498
499        let tree = PageLabelTree::from_dict(&dict);
500        assert!(tree.is_some());
501    }
502
503    #[test]
504    fn test_get_label_no_applicable_range() {
505        let mut tree = PageLabelTree::new();
506        // Add range starting at page 5
507        tree.add_range(5, PageLabel::decimal());
508
509        // Pages before the first range should return None
510        assert!(tree.get_label(0).is_none());
511        assert!(tree.get_label(4).is_none());
512
513        // Pages at or after the range should have labels
514        assert_eq!(tree.get_label(5), Some("1".to_string()));
515        assert_eq!(tree.get_label(6), Some("2".to_string()));
516    }
517
518    #[test]
519    fn test_page_label_builder_chained() {
520        let tree = PageLabelBuilder::new()
521            .prefix_pages(1, "TOC")
522            .roman_pages(2, false)
523            .roman_pages(2, true)
524            .letter_pages(2, false)
525            .letter_pages(2, true)
526            .decimal_pages(5)
527            .build();
528
529        assert_eq!(tree.get_label(0), Some("TOC".to_string()));
530        assert_eq!(tree.get_label(1), Some("i".to_string()));
531        assert_eq!(tree.get_label(3), Some("I".to_string()));
532        assert_eq!(tree.get_label(5), Some("a".to_string()));
533        assert_eq!(tree.get_label(7), Some("A".to_string()));
534        assert_eq!(tree.get_label(9), Some("1".to_string()));
535    }
536
537    #[test]
538    fn test_to_dict_round_trip() {
539        let mut original = PageLabelTree::new();
540        original.add_range(0, PageLabel::roman_lowercase());
541        original.add_range(5, PageLabel::decimal());
542
543        let dict = original.to_dict();
544        let restored = PageLabelTree::from_dict(&dict);
545
546        assert!(restored.is_some());
547        // Note: exact round-trip may not work perfectly due to style mapping differences
548        // but basic structure should be preserved
549    }
550}
oxidize_pdf/page_labels/page_label_tree.rs

oxidize_pdf/page_labels/
page_label_tree.rs
