1use std::cmp::Ordering;
23
24use minicbor::{Decoder, data::Type, decode::Error};
25
26pub enum Arity {
28 Leaf,
29 Fixed(u64),
30 Indefinite,
32}
33
34impl From<Option<u64>> for Arity {
35 fn from(len: Option<u64>) -> Self {
37 match len {
38 Some(n) => Self::Fixed(n),
39 None => Self::Indefinite,
40 }
41 }
42}
43
44pub trait TreeDecode<'b, C>: Sized {
49 type Builder;
53
54 fn begin(d: &mut Decoder<'b>, ctx: &mut C) -> Result<(Self::Builder, Arity), Error>;
57
58 fn child(builder: &mut Self::Builder, child: Self) -> Result<(), Error>;
61
62 fn end(builder: Self::Builder, d: &mut Decoder<'b>, ctx: &mut C) -> Result<Self, Error>;
64}
65
66struct Frame<B> {
67 builder: B,
68 arity: Arity,
69}
70
71impl<B> Frame<B> {
72 fn begin<'b, C, T>(d: &mut Decoder<'b>, ctx: &mut C) -> Result<Self, Error>
73 where
74 T: TreeDecode<'b, C, Builder = B>,
75 {
76 let (builder, arity) = T::begin(d, ctx)?;
77 Ok(Self { builder, arity })
78 }
79
80 fn expects_child(&mut self, d: &mut Decoder<'_>) -> Result<bool, Error> {
81 match self.arity {
82 Arity::Leaf | Arity::Fixed(0) => Ok(false),
83 Arity::Fixed(_) => Ok(true),
84 Arity::Indefinite if d.datatype()? == Type::Break => {
85 d.skip()?;
86 self.arity = Arity::Leaf;
87 Ok(false)
88 }
89 Arity::Indefinite => Ok(true),
90 }
91 }
92}
93
94pub fn decode_tree<'b, C, T>(d: &mut Decoder<'b>, ctx: &mut C) -> Result<T, Error>
96where
97 T: TreeDecode<'b, C>,
98{
99 let mut parents: Vec<Frame<T::Builder>> = Vec::new();
100 let mut current = Frame::begin::<C, T>(d, ctx)?;
101 loop {
102 if current.expects_child(d)? {
103 parents.push(current);
104 current = Frame::begin::<C, T>(d, ctx)?;
105 continue;
106 }
107 let node = T::end(current.builder, d, ctx)?;
108 let Some(mut parent) = parents.pop() else {
109 return Ok(node);
110 };
111 T::child(&mut parent.builder, node)?;
112 if let Arity::Fixed(remaining) = &mut parent.arity {
113 *remaining -= 1;
114 }
115 current = parent;
116 }
117}
118
119pub trait TreeNode: Sized {
121 fn children(&self) -> &[Self];
122
123 fn children_mut(&mut self) -> Option<&mut Vec<Self>>;
125}
126
127pub trait IndexedNode: Sized {
132 fn child_count(&self) -> usize;
133
134 fn child(&self, index: usize) -> &Self;
136}
137
138impl<T: TreeNode> IndexedNode for T {
139 fn child_count(&self) -> usize {
140 self.children().len()
141 }
142
143 fn child(&self, index: usize) -> &Self {
144 &self.children()[index]
145 }
146}
147
148pub fn map_tree<S, T>(
154 root: &S,
155 shallow: impl Fn(&S) -> T,
156 target_children: impl Fn(&mut T) -> Option<&mut Vec<T>>,
157) -> T
158where
159 S: TreeNode,
160{
161 let mut target_root = shallow(root);
162 let mut pending = vec![(root, &mut target_root)];
163 while let Some((source, target)) = pending.pop() {
164 let Some(children) = target_children(target) else {
165 continue;
166 };
167 let source_children = source.children();
168 *children = source_children.iter().map(&shallow).collect();
169 pending.extend(source_children.iter().zip(children.iter_mut()));
170 }
171 target_root
172}
173
174pub fn fold_tree<S, T>(root: &S, mut finish: impl FnMut(&S, Vec<T>) -> T) -> T
181where
182 S: IndexedNode,
183{
184 struct Frame<'a, S, T> {
185 node: &'a S,
186 next: usize,
187 results: Vec<T>,
188 }
189
190 fn open<S: IndexedNode, T>(node: &S) -> Frame<'_, S, T> {
191 Frame {
192 node,
193 next: 0,
194 results: Vec::with_capacity(node.child_count()),
195 }
196 }
197
198 let mut stack = vec![open(root)];
199 loop {
200 let top = stack.last_mut().expect("the root frame is popped last");
201 if top.next < top.node.child_count() {
202 let child = top.node.child(top.next);
203 top.next += 1;
204 stack.push(open(child));
205 continue;
206 }
207 let frame = stack.pop().expect("just observed");
208 let result = finish(frame.node, frame.results);
209 match stack.last_mut() {
210 Some(parent) => parent.results.push(result),
211 None => return result,
212 }
213 }
214}
215
216pub enum Visit<'a, S> {
218 Enter(&'a S),
220 Between(&'a S),
222 Exit(&'a S),
224}
225
226pub fn walk_tree<S, E>(
228 root: &S,
229 mut visit: impl FnMut(Visit<'_, S>) -> Result<(), E>,
230) -> Result<(), E>
231where
232 S: IndexedNode,
233{
234 let mut stack = vec![Visit::Enter(root)];
235 while let Some(step) = stack.pop() {
236 if let Visit::Enter(node) = step {
237 visit(Visit::Enter(node))?;
238 stack.push(Visit::Exit(node));
239 for i in (0..node.child_count()).rev() {
240 stack.push(Visit::Enter(node.child(i)));
241 if i > 0 {
242 stack.push(Visit::Between(node));
243 }
244 }
245 } else {
246 visit(step)?;
247 }
248 }
249 Ok(())
250}
251
252pub fn eq_tree<S>(left: &S, right: &S, same_node: impl Fn(&S, &S) -> bool) -> bool
255where
256 S: IndexedNode,
257{
258 let mut pending = vec![(left, right)];
259 while let Some((left, right)) = pending.pop() {
260 let count = left.child_count();
261 if !same_node(left, right) || count != right.child_count() {
262 return false;
263 }
264 pending.extend((0..count).map(|i| (left.child(i), right.child(i))));
265 }
266 true
267}
268
269pub fn cmp_tree<S>(left: &S, right: &S, cmp_node: impl Fn(&S, &S) -> Ordering) -> Ordering
273where
274 S: IndexedNode,
275{
276 enum Step<'a, S> {
277 Pair(&'a S, &'a S),
278 Counts(usize, usize),
280 }
281
282 let mut pending = vec![Step::Pair(left, right)];
283 while let Some(step) = pending.pop() {
284 let (left, right) = match step {
285 Step::Pair(left, right) => (left, right),
286 Step::Counts(left, right) => match left.cmp(&right) {
287 Ordering::Equal => continue,
288 ordering => return ordering,
289 },
290 };
291 match cmp_node(left, right) {
292 Ordering::Equal => {}
293 ordering => return ordering,
294 }
295 let counts = (left.child_count(), right.child_count());
296 pending.push(Step::Counts(counts.0, counts.1));
297 pending.extend(
298 (0..counts.0.min(counts.1))
299 .rev()
300 .map(|i| Step::Pair(left.child(i), right.child(i))),
301 );
302 }
303 Ordering::Equal
304}
305
306pub fn drop_children<S: TreeNode>(node: &mut S) {
310 let Some(children) = node.children_mut() else {
311 return;
312 };
313 let mut pending = std::mem::take(children);
314 while let Some(mut child) = pending.pop() {
315 if let Some(children) = child.children_mut() {
316 pending.append(children);
317 }
318 }
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324
325 #[derive(Debug, PartialEq)]
326 enum Node {
327 Leaf(u64),
328 List(Vec<Node>),
329 }
330
331 impl TreeNode for Node {
332 fn children(&self) -> &[Self] {
333 match self {
334 Node::List(children) => children,
335 Node::Leaf(_) => &[],
336 }
337 }
338
339 fn children_mut(&mut self) -> Option<&mut Vec<Self>> {
340 match self {
341 Node::List(children) => Some(children),
342 Node::Leaf(_) => None,
343 }
344 }
345 }
346
347 impl Drop for Node {
348 fn drop(&mut self) {
349 drop_children(self);
350 }
351 }
352
353 impl<'b, C> TreeDecode<'b, C> for Node {
354 type Builder = Node;
355
356 fn begin(d: &mut Decoder<'b>, _: &mut C) -> Result<(Node, Arity), Error> {
357 match d.datatype()? {
358 Type::Array | Type::ArrayIndef => Ok((Node::List(vec![]), d.array()?.into())),
359 _ => Ok((Node::Leaf(d.u64()?), Arity::Leaf)),
360 }
361 }
362
363 fn child(builder: &mut Node, child: Node) -> Result<(), Error> {
364 let Node::List(children) = builder else {
365 unreachable!()
366 };
367 children.push(child);
368 Ok(())
369 }
370
371 fn end(builder: Node, _: &mut Decoder<'b>, _: &mut C) -> Result<Node, Error> {
372 Ok(builder)
373 }
374 }
375
376 fn decode(bytes: &[u8]) -> Result<Node, Error> {
377 let mut d = Decoder::new(bytes);
378 let node = decode_tree(&mut d, &mut ())?;
379 assert_eq!(d.position(), bytes.len());
380 Ok(node)
381 }
382
383 #[test]
384 fn decodes_mixed_definite_and_indefinite_lists() {
385 let node = decode(&[0x84, 0x01, 0x80, 0x82, 0x02, 0x81, 0x03, 0x9f, 0xff]).unwrap();
387 let expected = Node::List(vec![
388 Node::Leaf(1),
389 Node::List(vec![]),
390 Node::List(vec![Node::Leaf(2), Node::List(vec![Node::Leaf(3)])]),
391 Node::List(vec![]),
392 ]);
393 assert_eq!(node, expected);
394 }
395
396 #[test]
397 fn rejects_truncated_input() {
398 assert!(decode(&[0x82, 0x01]).is_err());
399 assert!(decode(&[0x9f, 0x01]).is_err());
400 assert!(decode(&[0x81]).is_err());
401 }
402
403 #[test]
404 fn decodes_deep_nesting_on_a_small_stack() {
405 std::thread::Builder::new()
406 .stack_size(64 * 1024)
407 .spawn(|| {
408 let depth = 100_000;
409 let mut bytes = Vec::new();
410 for i in 0..depth {
411 bytes.push(if i % 2 == 0 { 0x81 } else { 0x9f });
412 }
413 bytes.push(0x00);
414 bytes.extend((0..depth).filter(|i| i % 2 == 1).map(|_| 0xff));
415 let mut cursor = &decode(&bytes).unwrap();
416 let mut seen = 0;
417 while let Node::List(children) = cursor {
418 assert_eq!(children.len(), 1);
419 cursor = &children[0];
420 seen += 1;
421 }
422 assert_eq!(seen, depth);
423 bytes.pop();
424 assert!(decode(&bytes).is_err());
425 })
426 .unwrap()
427 .join()
428 .unwrap();
429 }
430
431 fn mixed() -> Node {
432 Node::List(vec![
433 Node::Leaf(1),
434 Node::List(vec![]),
435 Node::List(vec![Node::Leaf(2), Node::List(vec![Node::Leaf(3)])]),
436 Node::List(vec![]),
437 ])
438 }
439
440 fn chain(depth: usize) -> Node {
441 let mut node = Node::Leaf(0);
442 for _ in 0..depth {
443 node = Node::List(vec![node]);
444 }
445 node
446 }
447
448 fn render(node: &Node) -> String {
449 let mut out = String::new();
450 walk_tree::<_, std::fmt::Error>(node, |visit| {
451 match visit {
452 Visit::Enter(Node::Leaf(n)) => out.push_str(&n.to_string()),
453 Visit::Enter(Node::List(_)) => out.push('['),
454 Visit::Between(_) => out.push(','),
455 Visit::Exit(Node::List(_)) => out.push(']'),
456 Visit::Exit(Node::Leaf(_)) => {}
457 }
458 Ok(())
459 })
460 .unwrap();
461 out
462 }
463
464 #[test]
465 fn walks_mixed_shapes_in_order() {
466 assert_eq!(render(&mixed()), "[1,[],[2,[3]],[]]");
467 assert_eq!(render(&Node::Leaf(7)), "7");
468 assert_eq!(render(&Node::List(vec![])), "[]");
469 }
470
471 #[test]
472 fn walk_propagates_errors() {
473 let result = walk_tree(&mixed(), |visit| match visit {
474 Visit::Enter(Node::Leaf(3)) => Err("three"),
475 _ => Ok(()),
476 });
477 assert_eq!(result, Err("three"));
478 }
479
480 #[test]
481 fn maps_mixed_shapes_positionally() {
482 #[derive(Debug, PartialEq)]
484 enum Target {
485 Leaf(u64),
486 List(Vec<Target>),
487 }
488 let mapped = map_tree(
489 &mixed(),
490 |node| match node {
491 Node::Leaf(n) => Target::Leaf(n * 2),
492 Node::List(_) => Target::List(vec![]),
493 },
494 |target| match target {
495 Target::List(children) => Some(children),
496 Target::Leaf(_) => None,
497 },
498 );
499 let expected = Target::List(vec![
500 Target::Leaf(2),
501 Target::List(vec![]),
502 Target::List(vec![Target::Leaf(4), Target::List(vec![Target::Leaf(6)])]),
503 Target::List(vec![]),
504 ]);
505 assert_eq!(mapped, expected);
506 }
507
508 #[test]
509 fn folds_children_in_order() {
510 let total = fold_tree(&mixed(), |node, children: Vec<u64>| match node {
511 Node::Leaf(n) => *n,
512 Node::List(_) => children.iter().sum(),
513 });
514 assert_eq!(total, 6);
515
516 let copy = fold_tree(&mixed(), |node, children| match node {
517 Node::Leaf(n) => Node::Leaf(*n),
518 Node::List(_) => Node::List(children),
519 });
520 assert_eq!(copy, mixed());
521 }
522
523 #[derive(Debug, PartialEq)]
525 enum Kv {
526 Leaf(u64),
527 Map(Vec<(Kv, Kv)>),
528 }
529
530 impl IndexedNode for Kv {
531 fn child_count(&self) -> usize {
532 match self {
533 Kv::Leaf(_) => 0,
534 Kv::Map(pairs) => pairs.len() * 2,
535 }
536 }
537
538 fn child(&self, index: usize) -> &Self {
539 let Kv::Map(pairs) = self else {
540 unreachable!("leaves have no children")
541 };
542 let (k, v) = &pairs[index / 2];
543 if index.is_multiple_of(2) { k } else { v }
544 }
545 }
546
547 fn kv_chain(depth: usize) -> Kv {
548 let mut node = Kv::Leaf(0);
549 for _ in 0..depth {
550 node = Kv::Map(vec![(Kv::Leaf(1), node)]);
551 }
552 node
553 }
554
555 fn render_kv(node: &Kv) -> String {
556 let mut out = String::new();
557 walk_tree::<_, std::fmt::Error>(node, |visit| {
558 match visit {
559 Visit::Enter(Kv::Leaf(n)) => out.push_str(&n.to_string()),
560 Visit::Enter(Kv::Map(_)) => out.push('{'),
561 Visit::Between(_) => out.push(','),
562 Visit::Exit(Kv::Map(_)) => out.push('}'),
563 Visit::Exit(Kv::Leaf(_)) => {}
564 }
565 Ok(())
566 })
567 .unwrap();
568 out
569 }
570
571 #[test]
572 fn indexed_children_interleave_keys_and_values() {
573 let node = Kv::Map(vec![
574 (Kv::Leaf(1), Kv::Leaf(2)),
575 (Kv::Leaf(3), Kv::Map(vec![(Kv::Leaf(4), Kv::Leaf(5))])),
576 ]);
577 assert_eq!(render_kv(&node), "{1,2,3,{4,5}}");
578
579 let copy = fold_tree(&node, |node, children| match node {
580 Kv::Leaf(n) => Kv::Leaf(*n),
581 Kv::Map(_) => {
582 let mut children = children.into_iter();
583 let mut pairs = Vec::new();
584 while let (Some(k), Some(v)) = (children.next(), children.next()) {
585 pairs.push((k, v));
586 }
587 Kv::Map(pairs)
588 }
589 });
590 assert_eq!(copy, node);
591
592 let same = |a: &Kv, b: &Kv| match (a, b) {
593 (Kv::Leaf(a), Kv::Leaf(b)) => a == b,
594 (Kv::Map(_), Kv::Map(_)) => true,
595 _ => false,
596 };
597 assert!(eq_tree(&node, ©, same));
598 assert!(!eq_tree(&node, &Kv::Map(vec![]), same));
599 assert!(!eq_tree(&kv_chain(3), &kv_chain(4), same));
600 }
601
602 #[test]
603 fn orders_like_a_derived_ord() {
604 let cmp = |a: &Node, b: &Node| match (a, b) {
605 (Node::Leaf(a), Node::Leaf(b)) => a.cmp(b),
606 (Node::Leaf(_), Node::List(_)) => Ordering::Less,
607 (Node::List(_), Node::Leaf(_)) => Ordering::Greater,
608 (Node::List(_), Node::List(_)) => Ordering::Equal,
609 };
610 let list = |xs: Vec<Node>| Node::List(xs);
611 let leaf = Node::Leaf;
612
613 assert_eq!(cmp_tree(&mixed(), &mixed(), cmp), Ordering::Equal);
614 assert_eq!(cmp_tree(&leaf(1), &leaf(2), cmp), Ordering::Less);
615 assert_eq!(cmp_tree(&leaf(1), &list(vec![]), cmp), Ordering::Less);
616 assert_eq!(
618 cmp_tree(&list(vec![leaf(2)]), &list(vec![leaf(1), leaf(9)]), cmp),
619 Ordering::Greater
620 );
621 assert_eq!(
622 cmp_tree(&list(vec![leaf(1)]), &list(vec![leaf(1), leaf(0)]), cmp),
623 Ordering::Less
624 );
625 assert_eq!(
627 cmp_tree(
628 &list(vec![list(vec![leaf(1)]), leaf(9)]),
629 &list(vec![list(vec![leaf(2)]), leaf(0)]),
630 cmp
631 ),
632 Ordering::Less
633 );
634 assert_eq!(cmp_tree(&chain(3), &chain(4), cmp), Ordering::Less);
635 }
636
637 #[test]
638 fn compares_structure_and_node_data() {
639 let same = |a: &Node, b: &Node| match (a, b) {
640 (Node::Leaf(a), Node::Leaf(b)) => a == b,
641 (Node::List(_), Node::List(_)) => true,
642 _ => false,
643 };
644 assert!(eq_tree(&mixed(), &mixed(), same));
645 assert!(!eq_tree(&mixed(), &Node::List(vec![]), same));
646 assert!(!eq_tree(&chain(3), &chain(4), same));
647 assert!(!eq_tree(&Node::Leaf(1), &Node::Leaf(2), same));
648 }
649
650 #[test]
651 fn traverses_deep_nesting_on_a_small_stack() {
652 std::thread::Builder::new()
653 .stack_size(64 * 1024)
654 .spawn(|| {
655 let depth = 100_000;
656 let node = chain(depth);
657 let text = render(&node);
658 assert_eq!(text, format!("{}0{}", "[".repeat(depth), "]".repeat(depth)));
659
660 let copy = map_tree(
661 &node,
662 |node| match node {
663 Node::Leaf(n) => Node::Leaf(*n),
664 Node::List(_) => Node::List(vec![]),
665 },
666 Node::children_mut,
667 );
668 assert!(eq_tree(&node, ©, |a, b| matches!(
669 (a, b),
670 (Node::Leaf(_), Node::Leaf(_)) | (Node::List(_), Node::List(_))
671 )));
672 let folded = fold_tree(&node, |node, children| match node {
673 Node::Leaf(n) => Node::Leaf(*n),
674 Node::List(_) => Node::List(children),
675 });
676 assert!(eq_tree(&node, &folded, |_, _| true));
677 assert_eq!(
678 cmp_tree(&node, &folded, |_, _| Ordering::Equal),
679 Ordering::Equal
680 );
681 assert_eq!(
682 cmp_tree(&node, &chain(depth - 1), |_, _| Ordering::Equal),
683 Ordering::Greater
684 );
685 drop(folded);
686 drop(copy);
687 drop(node);
688
689 let deep = std::mem::ManuallyDrop::new(kv_chain(depth));
692 let text = render_kv(&deep);
693 assert_eq!(
694 text,
695 format!("{}0{}", "{1,".repeat(depth), "}".repeat(depth))
696 );
697 let sum = fold_tree(&*deep, |node, children: Vec<u64>| match node {
698 Kv::Leaf(n) => *n,
699 Kv::Map(_) => children.iter().sum(),
700 });
701 assert_eq!(sum, depth as u64);
702 })
703 .unwrap()
704 .join()
705 .unwrap();
706 }
707}