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,
};

#[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
    }

    /// 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
    }
}

#[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());
    }
}