pallas-codec 1.4.0

//! Stack-safe decoding and traversal of recursive structures.
//!
//! Recursive types such as native scripts or Plutus data can nest thousands
//! of levels deep inside a small payload, and any operation that recurses per
//! level exhausts the thread stack. Everything here drives a heap-backed
//! stack instead.
//!
//! - [`decode_tree`](crate::tree::decode_tree) builds a value from CBOR. A
//!   type opts in by implementing [`TreeDecode`](crate::tree::TreeDecode),
//!   which splits decoding a node into a header, a sequence of children and
//!   a footer.
//! - [`map_tree`](crate::tree::map_tree), [`walk_tree`](crate::tree::walk_tree),
//!   [`fold_tree`](crate::tree::fold_tree), [`eq_tree`](crate::tree::eq_tree),
//!   [`cmp_tree`](crate::tree::cmp_tree) and
//!   [`drop_children`](crate::tree::drop_children) operate on an existing
//!   value. A type opts in by implementing [`TreeNode`](crate::tree::TreeNode),
//!   which exposes its child list, or [`IndexedNode`](crate::tree::IndexedNode)
//!   when its children live elsewhere, such as in key-value pairs; they cover
//!   copying into another tree (clone, protobuf or JSON values), emitting a
//!   linear encoding, comparison and destruction.

use std::cmp::Ordering;

use minicbor::{Decoder, data::Type, decode::Error};

/// Number of children that follow a node's header in the input.
pub enum Arity {
    Leaf,
    Fixed(u64),
    /// Children continue until a CBOR break.
    Indefinite,
}

impl From<Option<u64>> for Arity {
    /// Converts the result of [`Decoder::array`] or [`Decoder::map`].
    fn from(len: Option<u64>) -> Self {
        match len {
            Some(n) => Self::Fixed(n),
            None => Self::Indefinite,
        }
    }
}

/// A recursive type whose nodes can be decoded without recursion.
///
/// The driver never reserves memory from a CBOR length header; builders grow
/// only as children are actually decoded from the input.
pub trait TreeDecode<'b, C>: Sized {
    /// Partially decoded node, accumulating children until [`end`].
    ///
    /// [`end`]: TreeDecode::end
    type Builder;

    /// Decode a node's header: everything up to and including the header of
    /// its child list.
    fn begin(d: &mut Decoder<'b>, ctx: &mut C) -> Result<(Self::Builder, Arity), Error>;

    /// Attach the next fully decoded child. Map-like nodes receive keys and
    /// values alternately.
    fn child(builder: &mut Self::Builder, child: Self) -> Result<(), Error>;

    /// Decode anything that follows the child list and finish the node.
    fn end(builder: Self::Builder, d: &mut Decoder<'b>, ctx: &mut C) -> Result<Self, Error>;
}

struct Frame<B> {
    builder: B,
    arity: Arity,
}

impl<B> Frame<B> {
    fn begin<'b, C, T>(d: &mut Decoder<'b>, ctx: &mut C) -> Result<Self, Error>
    where
        T: TreeDecode<'b, C, Builder = B>,
    {
        let (builder, arity) = T::begin(d, ctx)?;
        Ok(Self { builder, arity })
    }

    fn expects_child(&mut self, d: &mut Decoder<'_>) -> Result<bool, Error> {
        match self.arity {
            Arity::Leaf | Arity::Fixed(0) => Ok(false),
            Arity::Fixed(_) => Ok(true),
            Arity::Indefinite if d.datatype()? == Type::Break => {
                d.skip()?;
                self.arity = Arity::Leaf;
                Ok(false)
            }
            Arity::Indefinite => Ok(true),
        }
    }
}

/// Decode a [`TreeDecode`] value using a heap-backed stack of open nodes.
pub fn decode_tree<'b, C, T>(d: &mut Decoder<'b>, ctx: &mut C) -> Result<T, Error>
where
    T: TreeDecode<'b, C>,
{
    let mut parents: Vec<Frame<T::Builder>> = Vec::new();
    let mut current = Frame::begin::<C, T>(d, ctx)?;
    loop {
        if current.expects_child(d)? {
            parents.push(current);
            current = Frame::begin::<C, T>(d, ctx)?;
            continue;
        }
        let node = T::end(current.builder, d, ctx)?;
        let Some(mut parent) = parents.pop() else {
            return Ok(node);
        };
        T::child(&mut parent.builder, node)?;
        if let Arity::Fixed(remaining) = &mut parent.arity {
            *remaining -= 1;
        }
        current = parent;
    }
}

/// A recursive type whose children live in a `Vec`.
pub trait TreeNode: Sized {
    fn children(&self) -> &[Self];

    /// The child list, or `None` for leaf variants.
    fn children_mut(&mut self) -> Option<&mut Vec<Self>>;
}

/// A recursive type whose children are reachable by index, wherever they
/// are stored. Map-like nodes expose keys and values alternately.
///
/// Every [`TreeNode`] is an `IndexedNode`.
pub trait IndexedNode: Sized {
    fn child_count(&self) -> usize;

    /// The child at `index`, which is below [`child_count`](Self::child_count).
    fn child(&self, index: usize) -> &Self;
}

impl<T: TreeNode> IndexedNode for T {
    fn child_count(&self) -> usize {
        self.children().len()
    }

    fn child(&self, index: usize) -> &Self {
        &self.children()[index]
    }
}

/// Build a target tree from a source tree without recursion.
///
/// `shallow` maps one node to its target with an empty child list, and
/// `target_children` exposes that list so the driver can fill it. The target
/// needs no trait: any type with a `Vec` of children fits.
pub fn map_tree<S, T>(
    root: &S,
    shallow: impl Fn(&S) -> T,
    target_children: impl Fn(&mut T) -> Option<&mut Vec<T>>,
) -> T
where
    S: TreeNode,
{
    let mut target_root = shallow(root);
    let mut pending = vec![(root, &mut target_root)];
    while let Some((source, target)) = pending.pop() {
        let Some(children) = target_children(target) else {
            continue;
        };
        let source_children = source.children();
        *children = source_children.iter().map(&shallow).collect();
        pending.extend(source_children.iter().zip(children.iter_mut()));
    }
    target_root
}

/// Post-order fold without recursion. `finish` receives each node with the
/// results of its children in order, an empty list for a leaf, and returns
/// the node's own result.
///
/// Unlike [`map_tree`] this builds bottom-up, so it fits targets whose
/// children are not a plain `Vec`, such as key-value pairs.
pub fn fold_tree<S, T>(root: &S, mut finish: impl FnMut(&S, Vec<T>) -> T) -> T
where
    S: IndexedNode,
{
    struct Frame<'a, S, T> {
        node: &'a S,
        next: usize,
        results: Vec<T>,
    }

    fn open<S: IndexedNode, T>(node: &S) -> Frame<'_, S, T> {
        Frame {
            node,
            next: 0,
            results: Vec::with_capacity(node.child_count()),
        }
    }

    let mut stack = vec![open(root)];
    loop {
        let top = stack.last_mut().expect("the root frame is popped last");
        if top.next < top.node.child_count() {
            let child = top.node.child(top.next);
            top.next += 1;
            stack.push(open(child));
            continue;
        }
        let frame = stack.pop().expect("just observed");
        let result = finish(frame.node, frame.results);
        match stack.last_mut() {
            Some(parent) => parent.results.push(result),
            None => return result,
        }
    }
}

/// One step of a [`walk_tree`] traversal.
pub enum Visit<'a, S> {
    /// A node, before any of its children.
    Enter(&'a S),
    /// The parent, before each of its children but the first.
    Between(&'a S),
    /// A node, after all of its children.
    Exit(&'a S),
}

/// Pre-order traversal without recursion, for encoders and renderers.
pub fn walk_tree<S, E>(
    root: &S,
    mut visit: impl FnMut(Visit<'_, S>) -> Result<(), E>,
) -> Result<(), E>
where
    S: IndexedNode,
{
    let mut stack = vec![Visit::Enter(root)];
    while let Some(step) = stack.pop() {
        if let Visit::Enter(node) = step {
            visit(Visit::Enter(node))?;
            stack.push(Visit::Exit(node));
            for i in (0..node.child_count()).rev() {
                stack.push(Visit::Enter(node.child(i)));
                if i > 0 {
                    stack.push(Visit::Between(node));
                }
            }
        } else {
            visit(step)?;
        }
    }
    Ok(())
}

/// Structural equality without recursion. `same_node` compares two nodes'
/// own data, ignoring their children.
pub fn eq_tree<S>(left: &S, right: &S, same_node: impl Fn(&S, &S) -> bool) -> bool
where
    S: IndexedNode,
{
    let mut pending = vec![(left, right)];
    while let Some((left, right)) = pending.pop() {
        let count = left.child_count();
        if !same_node(left, right) || count != right.child_count() {
            return false;
        }
        pending.extend((0..count).map(|i| (left.child(i), right.child(i))));
    }
    true
}

/// Lexicographic ordering without recursion, as a derived `Ord` over `Vec`
/// children would produce: `cmp_node` compares two nodes' own data, then the
/// children pairwise in order, then the child counts.
pub fn cmp_tree<S>(left: &S, right: &S, cmp_node: impl Fn(&S, &S) -> Ordering) -> Ordering
where
    S: IndexedNode,
{
    enum Step<'a, S> {
        Pair(&'a S, &'a S),
        /// Child counts, decided once every shared child compared equal.
        Counts(usize, usize),
    }

    let mut pending = vec![Step::Pair(left, right)];
    while let Some(step) = pending.pop() {
        let (left, right) = match step {
            Step::Pair(left, right) => (left, right),
            Step::Counts(left, right) => match left.cmp(&right) {
                Ordering::Equal => continue,
                ordering => return ordering,
            },
        };
        match cmp_node(left, right) {
            Ordering::Equal => {}
            ordering => return ordering,
        }
        let counts = (left.child_count(), right.child_count());
        pending.push(Step::Counts(counts.0, counts.1));
        pending.extend(
            (0..counts.0.min(counts.1))
                .rev()
                .map(|i| Step::Pair(left.child(i), right.child(i))),
        );
    }
    Ordering::Equal
}

/// Detach and destroy a node's descendants without recursion. Call from a
/// `Drop` impl; the node's own drop then has no children left to recurse
/// into.
pub fn drop_children<S: TreeNode>(node: &mut S) {
    let Some(children) = node.children_mut() else {
        return;
    };
    let mut pending = std::mem::take(children);
    while let Some(mut child) = pending.pop() {
        if let Some(children) = child.children_mut() {
            pending.append(children);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Debug, PartialEq)]
    enum Node {
        Leaf(u64),
        List(Vec<Node>),
    }

    impl TreeNode for Node {
        fn children(&self) -> &[Self] {
            match self {
                Node::List(children) => children,
                Node::Leaf(_) => &[],
            }
        }

        fn children_mut(&mut self) -> Option<&mut Vec<Self>> {
            match self {
                Node::List(children) => Some(children),
                Node::Leaf(_) => None,
            }
        }
    }

    impl Drop for Node {
        fn drop(&mut self) {
            drop_children(self);
        }
    }

    impl<'b, C> TreeDecode<'b, C> for Node {
        type Builder = Node;

        fn begin(d: &mut Decoder<'b>, _: &mut C) -> Result<(Node, Arity), Error> {
            match d.datatype()? {
                Type::Array | Type::ArrayIndef => Ok((Node::List(vec![]), d.array()?.into())),
                _ => Ok((Node::Leaf(d.u64()?), Arity::Leaf)),
            }
        }

        fn child(builder: &mut Node, child: Node) -> Result<(), Error> {
            let Node::List(children) = builder else {
                unreachable!()
            };
            children.push(child);
            Ok(())
        }

        fn end(builder: Node, _: &mut Decoder<'b>, _: &mut C) -> Result<Node, Error> {
            Ok(builder)
        }
    }

    fn decode(bytes: &[u8]) -> Result<Node, Error> {
        let mut d = Decoder::new(bytes);
        let node = decode_tree(&mut d, &mut ())?;
        assert_eq!(d.position(), bytes.len());
        Ok(node)
    }

    #[test]
    fn decodes_mixed_definite_and_indefinite_lists() {
        // [1, [], [2, [3]], []_]
        let node = decode(&[0x84, 0x01, 0x80, 0x82, 0x02, 0x81, 0x03, 0x9f, 0xff]).unwrap();
        let expected = Node::List(vec![
            Node::Leaf(1),
            Node::List(vec![]),
            Node::List(vec![Node::Leaf(2), Node::List(vec![Node::Leaf(3)])]),
            Node::List(vec![]),
        ]);
        assert_eq!(node, expected);
    }

    #[test]
    fn rejects_truncated_input() {
        assert!(decode(&[0x82, 0x01]).is_err());
        assert!(decode(&[0x9f, 0x01]).is_err());
        assert!(decode(&[0x81]).is_err());
    }

    #[test]
    fn decodes_deep_nesting_on_a_small_stack() {
        std::thread::Builder::new()
            .stack_size(64 * 1024)
            .spawn(|| {
                let depth = 100_000;
                let mut bytes = Vec::new();
                for i in 0..depth {
                    bytes.push(if i % 2 == 0 { 0x81 } else { 0x9f });
                }
                bytes.push(0x00);
                bytes.extend((0..depth).filter(|i| i % 2 == 1).map(|_| 0xff));
                let mut cursor = &decode(&bytes).unwrap();
                let mut seen = 0;
                while let Node::List(children) = cursor {
                    assert_eq!(children.len(), 1);
                    cursor = &children[0];
                    seen += 1;
                }
                assert_eq!(seen, depth);
                bytes.pop();
                assert!(decode(&bytes).is_err());
            })
            .unwrap()
            .join()
            .unwrap();
    }

    fn mixed() -> Node {
        Node::List(vec![
            Node::Leaf(1),
            Node::List(vec![]),
            Node::List(vec![Node::Leaf(2), Node::List(vec![Node::Leaf(3)])]),
            Node::List(vec![]),
        ])
    }

    fn chain(depth: usize) -> Node {
        let mut node = Node::Leaf(0);
        for _ in 0..depth {
            node = Node::List(vec![node]);
        }
        node
    }

    fn render(node: &Node) -> String {
        let mut out = String::new();
        walk_tree::<_, std::fmt::Error>(node, |visit| {
            match visit {
                Visit::Enter(Node::Leaf(n)) => out.push_str(&n.to_string()),
                Visit::Enter(Node::List(_)) => out.push('['),
                Visit::Between(_) => out.push(','),
                Visit::Exit(Node::List(_)) => out.push(']'),
                Visit::Exit(Node::Leaf(_)) => {}
            }
            Ok(())
        })
        .unwrap();
        out
    }

    #[test]
    fn walks_mixed_shapes_in_order() {
        assert_eq!(render(&mixed()), "[1,[],[2,[3]],[]]");
        assert_eq!(render(&Node::Leaf(7)), "7");
        assert_eq!(render(&Node::List(vec![])), "[]");
    }

    #[test]
    fn walk_propagates_errors() {
        let result = walk_tree(&mixed(), |visit| match visit {
            Visit::Enter(Node::Leaf(3)) => Err("three"),
            _ => Ok(()),
        });
        assert_eq!(result, Err("three"));
    }

    #[test]
    fn maps_mixed_shapes_positionally() {
        // Same shape, leaves doubled, into an unrelated target type.
        #[derive(Debug, PartialEq)]
        enum Target {
            Leaf(u64),
            List(Vec<Target>),
        }
        let mapped = map_tree(
            &mixed(),
            |node| match node {
                Node::Leaf(n) => Target::Leaf(n * 2),
                Node::List(_) => Target::List(vec![]),
            },
            |target| match target {
                Target::List(children) => Some(children),
                Target::Leaf(_) => None,
            },
        );
        let expected = Target::List(vec![
            Target::Leaf(2),
            Target::List(vec![]),
            Target::List(vec![Target::Leaf(4), Target::List(vec![Target::Leaf(6)])]),
            Target::List(vec![]),
        ]);
        assert_eq!(mapped, expected);
    }

    #[test]
    fn folds_children_in_order() {
        let total = fold_tree(&mixed(), |node, children: Vec<u64>| match node {
            Node::Leaf(n) => *n,
            Node::List(_) => children.iter().sum(),
        });
        assert_eq!(total, 6);

        let copy = fold_tree(&mixed(), |node, children| match node {
            Node::Leaf(n) => Node::Leaf(*n),
            Node::List(_) => Node::List(children),
        });
        assert_eq!(copy, mixed());
    }

    /// Children held in pairs rather than a `Vec`, as a map-like type has.
    #[derive(Debug, PartialEq)]
    enum Kv {
        Leaf(u64),
        Map(Vec<(Kv, Kv)>),
    }

    impl IndexedNode for Kv {
        fn child_count(&self) -> usize {
            match self {
                Kv::Leaf(_) => 0,
                Kv::Map(pairs) => pairs.len() * 2,
            }
        }

        fn child(&self, index: usize) -> &Self {
            let Kv::Map(pairs) = self else {
                unreachable!("leaves have no children")
            };
            let (k, v) = &pairs[index / 2];
            if index.is_multiple_of(2) { k } else { v }
        }
    }

    fn kv_chain(depth: usize) -> Kv {
        let mut node = Kv::Leaf(0);
        for _ in 0..depth {
            node = Kv::Map(vec![(Kv::Leaf(1), node)]);
        }
        node
    }

    fn render_kv(node: &Kv) -> String {
        let mut out = String::new();
        walk_tree::<_, std::fmt::Error>(node, |visit| {
            match visit {
                Visit::Enter(Kv::Leaf(n)) => out.push_str(&n.to_string()),
                Visit::Enter(Kv::Map(_)) => out.push('{'),
                Visit::Between(_) => out.push(','),
                Visit::Exit(Kv::Map(_)) => out.push('}'),
                Visit::Exit(Kv::Leaf(_)) => {}
            }
            Ok(())
        })
        .unwrap();
        out
    }

    #[test]
    fn indexed_children_interleave_keys_and_values() {
        let node = Kv::Map(vec![
            (Kv::Leaf(1), Kv::Leaf(2)),
            (Kv::Leaf(3), Kv::Map(vec![(Kv::Leaf(4), Kv::Leaf(5))])),
        ]);
        assert_eq!(render_kv(&node), "{1,2,3,{4,5}}");

        let copy = fold_tree(&node, |node, children| match node {
            Kv::Leaf(n) => Kv::Leaf(*n),
            Kv::Map(_) => {
                let mut children = children.into_iter();
                let mut pairs = Vec::new();
                while let (Some(k), Some(v)) = (children.next(), children.next()) {
                    pairs.push((k, v));
                }
                Kv::Map(pairs)
            }
        });
        assert_eq!(copy, node);

        let same = |a: &Kv, b: &Kv| match (a, b) {
            (Kv::Leaf(a), Kv::Leaf(b)) => a == b,
            (Kv::Map(_), Kv::Map(_)) => true,
            _ => false,
        };
        assert!(eq_tree(&node, &copy, same));
        assert!(!eq_tree(&node, &Kv::Map(vec![]), same));
        assert!(!eq_tree(&kv_chain(3), &kv_chain(4), same));
    }

    #[test]
    fn orders_like_a_derived_ord() {
        let cmp = |a: &Node, b: &Node| match (a, b) {
            (Node::Leaf(a), Node::Leaf(b)) => a.cmp(b),
            (Node::Leaf(_), Node::List(_)) => Ordering::Less,
            (Node::List(_), Node::Leaf(_)) => Ordering::Greater,
            (Node::List(_), Node::List(_)) => Ordering::Equal,
        };
        let list = |xs: Vec<Node>| Node::List(xs);
        let leaf = Node::Leaf;

        assert_eq!(cmp_tree(&mixed(), &mixed(), cmp), Ordering::Equal);
        assert_eq!(cmp_tree(&leaf(1), &leaf(2), cmp), Ordering::Less);
        assert_eq!(cmp_tree(&leaf(1), &list(vec![]), cmp), Ordering::Less);
        // A shared prefix decides before the length does.
        assert_eq!(
            cmp_tree(&list(vec![leaf(2)]), &list(vec![leaf(1), leaf(9)]), cmp),
            Ordering::Greater
        );
        assert_eq!(
            cmp_tree(&list(vec![leaf(1)]), &list(vec![leaf(1), leaf(0)]), cmp),
            Ordering::Less
        );
        // A nested difference is found before a later sibling.
        assert_eq!(
            cmp_tree(
                &list(vec![list(vec![leaf(1)]), leaf(9)]),
                &list(vec![list(vec![leaf(2)]), leaf(0)]),
                cmp
            ),
            Ordering::Less
        );
        assert_eq!(cmp_tree(&chain(3), &chain(4), cmp), Ordering::Less);
    }

    #[test]
    fn compares_structure_and_node_data() {
        let same = |a: &Node, b: &Node| match (a, b) {
            (Node::Leaf(a), Node::Leaf(b)) => a == b,
            (Node::List(_), Node::List(_)) => true,
            _ => false,
        };
        assert!(eq_tree(&mixed(), &mixed(), same));
        assert!(!eq_tree(&mixed(), &Node::List(vec![]), same));
        assert!(!eq_tree(&chain(3), &chain(4), same));
        assert!(!eq_tree(&Node::Leaf(1), &Node::Leaf(2), same));
    }

    #[test]
    fn traverses_deep_nesting_on_a_small_stack() {
        std::thread::Builder::new()
            .stack_size(64 * 1024)
            .spawn(|| {
                let depth = 100_000;
                let node = chain(depth);
                let text = render(&node);
                assert_eq!(text, format!("{}0{}", "[".repeat(depth), "]".repeat(depth)));

                let copy = map_tree(
                    &node,
                    |node| match node {
                        Node::Leaf(n) => Node::Leaf(*n),
                        Node::List(_) => Node::List(vec![]),
                    },
                    Node::children_mut,
                );
                assert!(eq_tree(&node, &copy, |a, b| matches!(
                    (a, b),
                    (Node::Leaf(_), Node::Leaf(_)) | (Node::List(_), Node::List(_))
                )));
                let folded = fold_tree(&node, |node, children| match node {
                    Node::Leaf(n) => Node::Leaf(*n),
                    Node::List(_) => Node::List(children),
                });
                assert!(eq_tree(&node, &folded, |_, _| true));
                assert_eq!(
                    cmp_tree(&node, &folded, |_, _| Ordering::Equal),
                    Ordering::Equal
                );
                assert_eq!(
                    cmp_tree(&node, &chain(depth - 1), |_, _| Ordering::Equal),
                    Ordering::Greater
                );
                drop(folded);
                drop(copy);
                drop(node);

                // Kv has no iterative Drop, so only the traversals are under
                // test here; leak the values rather than unwind them.
                let deep = std::mem::ManuallyDrop::new(kv_chain(depth));
                let text = render_kv(&deep);
                assert_eq!(
                    text,
                    format!("{}0{}", "{1,".repeat(depth), "}".repeat(depth))
                );
                let sum = fold_tree(&*deep, |node, children: Vec<u64>| match node {
                    Kv::Leaf(n) => *n,
                    Kv::Map(_) => children.iter().sum(),
                });
                assert_eq!(sum, depth as u64);
            })
            .unwrap()
            .join()
            .unwrap();
    }
}