monorepo/tree/src/lib.rs

//! This data structure is a generic tree which can contain any data. That data itself need to keep
//! track of its own tree structure.
//!
//! This surely already exists. I am created it to test my own ability to do things in Rust.

use std::{
    cell::{Ref, RefCell},
    collections::VecDeque,
    rc::Rc,
};

// I need to take what I learned about linked lists and about the other Tree data structure, and
// apply it here with arena allocation.
//
// Also, smarter node allocation and pointer handling in order to avoid clones.
#[derive(Clone, Debug, Default)]
pub enum Tree<T> {
    #[default]
    Empty,
    Root(Node<T>),
}

impl<T> Tree<T> {
    pub fn new(value: T) -> (Tree<T>, Node<T>) {
        let node = Node::new(value);
        let tree = Tree::Root(node.clone());
        (tree, node)
    }

    pub fn set_value(&mut self, value: T) -> Node<T> {
        let node = Node::new(value);
        *self = Tree::Root(node.clone());
        node
    }

    /// Use a breadth-first-search pattern to find a node, returning the node if found.
    pub fn find_bfs<F>(&self, op: F) -> Option<Node<T>>
    where
        F: FnOnce(&T) -> bool + Copy,
    {
        let mut queue: VecDeque<Node<T>> = match self {
            Tree::Empty => VecDeque::new(),
            Tree::Root(node) => {
                let mut queue = VecDeque::new();
                queue.push_back(node.clone());
                queue
            }
        };

        while let Some(node) = queue.pop_front() {
            if op(&node.value()) {
                return Some(node.clone());
            }

            for child in node.children().iter() {
                queue.push_back(child.clone())
            }
        }
        None
    }

    // Do a depth-first-search in order to get the path to a node. Start with a naive recursive
    // implementation, then switch to a stack-based implementation in order to avoid exceeding the
    // stack.
    pub fn path_to<F>(&self, f: F) -> Vec<Node<T>>
    where
        F: FnOnce(&T) -> bool + Copy,
    {
        unimplemented!()
    }

    /// Convert each node of a tree from type T to type U
    pub fn map<F, U>(&self, op: F) -> Tree<U>
    where
        F: FnOnce(&T) -> U + Copy,
    {
        // A key part of this is to avoid recursion. There is no telling how deep a tree may go (Go
        // game records can go hundreds of nodes deep), so we're going to just avoid recursion.
        match self {
            Tree::Empty => Tree::Empty,
            Tree::Root(root) => {
                let new_root = Node::new(op(&root.value()));

                // This queue serves as a work list. Each node in the queue needs to be converted,
                // and I've paired the node up with the one that it's supposed to be attached to.
                // So, as we look at a node A, we make sure that all of its children gets added to
                // the queue, and that the queue knows that the conversion of each child node
                // should get attached to A.
                let mut queue: VecDeque<(Node<T>, Node<U>)> = root
                    .children()
                    .iter()
                    .map(|child| (child.clone(), new_root.clone()))
                    .collect();

                while let Some((source, dest)) = queue.pop_front() {
                    let res = Node::new(op(&source.value()));
                    dest.add_child_node(res.clone());

                    for child in source.children().iter() {
                        queue.push_back((child.clone(), res.clone()));
                    }
                }
                Tree::Root(new_root)
            }
        }
    }
}

// By using the Rc<RefCell> container here, I'm able to make Node easily clonable while still
// having the contents be shared. This means that I can change the tree structure without having to
// make the visible objects mutable.
//
// This feels like cheating the type system.
//
// However, since I've moved the RefCell inside of the node, I can borrow the node multiple times
// in a traversal function and I can make changes to nodes that I find.
#[derive(Debug)]
pub struct Node<T>(Rc<RefCell<Node_<T>>>);

impl<T> Clone for Node<T> {
    fn clone(&self) -> Self {
        Node(Rc::clone(&self.0))
    }
}

#[derive(Debug)]
struct Node_<T> {
    value: T,
    children: Vec<Node<T>>,
}

impl<T> Node<T> {
    pub fn new(value: T) -> Self {
        Self(Rc::new(RefCell::new(Node_ {
            value,
            children: vec![],
        })))
    }

    /// Immutably retrieve the data in this node.
    pub fn value(&self) -> Ref<T> {
        // Ref::map is not actually a member function. I don't know why this was done, other than
        // maybe to avoid conflicting with other `map` declarations. Why that is necessary when
        // Option::map exists as a member, I don't know.
        Ref::map(self.0.borrow(), |v| &v.value)
    }

    pub fn children(&self) -> Ref<Vec<Node<T>>> {
        Ref::map(self.0.borrow(), |v| &v.children)
    }

    pub fn add_child_node(&self, child: Node<T>) {
        self.0.borrow_mut().children.push(child)
    }

    pub fn add_child_value(&self, value: T) -> Node<T> {
        let node = Node::new(value);
        self.0.borrow_mut().children.push(node.clone());
        node
    }
}

impl<T: PartialEq> PartialEq for Node<T> {
    fn eq(&self, other: &Node<T>) -> bool {
        self.0.borrow().value == other.0.borrow().value
            && self.0.borrow().children == other.0.borrow().children
    }
}

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

    #[test]
    fn can_find_node_in_tree() {
        let mut tree = Tree::default();
        tree.set_value(15);
        assert!(tree.find_bfs(|val| *val == 15).is_some());
        assert!(tree.find_bfs(|val| *val == 16).is_none());

        let node = tree.find_bfs(|val| *val == 15).unwrap();
        node.add_child_value(20);

        assert!(tree.find_bfs(|val| *val == 20).is_some());
    }

    #[test]
    fn node_can_add_children() {
        let n = Node::new(15);
        n.add_child_value(20);

        assert_eq!(*n.value(), 15);
        // assert_eq!(n.children(), vec![Rc::new(RefCell::new(Node::new(20)))]);
    }

    #[test]
    fn it_can_map_one_tree_to_another() {
        let (tree, n) = Tree::new(15);
        let n = n.add_child_value(16);
        let _ = n.add_child_value(17);

        let tree2 = tree.map(|v| v.to_string());

        assert!(tree2.find_bfs(|val| *val == "15").is_some());
        assert!(tree2.find_bfs(|val| *val == "16").is_some());
        assert!(tree2.find_bfs(|val| *val == "17").is_some());
    }

    #[test]
    fn path_to_on_empty_tree_returns_empty() {
        let tree: Tree<&str> = Tree::default();

        assert_eq!(tree.path_to(|val| *val == "i"), vec![]);
    }

    // A
    // B G   H
    // C     I
    // D E F
    #[test]
    fn it_can_find_a_path_to_a_node() {
        let (tree, a) = Tree::new("A");
        let b = a.add_child_value("B");
        let c = b.add_child_value("C");
        let _d = c.add_child_value("D");
        let _e = c.add_child_value("D");
        let _f = c.add_child_value("D");
        let _g = a.add_child_value("G");
        let h = a.add_child_value("H");
        let i = a.add_child_value("I");

        assert_eq!(tree.path_to(|val| *val == "z"), vec![]);
        assert_eq!(tree.path_to(|val| *val == "i"), vec![a, h, i]);
    }
}