pallas_codec/
tree.rs

1//! Stack-safe decoding and traversal of recursive structures.
2//!
3//! Recursive types such as native scripts or Plutus data can nest thousands
4//! of levels deep inside a small payload, and any operation that recurses per
5//! level exhausts the thread stack. Everything here drives a heap-backed
6//! stack instead.
7//!
8//! - [`decode_tree`](crate::tree::decode_tree) builds a value from CBOR. A
9//!   type opts in by implementing [`TreeDecode`](crate::tree::TreeDecode),
10//!   which splits decoding a node into a header, a sequence of children and
11//!   a footer.
12//! - [`map_tree`](crate::tree::map_tree), [`walk_tree`](crate::tree::walk_tree),
13//!   [`eq_tree`](crate::tree::eq_tree) and
14//!   [`drop_children`](crate::tree::drop_children) operate on an existing
15//!   value. A type opts in by implementing [`TreeNode`](crate::tree::TreeNode),
16//!   which exposes its child list; they cover copying into another tree
17//!   (clone, protobuf or JSON values), emitting a linear encoding, comparison
18//!   and destruction.
19
20use minicbor::{Decoder, data::Type, decode::Error};
21
22/// Number of children that follow a node's header in the input.
23pub enum Arity {
24    Leaf,
25    Fixed(u64),
26    /// Children continue until a CBOR break.
27    Indefinite,
28}
29
30impl From<Option<u64>> for Arity {
31    /// Converts the result of [`Decoder::array`] or [`Decoder::map`].
32    fn from(len: Option<u64>) -> Self {
33        match len {
34            Some(n) => Self::Fixed(n),
35            None => Self::Indefinite,
36        }
37    }
38}
39
40/// A recursive type whose nodes can be decoded without recursion.
41///
42/// The driver never reserves memory from a CBOR length header; builders grow
43/// only as children are actually decoded from the input.
44pub trait TreeDecode<'b, C>: Sized {
45    /// Partially decoded node, accumulating children until [`end`].
46    ///
47    /// [`end`]: TreeDecode::end
48    type Builder;
49
50    /// Decode a node's header: everything up to and including the header of
51    /// its child list.
52    fn begin(d: &mut Decoder<'b>, ctx: &mut C) -> Result<(Self::Builder, Arity), Error>;
53
54    /// Attach the next fully decoded child. Map-like nodes receive keys and
55    /// values alternately.
56    fn child(builder: &mut Self::Builder, child: Self) -> Result<(), Error>;
57
58    /// Decode anything that follows the child list and finish the node.
59    fn end(builder: Self::Builder, d: &mut Decoder<'b>, ctx: &mut C) -> Result<Self, Error>;
60}
61
62struct Frame<B> {
63    builder: B,
64    arity: Arity,
65}
66
67impl<B> Frame<B> {
68    fn begin<'b, C, T>(d: &mut Decoder<'b>, ctx: &mut C) -> Result<Self, Error>
69    where
70        T: TreeDecode<'b, C, Builder = B>,
71    {
72        let (builder, arity) = T::begin(d, ctx)?;
73        Ok(Self { builder, arity })
74    }
75
76    fn expects_child(&mut self, d: &mut Decoder<'_>) -> Result<bool, Error> {
77        match self.arity {
78            Arity::Leaf | Arity::Fixed(0) => Ok(false),
79            Arity::Fixed(_) => Ok(true),
80            Arity::Indefinite if d.datatype()? == Type::Break => {
81                d.skip()?;
82                self.arity = Arity::Leaf;
83                Ok(false)
84            }
85            Arity::Indefinite => Ok(true),
86        }
87    }
88}
89
90/// Decode a [`TreeDecode`] value using a heap-backed stack of open nodes.
91pub fn decode_tree<'b, C, T>(d: &mut Decoder<'b>, ctx: &mut C) -> Result<T, Error>
92where
93    T: TreeDecode<'b, C>,
94{
95    let mut parents: Vec<Frame<T::Builder>> = Vec::new();
96    let mut current = Frame::begin::<C, T>(d, ctx)?;
97    loop {
98        if current.expects_child(d)? {
99            parents.push(current);
100            current = Frame::begin::<C, T>(d, ctx)?;
101            continue;
102        }
103        let node = T::end(current.builder, d, ctx)?;
104        let Some(mut parent) = parents.pop() else {
105            return Ok(node);
106        };
107        T::child(&mut parent.builder, node)?;
108        if let Arity::Fixed(remaining) = &mut parent.arity {
109            *remaining -= 1;
110        }
111        current = parent;
112    }
113}
114
115/// A recursive type whose children live in a `Vec`.
116pub trait TreeNode: Sized {
117    fn children(&self) -> &[Self];
118
119    /// The child list, or `None` for leaf variants.
120    fn children_mut(&mut self) -> Option<&mut Vec<Self>>;
121}
122
123/// Build a target tree from a source tree without recursion.
124///
125/// `shallow` maps one node to its target with an empty child list, and
126/// `target_children` exposes that list so the driver can fill it. The target
127/// needs no trait: any type with a `Vec` of children fits.
128pub fn map_tree<S, T>(
129    root: &S,
130    shallow: impl Fn(&S) -> T,
131    target_children: impl Fn(&mut T) -> Option<&mut Vec<T>>,
132) -> T
133where
134    S: TreeNode,
135{
136    let mut target_root = shallow(root);
137    let mut pending = vec![(root, &mut target_root)];
138    while let Some((source, target)) = pending.pop() {
139        let Some(children) = target_children(target) else {
140            continue;
141        };
142        let source_children = source.children();
143        *children = source_children.iter().map(&shallow).collect();
144        pending.extend(source_children.iter().zip(children.iter_mut()));
145    }
146    target_root
147}
148
149/// One step of a [`walk_tree`] traversal.
150pub enum Visit<'a, S> {
151    /// A node, before any of its children.
152    Enter(&'a S),
153    /// The parent, before each of its children but the first.
154    Between(&'a S),
155    /// A node, after all of its children.
156    Exit(&'a S),
157}
158
159/// Pre-order traversal without recursion, for encoders and renderers.
160pub fn walk_tree<S, E>(
161    root: &S,
162    mut visit: impl FnMut(Visit<'_, S>) -> Result<(), E>,
163) -> Result<(), E>
164where
165    S: TreeNode,
166{
167    let mut stack = vec![Visit::Enter(root)];
168    while let Some(step) = stack.pop() {
169        if let Visit::Enter(node) = step {
170            visit(Visit::Enter(node))?;
171            stack.push(Visit::Exit(node));
172            for (i, child) in node.children().iter().enumerate().rev() {
173                stack.push(Visit::Enter(child));
174                if i > 0 {
175                    stack.push(Visit::Between(node));
176                }
177            }
178        } else {
179            visit(step)?;
180        }
181    }
182    Ok(())
183}
184
185/// Structural equality without recursion. `same_node` compares two nodes'
186/// own data, ignoring their children.
187pub fn eq_tree<S>(left: &S, right: &S, same_node: impl Fn(&S, &S) -> bool) -> bool
188where
189    S: TreeNode,
190{
191    let mut pending = vec![(left, right)];
192    while let Some((left, right)) = pending.pop() {
193        if !same_node(left, right) || left.children().len() != right.children().len() {
194            return false;
195        }
196        pending.extend(left.children().iter().zip(right.children()));
197    }
198    true
199}
200
201/// Detach and destroy a node's descendants without recursion. Call from a
202/// `Drop` impl; the node's own drop then has no children left to recurse
203/// into.
204pub fn drop_children<S: TreeNode>(node: &mut S) {
205    let Some(children) = node.children_mut() else {
206        return;
207    };
208    let mut pending = std::mem::take(children);
209    while let Some(mut child) = pending.pop() {
210        if let Some(children) = child.children_mut() {
211            pending.append(children);
212        }
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[derive(Debug, PartialEq)]
221    enum Node {
222        Leaf(u64),
223        List(Vec<Node>),
224    }
225
226    impl TreeNode for Node {
227        fn children(&self) -> &[Self] {
228            match self {
229                Node::List(children) => children,
230                Node::Leaf(_) => &[],
231            }
232        }
233
234        fn children_mut(&mut self) -> Option<&mut Vec<Self>> {
235            match self {
236                Node::List(children) => Some(children),
237                Node::Leaf(_) => None,
238            }
239        }
240    }
241
242    impl Drop for Node {
243        fn drop(&mut self) {
244            drop_children(self);
245        }
246    }
247
248    impl<'b, C> TreeDecode<'b, C> for Node {
249        type Builder = Node;
250
251        fn begin(d: &mut Decoder<'b>, _: &mut C) -> Result<(Node, Arity), Error> {
252            match d.datatype()? {
253                Type::Array | Type::ArrayIndef => Ok((Node::List(vec![]), d.array()?.into())),
254                _ => Ok((Node::Leaf(d.u64()?), Arity::Leaf)),
255            }
256        }
257
258        fn child(builder: &mut Node, child: Node) -> Result<(), Error> {
259            let Node::List(children) = builder else {
260                unreachable!()
261            };
262            children.push(child);
263            Ok(())
264        }
265
266        fn end(builder: Node, _: &mut Decoder<'b>, _: &mut C) -> Result<Node, Error> {
267            Ok(builder)
268        }
269    }
270
271    fn decode(bytes: &[u8]) -> Result<Node, Error> {
272        let mut d = Decoder::new(bytes);
273        let node = decode_tree(&mut d, &mut ())?;
274        assert_eq!(d.position(), bytes.len());
275        Ok(node)
276    }
277
278    #[test]
279    fn decodes_mixed_definite_and_indefinite_lists() {
280        // [1, [], [2, [3]], []_]
281        let node = decode(&[0x84, 0x01, 0x80, 0x82, 0x02, 0x81, 0x03, 0x9f, 0xff]).unwrap();
282        let expected = Node::List(vec![
283            Node::Leaf(1),
284            Node::List(vec![]),
285            Node::List(vec![Node::Leaf(2), Node::List(vec![Node::Leaf(3)])]),
286            Node::List(vec![]),
287        ]);
288        assert_eq!(node, expected);
289    }
290
291    #[test]
292    fn rejects_truncated_input() {
293        assert!(decode(&[0x82, 0x01]).is_err());
294        assert!(decode(&[0x9f, 0x01]).is_err());
295        assert!(decode(&[0x81]).is_err());
296    }
297
298    #[test]
299    fn decodes_deep_nesting_on_a_small_stack() {
300        std::thread::Builder::new()
301            .stack_size(64 * 1024)
302            .spawn(|| {
303                let depth = 100_000;
304                let mut bytes = Vec::new();
305                for i in 0..depth {
306                    bytes.push(if i % 2 == 0 { 0x81 } else { 0x9f });
307                }
308                bytes.push(0x00);
309                bytes.extend((0..depth).filter(|i| i % 2 == 1).map(|_| 0xff));
310                let mut cursor = &decode(&bytes).unwrap();
311                let mut seen = 0;
312                while let Node::List(children) = cursor {
313                    assert_eq!(children.len(), 1);
314                    cursor = &children[0];
315                    seen += 1;
316                }
317                assert_eq!(seen, depth);
318                bytes.pop();
319                assert!(decode(&bytes).is_err());
320            })
321            .unwrap()
322            .join()
323            .unwrap();
324    }
325
326    fn mixed() -> Node {
327        Node::List(vec![
328            Node::Leaf(1),
329            Node::List(vec![]),
330            Node::List(vec![Node::Leaf(2), Node::List(vec![Node::Leaf(3)])]),
331            Node::List(vec![]),
332        ])
333    }
334
335    fn chain(depth: usize) -> Node {
336        let mut node = Node::Leaf(0);
337        for _ in 0..depth {
338            node = Node::List(vec![node]);
339        }
340        node
341    }
342
343    fn render(node: &Node) -> String {
344        let mut out = String::new();
345        walk_tree::<_, std::fmt::Error>(node, |visit| {
346            match visit {
347                Visit::Enter(Node::Leaf(n)) => out.push_str(&n.to_string()),
348                Visit::Enter(Node::List(_)) => out.push('['),
349                Visit::Between(_) => out.push(','),
350                Visit::Exit(Node::List(_)) => out.push(']'),
351                Visit::Exit(Node::Leaf(_)) => {}
352            }
353            Ok(())
354        })
355        .unwrap();
356        out
357    }
358
359    #[test]
360    fn walks_mixed_shapes_in_order() {
361        assert_eq!(render(&mixed()), "[1,[],[2,[3]],[]]");
362        assert_eq!(render(&Node::Leaf(7)), "7");
363        assert_eq!(render(&Node::List(vec![])), "[]");
364    }
365
366    #[test]
367    fn walk_propagates_errors() {
368        let result = walk_tree(&mixed(), |visit| match visit {
369            Visit::Enter(Node::Leaf(3)) => Err("three"),
370            _ => Ok(()),
371        });
372        assert_eq!(result, Err("three"));
373    }
374
375    #[test]
376    fn maps_mixed_shapes_positionally() {
377        // Same shape, leaves doubled, into an unrelated target type.
378        #[derive(Debug, PartialEq)]
379        enum Target {
380            Leaf(u64),
381            List(Vec<Target>),
382        }
383        let mapped = map_tree(
384            &mixed(),
385            |node| match node {
386                Node::Leaf(n) => Target::Leaf(n * 2),
387                Node::List(_) => Target::List(vec![]),
388            },
389            |target| match target {
390                Target::List(children) => Some(children),
391                Target::Leaf(_) => None,
392            },
393        );
394        let expected = Target::List(vec![
395            Target::Leaf(2),
396            Target::List(vec![]),
397            Target::List(vec![Target::Leaf(4), Target::List(vec![Target::Leaf(6)])]),
398            Target::List(vec![]),
399        ]);
400        assert_eq!(mapped, expected);
401    }
402
403    #[test]
404    fn compares_structure_and_node_data() {
405        let same = |a: &Node, b: &Node| match (a, b) {
406            (Node::Leaf(a), Node::Leaf(b)) => a == b,
407            (Node::List(_), Node::List(_)) => true,
408            _ => false,
409        };
410        assert!(eq_tree(&mixed(), &mixed(), same));
411        assert!(!eq_tree(&mixed(), &Node::List(vec![]), same));
412        assert!(!eq_tree(&chain(3), &chain(4), same));
413        assert!(!eq_tree(&Node::Leaf(1), &Node::Leaf(2), same));
414    }
415
416    #[test]
417    fn traverses_deep_nesting_on_a_small_stack() {
418        std::thread::Builder::new()
419            .stack_size(64 * 1024)
420            .spawn(|| {
421                let depth = 100_000;
422                let node = chain(depth);
423                let text = render(&node);
424                assert_eq!(text, format!("{}0{}", "[".repeat(depth), "]".repeat(depth)));
425
426                let copy = map_tree(
427                    &node,
428                    |node| match node {
429                        Node::Leaf(n) => Node::Leaf(*n),
430                        Node::List(_) => Node::List(vec![]),
431                    },
432                    Node::children_mut,
433                );
434                assert!(eq_tree(&node, &copy, |a, b| matches!(
435                    (a, b),
436                    (Node::Leaf(_), Node::Leaf(_)) | (Node::List(_), Node::List(_))
437                )));
438                drop(copy);
439                drop(node);
440            })
441            .unwrap()
442            .join()
443            .unwrap();
444    }
445}
pallas_codec/tree.rs

pallas_codec/
tree.rs
