writing a tree in rust

1
2
3
4
5
6
7

#[derive(PartialEq)]
struct TreeNode<T> {
    pub value: Option<T>,
    pub children: Vec<Rc<RefCell<TreeNode<T>>>>,
    pub parent: Option<Rc<RefCell<TreeNode<T>>>>,
}

1
2
3
4
5
6
7

 pub fn new() -> TreeNode<T> {
        TreeNode {
            value: None,
            children: vec![],
            parent: None,
        }
    }

1
2
3

    pub fn add_child(&mut self, new_node: Rc<RefCell<TreeNode<T>>>) {
        self.children.push(new_node);
    }

1
2
3
4

    pub fn del_child(&mut self, node_to_remove: Rc<RefCell<TreeNode<T>>>) {
        self.children
            .retain(|child| !Rc::ptr_eq(child, &node_to_remove));
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

pub fn print(&self) -> String
    where
        T: ToString,
    {
        if let Some(ref value) = self.value {
            value.to_string()
        } else {
            String::from("[")
                + &self
                    .children
                    .iter()
                    .map(|tn| tn.borrow().print())
                    .collect::<Vec<String>>()
                    .join(",")
                + "]"
        }
    }

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use std::cell::RefCell;
use std::rc::Rc;

#[derive(PartialEq)]
struct TreeNode<T> {
    pub value: Option<T>,
    pub children: Vec<Rc<RefCell<TreeNode<T>>>>,
    pub parent: Option<Rc<RefCell<TreeNode<T>>>>,
}

impl<T> TreeNode<T> {
    pub fn new() -> TreeNode<T> {
        TreeNode {
            value: None,
            children: vec![],
            parent: None,
        }
    }

    pub fn add_child(&mut self, new_node: Rc<RefCell<TreeNode<T>>>) {
        self.children.push(new_node);
    }

    pub fn del_child(&mut self, node_to_remove: Rc<RefCell<TreeNode<T>>>) {
        self.children
            .retain(|child| !Rc::ptr_eq(child, &node_to_remove));
    }
    pub fn print(&self) -> String
    where
        T: ToString,
    {
        if let Some(ref value) = self.value {
            value.to_string()
        } else {
            String::from("[")
                + &self
                    .children
                    .iter()
                    .map(|tn| tn.borrow().print())
                    .collect::<Vec<String>>()
                    .join(",")
                + "]"
        }
    }
}

fn init_tree(s: String) -> Rc<RefCell<TreeNode<u32>>> {
    let root = Rc::new(RefCell::new(TreeNode::new()));
    let mut current = Rc::clone(&root);
    let chars = s.chars().collect::<Vec<char>>();
    for (_, c) in chars
        .iter()
        .enumerate()
        .filter(|(idx, _)| *idx > 0 && *idx + 1 < chars.len())
    {
        if *c == '[' || c.is_numeric() {
            let child = Rc::new(RefCell::new(TreeNode::new()));
            current.borrow_mut().children.push(Rc::clone(&child));
            {
                let mut mut_child = child.borrow_mut();
                mut_child.parent = Some(Rc::clone(&current));
                if c.is_numeric() {
                    mut_child.value = c.to_digit(10);
                }
            }
            current = child;
        } else if *c == ',' || *c == ']' {
            let current_clone = Rc::clone(&current);
            current = Rc::clone(current_clone.borrow().parent.as_ref().unwrap());
        } else {
            panic!("Unknown character: {}", c);
        }
    }
    root
}

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

    #[test]
    fn test_init_tree_1() {
        let tree = init_tree(String::from("[1,2]"));
        assert_eq!(tree.borrow().children[0].borrow().value.unwrap(), 1);
    }

    #[test]
    fn test_init_tree_2() {
        let tree = init_tree(String::from("[1,2]"));
        assert_eq!(tree.borrow().children[1].borrow().value.unwrap(), 2);
    }

    #[test]
    fn test_init_tree_3() {
        let tree = init_tree(String::from("[0,1,[3,4,5,[7,8]],2]"));
        assert_eq!(tree.borrow().print(), "[0,1,[3,4,5,[7,8]],2]");
    }

    #[test]
    fn test_add_child() {
        let tree = init_tree(String::from("[0,1,[3,4,5,[7,8]],2]"));
        let new_node = Rc::new(RefCell::new(TreeNode::new()));
        new_node.borrow_mut().value = Some(9);
        let child = &tree.borrow().children[2];
        child.borrow_mut().add_child(new_node);
        assert_eq!(tree.borrow().print(), "[0,1,[3,4,5,[7,8],9],2]");
    }

    #[test]
    fn test_del_child() {
        let tree = init_tree(String::from("[0,1,[3,4,5,[7,8]],2]"));
        let child_to_remove = Rc::clone(&tree.borrow().children[2].borrow().children[2]);
        tree.borrow_mut().children[2]
            .borrow_mut()
            .del_child(child_to_remove);
        assert_eq!(tree.borrow().print(), "[0,1,[3,4,[7,8]],2]");
    }
}