1use std::collections::HashSet;
22
23use pdfrum_common::{DiagKind, Diagnostics, Limits, Severity};
24use pdfrum_object::{Array, Dict, Name, Object, Resolve, decode_text};
25
26use crate::names;
27
28#[derive(Debug, Clone, PartialEq)]
30pub struct NameTree {
31 pub root: Dict,
33}
34
35impl NameTree {
36 #[must_use]
67 pub fn open<R: Resolve>(catalog: &Dict, category: &Name, r: &R) -> Option<NameTree> {
68 let names = catalog.dict(names::NAMES, r)?;
69 Some(NameTree {
70 root: names.dict(category, r)?,
71 })
72 }
73
74 #[must_use]
101 pub fn lookup<R: Resolve>(
102 &self,
103 key: &str,
104 r: &R,
105 limits: &Limits,
106 diags: &mut Diagnostics,
107 ) -> Option<Object> {
108 let mut seen = HashSet::new();
109 let mut index = 0;
110 search_by_name(&self.root, key, 0, &mut seen, &mut index, r, limits, diags)
111 }
112
113 #[must_use]
149 pub fn lookup_by_index<R: Resolve>(
150 &self,
151 target: usize,
152 r: &R,
153 limits: &Limits,
154 diags: &mut Diagnostics,
155 ) -> Option<(String, Object)> {
156 let mut cursor = 0;
157 search_by_index(&self.root, target, &mut cursor, 0, r, limits, diags)
158 }
159
160 #[must_use]
189 pub fn count<R: Resolve>(&self, r: &R, limits: &Limits, diags: &mut Diagnostics) -> usize {
190 let mut seen = Vec::new();
191 count_names(&self.root, 0, &mut seen, r, limits, diags)
192 }
193}
194
195fn node_limits(limits: &Array, diags: &mut Diagnostics) -> (String, String) {
201 if limits.len() < 2 {
202 diags.record(Severity::Recovered, DiagKind::NameTreeLimitsRepaired, None);
203 }
204 let read = |index: usize| {
205 limits
206 .string_at(index)
207 .map(|s| decode_text(s.as_bytes()).into_owned())
208 .unwrap_or_default()
209 };
210 let (low, high) = (read(0), read(1));
211 if low > high {
212 diags.record(Severity::Recovered, DiagKind::NameTreeLimitsRepaired, None);
213 return (high, low);
214 }
215 (low, high)
216}
217
218fn traversed(object: Option<&Object>, seen: &mut HashSet<u32>) -> bool {
224 let Some(num) = object.and_then(Object::as_ref_id).map(|r| r.num) else {
225 return false;
226 };
227 !seen.insert(num)
228}
229
230fn array_traversed(raw: Option<&Object>, array: &Array, seen: &mut HashSet<u32>) -> bool {
235 let mut hit = traversed(raw, seen);
236 for element in array.iter() {
237 if traversed(Some(element), seen) {
238 hit = true;
239 }
240 }
241 hit
242}
243
244#[allow(clippy::too_many_arguments)]
245fn search_by_name<R: Resolve>(
246 node: &Dict,
247 key: &str,
248 depth: u32,
249 seen: &mut HashSet<u32>,
250 index: &mut usize,
251 r: &R,
252 limits: &Limits,
253 diags: &mut Diagnostics,
254) -> Option<Object> {
255 if depth > limits.max_name_tree_depth {
256 diags.record(Severity::Suspicious, DiagKind::TreeDepthExceeded, None);
257 return None;
258 }
259 let mut leaf = node.array(names::NAMES, r);
260 let mut bounds = node.array(names::LIMITS, r);
261 if let Some(array) = &leaf
262 && array_traversed(node.raw(names::NAMES), array, seen)
263 {
264 leaf = None;
265 }
266 if let Some(array) = &bounds
267 && array_traversed(node.raw(names::LIMITS), array, seen)
268 {
269 bounds = None;
270 }
271
272 if let Some(bounds) = &bounds {
273 let (low, high) = node_limits(bounds, diags);
274 if key < low.as_str() || key > high.as_str() {
275 return None;
276 }
277 }
278
279 if let Some(leaf) = &leaf {
280 if leaf.len() % 2 == 1 {
281 diags.record(Severity::Suspicious, DiagKind::NameTreeMalformed, None);
282 }
283 for pair in 0..leaf.len() / 2 {
284 let at = leaf
285 .string_at(pair * 2)
286 .map(|s| decode_text(s.as_bytes()).into_owned())
287 .unwrap_or_default();
288 if at.as_str() > key {
289 break;
290 }
291 if at.as_str() < key {
292 continue;
293 }
294 *index += pair;
295 return leaf.get(pair * 2 + 1, r).map(|v| v.get().clone());
296 }
297 *index += leaf.len() / 2;
298 return None;
299 }
300
301 let kids = node.array(names::KIDS, r)?;
302 if traversed(node.raw(names::KIDS), seen) {
303 return None;
304 }
305 for slot in 0..kids.len() {
306 if traversed(kids.raw_at(slot), seen) {
307 continue;
308 }
309 let Some(kid) = kids.dict_at(slot, r) else {
310 continue;
311 };
312 if let Some(found) = search_by_name(&kid, key, depth + 1, seen, index, r, limits, diags) {
313 return Some(found);
314 }
315 }
316 None
317}
318
319fn search_by_index<R: Resolve>(
320 node: &Dict,
321 target: usize,
322 cursor: &mut usize,
323 depth: u32,
324 r: &R,
325 limits: &Limits,
326 diags: &mut Diagnostics,
327) -> Option<(String, Object)> {
328 if depth > limits.max_name_tree_depth {
329 diags.record(Severity::Suspicious, DiagKind::TreeDepthExceeded, None);
330 return None;
331 }
332 if let Some(leaf) = node.array(names::NAMES, r) {
333 let count = leaf.len() / 2;
334 if target >= *cursor + count {
335 *cursor += count;
336 return None;
337 }
338 let slot = (target - *cursor) * 2;
339 let value = leaf.get(slot + 1, r)?.get().clone();
340 let key = leaf
341 .string_at(slot)
342 .map(|s| decode_text(s.as_bytes()).into_owned())
343 .unwrap_or_default();
344 return Some((key, value));
345 }
346 let kids = node.array(names::KIDS, r)?;
347 for slot in 0..kids.len() {
348 let Some(kid) = kids.dict_at(slot, r) else {
349 continue;
350 };
351 if let Some(found) = search_by_index(&kid, target, cursor, depth + 1, r, limits, diags) {
352 return Some(found);
353 }
354 }
355 None
356}
357
358fn count_names<R: Resolve>(
359 node: &Dict,
360 depth: u32,
361 seen: &mut Vec<Dict>,
362 r: &R,
363 limits: &Limits,
364 diags: &mut Diagnostics,
365) -> usize {
366 if depth > limits.max_name_tree_depth {
367 diags.record(Severity::Suspicious, DiagKind::TreeDepthExceeded, None);
368 return 0;
369 }
370 if seen.contains(node) {
371 return 0;
372 }
373 seen.push(node.clone());
374 if let Some(leaf) = node.array(names::NAMES, r) {
376 return leaf.len() / 2;
377 }
378 let Some(kids) = node.array(names::KIDS, r) else {
379 return 0;
380 };
381 (0..kids.len())
382 .filter_map(|slot| kids.dict_at(slot, r))
383 .map(|kid| count_names(&kid, depth + 1, seen, r, limits, diags))
384 .sum()
385}
386
387#[must_use]
419pub fn lookup_named_dest<R: Resolve>(
420 catalog: &Dict,
421 key: &[u8],
422 r: &R,
423 limits: &Limits,
424 diags: &mut Diagnostics,
425) -> Option<Array> {
426 if let Some(tree) = NameTree::open(catalog, names::DESTS, r) {
427 let text = decode_text(key).into_owned();
428 if let Some(found) = tree
429 .lookup(&text, r, limits, diags)
430 .and_then(|value| named_dest_from(&value, r))
431 {
432 return Some(found);
433 }
434 }
435 let old_style = catalog.dict(names::DESTS, r)?;
436 let value = old_style.get(&Name::new(key.to_vec()), r)?.get().clone();
437 let found = named_dest_from(&value, r);
438 if found.is_some() {
439 diags.record(Severity::Recovered, DiagKind::LegacyNamedDest, None);
440 }
441 found
442}
443
444fn named_dest_from<R: Resolve>(object: &Object, r: &R) -> Option<Array> {
447 match object {
448 Object::Array(array) => Some(array.clone()),
449 Object::Dict(dict) => dict.array(names::D, r),
450 _ => None,
451 }
452}
453
454#[cfg(test)]
455mod tests {
456 use super::{NameTree, lookup_named_dest};
457 use pdfrum_common::{Diagnostics, Limits};
458 use pdfrum_object::{Array, Dict, Name, NoResolve, Object, PdfString};
459
460 fn dict(pairs: &[(&str, Object)]) -> Dict {
461 Dict::from_pairs(
462 pairs
463 .iter()
464 .map(|(k, v)| (Name::from(*k), v.clone()))
465 .collect::<Vec<_>>(),
466 )
467 }
468
469 fn leaf(pairs: &[(&[u8], i64)]) -> Dict {
470 let mut array = Array::new();
471 for (key, value) in pairs {
472 array.push(Object::Str(PdfString::literal(key)));
473 array.push(Object::Int(*value));
474 }
475 dict(&[("Names", Object::Array(array))])
476 }
477
478 #[test]
479 fn a_missing_link_yields_no_tree_at_all() {
480 assert!(NameTree::open(&Dict::new(), pdfrum_object::names::DESTS, &NoResolve).is_none());
481 let catalog = dict(&[("Names", Object::Dict(Dict::new()))]);
482 assert!(NameTree::open(&catalog, pdfrum_object::names::DESTS, &NoResolve).is_none());
483 }
484
485 #[test]
486 fn a_leaf_finds_its_keys_and_stops_at_the_first_greater_one() {
487 let tree = NameTree {
488 root: leaf(&[(b"a", 1), (b"c", 3), (b"e", 5)]),
489 };
490 let (l, mut d) = (Limits::default(), Diagnostics::default());
491 assert_eq!(
492 tree.lookup("c", &NoResolve, &l, &mut d),
493 Some(Object::Int(3))
494 );
495 assert_eq!(tree.lookup("b", &NoResolve, &l, &mut d), None);
496 assert_eq!(
497 tree.lookup("e", &NoResolve, &l, &mut d),
498 Some(Object::Int(5))
499 );
500 }
501
502 #[test]
503 fn a_key_with_a_byte_order_mark_reads_back_as_its_text() {
504 let tree = NameTree {
505 root: leaf(&[(b"\xFE\xFF\x00\x31", 100)]),
506 };
507 let (l, mut d) = (Limits::default(), Diagnostics::default());
508 assert_eq!(
509 tree.lookup_by_index(0, &NoResolve, &l, &mut d),
510 Some(("1".to_owned(), Object::Int(100)))
511 );
512 assert_eq!(
513 tree.lookup("1", &NoResolve, &l, &mut d),
514 Some(Object::Int(100))
515 );
516 }
517
518 #[test]
519 fn a_lookup_reads_a_four_element_limits_array_without_shortening_it() {
520 let mut node = leaf(&[(b"1", 111), (b"9", 999)]);
521 let bounds = Array::of([
522 Object::Str(PdfString::literal(b"1")),
523 Object::Str(PdfString::literal(b"9")),
524 Object::Str(PdfString::literal(b"extra")),
525 Object::Str(PdfString::literal(b"more")),
526 ]);
527 node.push(Name::from("Limits"), Object::Array(bounds));
528 let tree = NameTree { root: node.clone() };
529 let (l, mut d) = (Limits::default(), Diagnostics::default());
530 assert_eq!(
531 tree.lookup("9", &NoResolve, &l, &mut d),
532 Some(Object::Int(999))
533 );
534 assert_eq!(
535 tree.lookup("1", &NoResolve, &l, &mut d),
536 Some(Object::Int(111))
537 );
538 assert_eq!(
540 tree.root
541 .array(pdfrum_object::names::LENGTH, &NoResolve)
542 .map_or(4, |a| a.len()),
543 4
544 );
545 }
546
547 #[test]
548 fn a_node_holding_both_names_and_kids_is_a_leaf() {
549 let mut node = leaf(&[(b"a", 1)]);
550 node.push(
551 Name::from("Kids"),
552 Object::Array(Array::of([Object::Dict(leaf(&[(b"z", 26)]))])),
553 );
554 let tree = NameTree { root: node };
555 let (l, mut d) = (Limits::default(), Diagnostics::default());
556 assert_eq!(tree.lookup("z", &NoResolve, &l, &mut d), None);
557 assert_eq!(tree.count(&NoResolve, &l, &mut d), 1);
558 }
559
560 #[test]
561 fn limits_the_wrong_way_round_are_read_swapped() {
562 let mut node = leaf(&[(b"a", 1)]);
563 node.push(
564 Name::from("Limits"),
565 Object::Array(Array::of([
566 Object::Str(PdfString::literal(b"z")),
567 Object::Str(PdfString::literal(b"a")),
568 ])),
569 );
570 let tree = NameTree { root: node };
571 let (l, mut d) = (Limits::default(), Diagnostics::default());
572 assert_eq!(
573 tree.lookup("a", &NoResolve, &l, &mut d),
574 Some(Object::Int(1))
575 );
576 assert!(d.contains(&pdfrum_common::DiagKind::NameTreeLimitsRepaired));
577 }
578
579 #[test]
580 fn the_old_style_dests_dictionary_is_keyed_on_raw_bytes() {
581 let dest = Array::of([Object::Int(11), Object::Name(Name::from("Fit"))]);
582 let catalog = dict(&[(
583 "Dests",
584 Object::Dict(dict(&[("FirstAlternate", Object::Array(dest.clone()))])),
585 )]);
586 let (l, mut d) = (Limits::default(), Diagnostics::default());
587 assert_eq!(
588 lookup_named_dest(&catalog, b"FirstAlternate", &NoResolve, &l, &mut d),
589 Some(dest)
590 );
591 assert!(d.contains(&pdfrum_common::DiagKind::LegacyNamedDest));
592 assert_eq!(
593 lookup_named_dest(&catalog, b"Missing", &NoResolve, &l, &mut d),
594 None
595 );
596 }
597
598 #[test]
599 fn a_dictionary_valued_entry_unwraps_its_d_array() {
600 let dest = Array::of([Object::Int(1), Object::Name(Name::from("Fit"))]);
601 let wrapped = Object::Dict(dict(&[("D", Object::Array(dest.clone()))]));
602 let catalog = dict(&[("Dests", Object::Dict(dict(&[("Here", wrapped)])))]);
603 let (l, mut d) = (Limits::default(), Diagnostics::default());
604 assert_eq!(
605 lookup_named_dest(&catalog, b"Here", &NoResolve, &l, &mut d),
606 Some(dest)
607 );
608 }
609}