ego_binary_tree/
binary_tree.rs

1use ego_tree::{NodeMut, NodeRef, Tree};
2
3/// Wrapper around a ego_tree::Tree that constrains functionality / API
4/// to a binary tree. Always contains at least one node.
5pub struct BinaryTree<T> {
6    inner: Tree<Option<T>>,
7}
8
9impl<T> BinaryTree<T> {
10    pub fn new(root_value: T) -> Self {
11        let mut tree = Tree::new(Some(root_value));
12        let mut root = tree.root_mut();
13        root.append(None);
14        root.append(None);
15        Self { inner: tree }
16    }
17
18    /// Returns a reference to the root node.
19    pub fn root(&self) -> BinaryNodeRef<T> {
20        BinaryNodeRef::wrap(self.inner.root())
21    }
22
23    /// Returns a mutator of the root node.
24    pub fn root_mut(&mut self) -> BinaryNodeMut<T> {
25        BinaryNodeMut::wrap(self.inner.root_mut())
26    }
27}
28
29#[derive(Debug, PartialEq, Clone, Copy)]
30pub struct BinaryNodeRef<'a, T> {
31    inner: NodeRef<'a, Option<T>>,
32}
33
34impl<'a, T> BinaryNodeRef<'a, T> {
35    /// Return the left child, if exists.
36    pub fn left(&self) -> Option<BinaryNodeRef<'a, T>> {
37        let left = self.inner.children().next().expect("always has children");
38        if left.value().is_none() {
39            return None;
40        }
41
42        Some(BinaryNodeRef::wrap(left))
43    }
44
45    /// Return the right child, if exists.
46    pub fn right(&self) -> Option<BinaryNodeRef<'a, T>> {
47        let mut children = self.inner.children();
48        let _left = children.next().expect("always has children");
49
50        let right = children.next().expect("always has children");
51        if right.value().is_none() {
52            return None;
53        }
54
55        Some(BinaryNodeRef::wrap(right))
56    }
57
58    /// Get the value for this node.
59    pub fn value(&self) -> &T {
60        self.inner.value().as_ref().expect("exists")
61    }
62
63    fn wrap(node: NodeRef<'a, Option<T>>) -> Self {
64        Self { inner: node }
65    }
66}
67
68#[derive(Debug)]
69pub struct BinaryNodeMut<'a, T> {
70    inner: NodeMut<'a, Option<T>>,
71}
72
73impl<'a, T> BinaryNodeMut<'a, T> {
74    fn left_inner(&mut self) -> NodeMut<Option<T>> {
75        self.inner.first_child().expect("exists")
76    }
77
78    fn right_inner(&mut self) -> NodeMut<Option<T>> {
79        self.inner.last_child().expect("exists")
80    }
81
82    /// Return the left child, if exists.
83    pub fn left(&mut self) -> Option<BinaryNodeMut<T>> {
84        let mut left_inner = self.left_inner();
85        if left_inner.value().is_none() {
86            return None;
87        }
88
89        Some(BinaryNodeMut::wrap(left_inner))
90    }
91
92    /// Return the right child, if exists.
93    pub fn right(&mut self) -> Option<BinaryNodeMut<T>> {
94        let mut right_inner = self.right_inner();
95        if right_inner.value().is_none() {
96            return None;
97        }
98
99        Some(BinaryNodeMut::wrap(right_inner))
100    }
101
102    /// Get the value for this node.
103    pub fn value(&mut self) -> &mut T {
104        self.inner.value().as_mut().expect("exists")
105    }
106
107    /// Set the right child to value and return the node.
108    pub fn set_right(&mut self, value: T) -> BinaryNodeMut<T> {
109        let mut right_inner = self.right_inner();
110        *right_inner.value() = Some(value);
111        // Create left / right nodes for new node if not present.
112        if !right_inner.has_children() {
113            right_inner.append(None);
114            right_inner.append(None);
115        }
116        BinaryNodeMut::wrap(right_inner)
117    }
118
119    /// Set the left child to value and return the node.
120    pub fn set_left(&mut self, value: T) -> BinaryNodeMut<T> {
121        let mut left_inner = self.left_inner();
122        *left_inner.value() = Some(value);
123        // Create left / right nodes for new node if not present.
124        if !left_inner.has_children() {
125            left_inner.append(None);
126            left_inner.append(None);
127        }
128        BinaryNodeMut::wrap(left_inner)
129    }
130
131    fn wrap(node: NodeMut<'a, Option<T>>) -> Self {
132        Self { inner: node }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn binary_tree_api_works() {
142        let mut tree = BinaryTree::new(5);
143        assert!(tree.root_mut().left().is_none());
144        assert!(tree.root_mut().right().is_none());
145
146        let mut root = tree.root_mut();
147        let mut left = root.set_left(3);
148        assert_eq!(left.value(), &3);
149        assert!(left.left().is_none());
150        assert!(left.right().is_none());
151    }
152}
ego_binary_tree/binary_tree.rs

ego_binary_tree/
binary_tree.rs
