atlas_rb_tree/
tree.rs

1use crate::node::{Node, NodeColor};
2use std::{
3    cell::RefCell,
4    fmt::{Debug, Formatter},
5    rc::Rc,
6};
7
8#[cfg(test)]
9mod tree_tests;
10
11pub struct Tree<T> {
12    root: Rc<RefCell<Node<T>>>,
13    sentinel: Rc<RefCell<Node<T>>>,
14    length: usize,
15}
16
17// TODO cleanup traits. For instance, debug might be to strict.
18impl<T: PartialOrd + Clone + PartialEq + Debug + Default> Tree<T> {
19    pub fn new() -> Tree<T> {
20        let sentinel = Rc::new(RefCell::new(Node::new_sentinel()));
21        Self {
22            root: sentinel.clone(),
23            sentinel,
24            length: 0,
25        }
26    }
27
28    pub fn insert(&mut self, key: T) {
29        let mut z = Node::new(key);
30        let mut x = self.root.clone();
31        let mut y = self.sentinel.clone();
32
33        while !x.borrow().is_nil() {
34            y = x.clone();
35            if z.key < x.borrow().key {
36                let x_tmp = x.borrow().left().clone();
37                x = x_tmp
38            } else {
39                let x_tmp = x.borrow().right().clone();
40                x = x_tmp;
41            }
42        }
43        z.set_parent(y.clone());
44        // Z is now Reference counted for
45        let z = Rc::new(RefCell::new(z));
46
47        if y.borrow().is_nil() {
48            self.root = z.clone();
49        } else if z.borrow().key < y.borrow().key {
50            y.borrow_mut().set_left_child(z.clone());
51        } else {
52            y.borrow_mut().set_right_child(z.clone());
53        }
54
55        z.borrow_mut().set_left_child(self.sentinel.clone());
56        z.borrow_mut().set_right_child(self.sentinel.clone());
57        z.borrow_mut().color = NodeColor::Red;
58        self.insert_fix_up(z);
59        self.length += 1;
60    }
61
62    fn insert_fix_up(&mut self, mut z: Rc<RefCell<Node<T>>>) {
63        while z.borrow().parent().borrow().color == NodeColor::Red {
64            if z.borrow().parent() == z.borrow().parent().borrow().parent().borrow().left() {
65                let y = z
66                    .borrow()
67                    .parent()
68                    .borrow()
69                    .parent()
70                    .borrow()
71                    .right()
72                    .clone();
73                // Case 1
74                if y.borrow().color == NodeColor::Red {
75                    z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
76                    y.borrow_mut().color = NodeColor::Black;
77                    z.borrow_mut()
78                        .parent_mut()
79                        .borrow_mut()
80                        .parent_mut()
81                        .borrow_mut()
82                        .color = NodeColor::Red;
83                    let z_tmp = z.borrow().parent().borrow().parent().clone();
84                    z = z_tmp;
85                } else {
86                    // Case 2
87                    if &z == z.borrow().parent().borrow().right() {
88                        let z_tmp = z.borrow().parent().clone();
89                        z = z_tmp;
90                        self.left_rotate(z.clone());
91                    }
92                    // Case 3
93                    z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
94                    z.borrow_mut()
95                        .parent_mut()
96                        .borrow_mut()
97                        .parent_mut()
98                        .borrow_mut()
99                        .color = NodeColor::Red;
100                    let x = z.borrow().parent().borrow().parent().clone();
101                    self.right_rotate(x);
102                }
103            } else {
104                let y = z
105                    .borrow()
106                    .parent()
107                    .borrow()
108                    .parent()
109                    .borrow()
110                    .left()
111                    .clone();
112                // Case 4
113                if y.borrow().color == NodeColor::Red {
114                    z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
115                    y.borrow_mut().color = NodeColor::Black;
116                    z.borrow_mut()
117                        .parent_mut()
118                        .borrow_mut()
119                        .parent_mut()
120                        .borrow_mut()
121                        .color = NodeColor::Red;
122                    let z_tmp = z.borrow().parent().borrow().parent().clone();
123                    z = z_tmp;
124                } else {
125                    // Case 5
126                    if &z == z.borrow().parent().borrow().left() {
127                        let z_tmp = z.borrow().parent().clone();
128                        z = z_tmp;
129                        self.right_rotate(z.clone());
130                    }
131                    // Case 6
132                    z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
133                    z.borrow_mut()
134                        .parent_mut()
135                        .borrow_mut()
136                        .parent_mut()
137                        .borrow_mut()
138                        .color = NodeColor::Red;
139                    let x = z.borrow().parent().borrow().parent().clone();
140                    self.left_rotate(x);
141                }
142            }
143        }
144        self.root.borrow_mut().color = NodeColor::Black;
145    }
146
147    fn left_rotate(&mut self, x: Rc<RefCell<Node<T>>>) {
148        // Assumes that x.right != T.nil
149        if !x.borrow().right().borrow().is_nil() {
150            // let y = x.borrow_mut().right.take().unwrap();
151            let y = x.borrow().right().clone();
152            x.borrow_mut().set_right_child(y.borrow().left().clone());
153            if !y.borrow().left().borrow().is_nil() {
154                y.borrow_mut().left_mut().borrow_mut().set_parent(x.clone());
155            }
156            y.borrow_mut().set_parent(x.borrow().parent().clone());
157            if x.borrow().parent().borrow().is_nil() {
158                self.root = y.clone();
159            } else if &x == x.borrow().parent().borrow().left() {
160                x.borrow_mut()
161                    .parent_mut()
162                    .borrow_mut()
163                    .set_left_child(y.clone());
164            } else {
165                x.borrow_mut()
166                    .parent_mut()
167                    .borrow_mut()
168                    .set_right_child(y.clone());
169            }
170
171            y.borrow_mut().set_left_child(x.clone());
172            x.borrow_mut().set_parent(y);
173        } else {
174            panic!(
175                "Invariant violated. The right child of {:?} must not be T.nil.",
176                x
177            );
178        }
179    }
180
181    fn right_rotate(&mut self, y: Rc<RefCell<Node<T>>>) {
182        // Assumes that x.left != T.nil
183        if !y.borrow().left().borrow().is_nil() {
184            // let x = y.clone().borrow_mut().left.take().unwrap();
185            let x = y.borrow().left().clone();
186            y.borrow_mut().set_left_child(x.borrow().right().clone());
187            if !x.borrow().right().borrow().is_nil() {
188                x.borrow_mut()
189                    .right_mut()
190                    .borrow_mut()
191                    .set_parent(y.clone());
192            }
193            x.borrow_mut().set_parent(y.borrow().parent().clone());
194            if y.borrow().parent().borrow().is_nil() {
195                self.root = x.clone();
196            } else if &y.clone() == y.borrow().parent().borrow().right() {
197                y.borrow_mut()
198                    .parent_mut()
199                    .borrow_mut()
200                    .set_right_child(x.clone());
201            } else {
202                y.borrow_mut()
203                    .parent_mut()
204                    .borrow_mut()
205                    .set_left_child(x.clone());
206            }
207            x.borrow_mut().set_right_child(y.clone());
208            y.borrow_mut().set_parent(x);
209        } else {
210            panic!(
211                "Invariant violated. The left child of {:?} must not be T.nil.",
212                y
213            );
214        }
215    }
216
217    pub fn delete(&mut self, key: T) {
218        let node = self.search(key);
219        if node.is_some() {
220            self.delete_node(node.as_ref().unwrap().clone())
221        }
222        self.length -= 1;
223    }
224
225    fn delete_node(&mut self, z: Rc<RefCell<Node<T>>>) {
226        let mut y = z.clone();
227        let mut y_color = y.borrow().color.clone();
228        let x;
229        if z.borrow().left().borrow().is_nil() {
230            x = z.borrow().right().clone();
231            let u = z.clone();
232            let v = z.borrow().right().clone();
233            self.transplant(u, v);
234        } else if z.borrow().right().borrow().is_nil() {
235            x = z.borrow().left().clone();
236            let u = z.clone();
237            let v = z.borrow().left().clone();
238            self.transplant(u, v);
239        } else {
240            y = self
241                .minimum_node(z.borrow().right().clone())
242                .expect("Expected this to be set");
243            y_color = y.borrow().color.clone();
244            x = y.borrow().right().clone();
245            if &y != z.borrow().right() {
246                let u = y.clone();
247                let v = y.borrow().right().clone();
248                self.transplant(u, v);
249                y.borrow_mut().set_right_child(z.borrow().right().clone());
250                y.borrow_mut()
251                    .right_mut()
252                    .borrow_mut()
253                    .set_parent(y.clone());
254            } else {
255                x.borrow_mut().set_parent(y.clone());
256            }
257            let u = z.clone();
258            let v = y.clone();
259            self.transplant(u, v);
260            y.borrow_mut().set_left_child(z.borrow().left().clone());
261            y.borrow_mut().left_mut().borrow_mut().set_parent(y.clone());
262            y.borrow_mut().color = z.borrow().color.clone();
263        }
264
265        if y_color == NodeColor::Black {
266            self.delete_fix_up(x);
267        }
268    }
269
270    fn transplant(&mut self, u: Rc<RefCell<Node<T>>>, v: Rc<RefCell<Node<T>>>) {
271        if u.borrow().parent().borrow().is_nil() {
272            self.root = v.clone();
273        } else if &u == u.borrow().parent().borrow().left() {
274            u.borrow_mut()
275                .parent_mut()
276                .borrow_mut()
277                .set_left_child(v.clone());
278        } else {
279            u.borrow_mut()
280                .parent_mut()
281                .borrow_mut()
282                .set_right_child(v.clone());
283        }
284        v.borrow_mut().set_parent(u.borrow().parent().clone());
285    }
286
287    fn delete_fix_up(&mut self, mut x: Rc<RefCell<Node<T>>>) {
288        while x != self.root && x.borrow().color == NodeColor::Black {
289            if &x == x.borrow().parent().borrow().left() {
290                let mut w = x.borrow().parent().borrow().right().clone();
291                // Case 1
292                if w.borrow().color == NodeColor::Red {
293                    w.borrow_mut().color = NodeColor::Black;
294                    x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Red;
295                    self.left_rotate(x.borrow().parent().clone());
296                    w = x.borrow().parent().borrow().right().clone();
297                }
298
299                if w.borrow().left().borrow().color == NodeColor::Black
300                    && w.borrow().right().borrow().color == NodeColor::Black
301                {
302                    w.borrow_mut().color = NodeColor::Red;
303                    let x_tmp = x.borrow().parent().clone();
304                    x = x_tmp;
305                } else {
306                    // Case 3
307                    if w.borrow_mut().right().borrow().color == NodeColor::Black {
308                        w.borrow_mut().left().borrow_mut().color = NodeColor::Black;
309                        w.borrow_mut().color = NodeColor::Red;
310                        self.right_rotate(w.clone());
311                        w = x.borrow().parent().borrow().right().clone();
312                    }
313                    // Case 4
314                    w.borrow_mut().color = x.borrow().parent().borrow().color.clone();
315                    x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
316                    w.borrow_mut().right_mut().borrow_mut().color = NodeColor::Black;
317                    self.left_rotate(x.borrow().parent().clone());
318                    x = self.root.clone();
319                }
320            } else {
321                let mut w = x.borrow().parent().borrow().left().clone();
322                // Case 5
323                if w.borrow_mut().color == NodeColor::Red {
324                    w.borrow_mut().color = NodeColor::Black;
325                    x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Red;
326                    self.right_rotate(x.borrow().parent().clone());
327                    w = x.borrow().parent().borrow().left().clone();
328                }
329                // Case 6
330                if w.borrow().right().borrow().color == NodeColor::Black
331                    && w.borrow().left().borrow().color == NodeColor::Black
332                {
333                    w.borrow_mut().color = NodeColor::Red;
334                    let x_tmp = x.borrow().parent().clone();
335                    x = x_tmp;
336                } else {
337                    // Case 7
338                    if w.borrow().left().borrow().color == NodeColor::Black {
339                        w.borrow_mut().right_mut().borrow_mut().color = NodeColor::Black;
340                        w.borrow_mut().color = NodeColor::Red;
341                        self.left_rotate(w.clone());
342                        w = x.borrow_mut().parent().borrow().left().clone();
343                    }
344                    // Case 8
345                    w.borrow_mut().color = x.borrow().parent().borrow().color.clone();
346                    x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
347                    w.borrow_mut().left_mut().borrow_mut().color = NodeColor::Black;
348                    self.right_rotate(x.borrow().parent().clone());
349                    x = self.root.clone();
350                }
351            }
352        }
353        x.borrow_mut().color = NodeColor::Black;
354    }
355
356    pub fn contains_key(&self, key: T) -> bool {
357        self.search(key).is_some()
358    }
359
360    fn search(&self, key: T) -> Option<Rc<RefCell<Node<T>>>> {
361        let mut node = self.root.clone();
362        while !node.borrow().is_nil() {
363            let node_key = node.borrow().key.clone();
364            if node_key == key {
365                return Some(node);
366            } else if key < node_key {
367                let node_tmp = node.borrow().left().clone();
368                node = node_tmp;
369            } else {
370                let node_tmp = node.borrow().right().clone();
371                node = node_tmp;
372            }
373        }
374        None
375    }
376
377    pub fn is_empty(&self) -> bool {
378        self.length == 0
379    }
380
381    pub fn len(&self) -> usize {
382        self.length
383    }
384
385    pub fn clear(&mut self) {
386        self.root = self.sentinel.clone();
387        self.length = 0;
388    }
389
390    pub fn minimum(&self) -> Option<T> {
391        if self.root.borrow().is_nil() {
392            return None;
393        }
394
395        self.minimum_node(self.root.clone())
396            .and_then(|node| Some(node.borrow().key.clone()))
397    }
398    fn minimum_node(&self, node: Rc<RefCell<Node<T>>>) -> Option<Rc<RefCell<Node<T>>>> {
399        if node.borrow().left().borrow().is_nil() {
400            return Some(node);
401        }
402
403        let mut x = node.clone();
404        while !x.borrow().left().borrow().is_nil() {
405            let x_tmp = x.borrow().left().clone();
406            x = x_tmp;
407        }
408        Some(x)
409    }
410
411    pub fn maximum(&self) -> Option<T> {
412        if self.root.borrow().is_nil() {
413            return None;
414        }
415
416        self.maximum_node(self.root.clone())
417            .and_then(|node| Some(node.borrow().key.clone()))
418    }
419    fn maximum_node(&self, node: Rc<RefCell<Node<T>>>) -> Option<Rc<RefCell<Node<T>>>> {
420        if node.borrow().right().borrow().is_nil() {
421            return Some(node);
422        }
423
424        let mut x = node.clone();
425        while !x.borrow().right().borrow().is_nil() {
426            let x_tmp = x.borrow().right().clone();
427            x = x_tmp;
428        }
429        Some(x)
430    }
431}
432
433impl<T: Debug> Debug for Tree<T> {
434    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
435        write!(f, "{:?}", self.root)
436    }
437}
438
439/// A DFS implementation using recursion that iterates the
440/// entire tree for equality. There are a few speedups I've included,
441/// like eliminating base cases and greedily failing.
442fn tree_equality_dfs<T: PartialEq>(
443    me: Option<Rc<RefCell<Node<T>>>>,
444    other: Option<Rc<RefCell<Node<T>>>>,
445) -> bool {
446    // Solve base cases
447    match (me.clone(), other.clone()) {
448        (None, None) => panic!("Invariant violated. Node's cannot be None"),
449        (Some(_), Some(_)) => (),
450        _ => return false,
451    }
452
453    if me.as_ref().unwrap().borrow().color != other.as_ref().unwrap().borrow().color {
454        return false;
455    }
456
457    if me.as_ref().unwrap().borrow().key != other.as_ref().unwrap().borrow().key {
458        return false;
459    }
460
461    match (
462        me.as_ref().unwrap().borrow().is_sentinel,
463        other.as_ref().unwrap().borrow().is_sentinel,
464    ) {
465        (true, true) => return true,
466        (false, false) => (),
467        _ => return false,
468    }
469
470    let left_subtree = tree_equality_dfs(
471        me.as_ref().unwrap().borrow().left.clone(),
472        other.as_ref().unwrap().borrow().left.clone(),
473    );
474    if !left_subtree {
475        return false;
476    }
477
478    let right_subtree = tree_equality_dfs(
479        me.as_ref().unwrap().borrow().right.clone(),
480        other.as_ref().unwrap().borrow().right.clone(),
481    );
482    if !right_subtree {
483        return false;
484    }
485
486    true
487}
488
489impl<T: PartialEq> PartialEq<Self> for Tree<T> {
490    fn eq(&self, other: &Self) -> bool {
491        tree_equality_dfs(Some(self.root.clone()), Some(other.root.clone()))
492    }
493}
atlas_rb_tree/tree.rs

atlas_rb_tree/
tree.rs
