rust_black_trees/
rbtree.rs

1use std::cell::RefCell;
2use std::rc::Rc;
3
4use super::tree::BaseTree;
5use super::tree::Tree;
6
7use super::node::Node;
8use super::node::{paint, endpaint};
9use super::node::*;
10
11/// a nice convenient macro which allows a user to initialize a tree with
12/// a number of elements
13/// usage: redblack!{1, 2, 3, 4, 5, 6, 7, 8, 9, 0};
14#[macro_export]
15macro_rules! redblack {
16    ( $( $x:expr ),* ) => {
17        {
18            let mut temp_tree = RBTree::new();
19            $(
20                temp_tree.insert($x);
21            )*
22            temp_tree
23        }
24    };
25}
26
27#[derive(Debug)]
28pub struct ColorNode<T> {
29    pub value: T,
30    pub ptr: usize,
31    pub parent: Option<usize>,
32    pub lchild: Option<usize>,
33    pub rchild: Option<usize>,
34    pub color: Color,
35    data: Rc<RefCell<Vec<ColorNode<T>>>>,
36}
37
38pub trait ColoredNode<T>: Node<T> {
39    fn new(val: T, selfptr: usize, data: Rc<RefCell<Vec<ColorNode<T>>>>) -> Self;
40    fn is_red(&self) -> bool;
41    fn is_child_black(&self, side: Side) -> bool;
42    fn is_parent_black(&self) -> bool;
43    fn is_sibling_black(&self) -> bool;
44}
45
46impl<T> ColoredNode<T> for ColorNode<T>
47where
48    T: std::fmt::Debug,
49    T: std::cmp::PartialOrd,
50{
51    fn new(val: T, selfptr: usize, data: Rc<RefCell<Vec<ColorNode<T>>>>) -> Self {
52        Self {
53            value: val,
54            ptr: selfptr,
55            parent: None,
56            lchild: None,
57            rchild: None,
58            color: Color::Black,
59            data: data,
60        }
61    }
62
63    fn is_red(&self) -> bool {
64        match self.color {
65            Color::Red => true,
66            Color::Black => false,
67        }
68    }
69
70    // Nil nodes are black children too
71    fn is_child_black(&self, side: Side) -> bool {
72        let child = self.get_child(side);
73        if child.is_some() && self.get(child.unwrap()).is_red() {
74            false
75        } else {
76            true
77        }
78    }
79
80    // this will panic of called on root node
81    fn is_parent_black(&self) -> bool {
82        let p = self.parent.unwrap();
83        !self.get(p).is_red()
84    }
85
86    // Nil nodes are black children too
87    fn is_sibling_black(&self) -> bool {
88        let sib = self.get_sibling();
89        if sib.is_some() && self.get(sib.unwrap()).is_red() {
90            false
91        } else {
92            true
93        }
94    }
95}
96
97impl<T: std::fmt::Debug + std::cmp::PartialOrd> Node<T> for ColorNode<T> {
98    fn to_self_string(&self) -> String {
99        format!(
100            "[P:{:?} C:{:?} V:{:?}]",
101            self.parent, self.color, self.value
102        )
103    }
104    fn to_self_string_display(&self) -> (String, usize) {
105        const RED: usize = 1;
106        const BLK: usize = 0;
107        const FG: usize = 30;
108        const BG: usize = 40;
109        if self.is_red() {
110            (
111                format!(
112                    "{}{:?}{}",
113                    paint(FG + BLK, BG + RED),
114                    self.value,
115                    endpaint()
116                ),
117                format!("{:?}", self.value).len(),
118            )
119        } else {
120            (
121                format!(
122                    "{}{:?}{}",
123                    paint(FG + RED, BG + BLK),
124                    self.value,
125                    endpaint()
126                ),
127                format!("{:?}", self.value).len(),
128            )
129        }
130    }
131
132    fn is(&self, val: &T) -> bool {
133        &self.value == val
134    }
135    fn greater(&self, val: &T) -> bool {
136        &self.value > val
137    }
138    fn lesser(&self, val: &T) -> bool {
139        &self.value < val
140    }
141
142    fn get_value(&self) -> &T {
143        &self.value
144    }
145    /**
146     * In order to return a reference to a value of a vector contained within a
147     * refcell, a raw pointer is used. The unsafe code could be avoided by
148     * replacing each call to self.get(n) with &self.data.borrow()[n] and each call
149     * to self.get_mut(n) with &mut self.data.borrow()[n]
150     */
151    fn get(&self, ptr: usize) -> &ColorNode<T> {
152        unsafe { &(*self.data.as_ptr())[ptr] }
153    }
154
155    fn get_mut(&self, ptr: usize) -> &mut ColorNode<T> {
156        unsafe { &mut (*self.data.as_ptr())[ptr] }
157    }
158
159    fn get_child(&self, side: Side) -> Option<usize> {
160        match side {
161            Side::Left => self.lchild,
162            Side::Right => self.rchild,
163        }
164    }
165
166    fn set_child(&mut self, child: usize, side: Side) {
167        self.set_child_opt(Some(child), side)
168    }
169
170    fn set_child_opt(&mut self, c: Option<usize>, side: Side) {
171        match side {
172            Side::Left => self.lchild = c,
173            Side::Right => self.rchild = c,
174        };
175        if let Some(child) = c {
176            self.get_mut(child).parent = Some(self.location());
177        }
178    }
179    fn set_parent(&mut self, p: Option<usize>) {
180        self.parent = p;
181    }
182
183    fn get_parent(&self) -> Option<usize> {
184        self.parent
185    }
186
187    fn location(&self) -> usize {
188        self.ptr
189    }
190}
191
192/**
193 * Arena based memory tree structure
194*/
195#[derive(Debug)]
196pub struct RBTree<T> {
197    root: Option<usize>,
198    size: usize,
199    data: Rc<RefCell<Vec<ColorNode<T>>>>,
200    free: Vec<usize>,
201}
202
203impl<T> Tree<T> for RBTree<T>
204where
205    T: PartialOrd,
206    T: PartialEq,
207    T: std::fmt::Debug,
208{
209    fn new() -> Self {
210        Self {
211            root: None,
212            data: Rc::new(RefCell::new(Vec::new())),
213            size: 0,
214            free: Vec::new(),
215        }
216    }
217}
218const TREE_END: usize = 0xFFFFFFFF;
219impl<T> BaseTree<T> for RBTree<T>
220where
221    T: PartialOrd,
222    T: PartialEq,
223    T: std::fmt::Debug,
224{
225    type MNode = ColorNode<T>;
226    /**
227     * In order to return a reference to a value of a vector contained within a refcell, a raw
228     * pointer is used. The unsafe code could be avoided by replacing each call to self.get(n) with
229     * &self.data.borrow()[n] and each call to self.get_mut(n) with &mut self.data.borrow()[n]. This
230     * allows us to do the same thing with less keystrokes. It does make the program not
231     * thread-safe, but a this data structure is a pretty terrible choice for a multi-threaded data
232     * structure anyways, since re-balancing can require that most of the tree be locked to one
233     * thread during an insertion or deletion
234     */
235    fn get(&self, val: usize) -> &Self::MNode {
236        unsafe { &(*self.data.as_ptr())[val] }
237    }
238
239    fn get_mut(&self, val: usize) -> &mut Self::MNode {
240        unsafe { &mut (*self.data.as_ptr())[val] }
241    }
242
243    fn get_root(&self) -> Option<usize> {
244        self.root
245    }
246
247    fn set_root(&mut self, new_root: Option<usize>) {
248        self.root = new_root
249    }
250
251    fn crement_size(&mut self, amount: isize) {
252        self.size = (self.size as isize + amount) as usize;
253    }
254
255    fn attach_child(&self, p: usize, c: usize, side: Side) {
256        self.get_mut(p).set_child(c, side)
257    }
258
259    fn rebalance_ins(&mut self, n: usize) {
260        self.fix_ins_color(n);
261    }
262
263    fn rebalance_del(&mut self, n: usize, child: usize) {
264        /*
265        println!("Deleting {} with {}: ", n, child);
266        for n in self.data.borrow().iter() {
267            print!("({} -> {:?})", &n.ptr, &n.value);
268        }
269        println!();
270        println!("From tree:\n{}", self.to_pretty_string());
271        */
272        if self.get(n).ptr == TREE_END || self.get(child).ptr == TREE_END {
273            self.fix_del_color_long();
274        } else {
275            self.fix_del_color(n, child)
276        }
277    }
278
279    fn delete_replace(&mut self, n: usize) -> usize {
280        self.get_mut(n).ptr = TREE_END;
281        n
282    }
283
284    fn replace_node(&mut self, to_delete: usize, to_attach: Option<usize>) {
285        let node = self.get(to_delete);
286        self.get_mut(to_delete).ptr = TREE_END;
287        if let Some(p) = node.parent {
288            if node.is_child(Side::Left) {
289                self.get_mut(p).lchild = to_attach;
290            } else {
291                self.get_mut(p).rchild = to_attach;
292            }
293        } else {
294            self.root = to_attach;
295        }
296    }
297
298    fn get_size(&self) -> usize {
299        return self.size;
300    }
301
302    fn create_node(&mut self, val: T) -> usize {
303        let loc = self.data.borrow().len();
304        self.data
305            .borrow_mut()
306            .push(ColorNode::new(val, loc, self.data.clone()));
307        loc
308    }
309
310    fn delete_node(&mut self, index: usize) {
311        self.free.push(index);
312    }
313}
314
315impl<T> RBTree<T>
316where
317    T: PartialOrd,
318    T: PartialEq,
319    T: std::fmt::Debug,
320{
321    // child is the new node in the location, n is being deleted
322    fn fix_del_color(&mut self, n: usize, child: usize) {
323        if !self.get(n).is_red() {
324            if self.get(child).is_red() {
325                self.get_mut(child).color = Color::Black;
326            } else {
327                self.delete_case_1(child);
328            }
329        }
330    }
331
332    // sets a node to black if it exists. This is fine, cause all
333    // nodes that don't exist are by definition black anyways
334    fn set_maybe_black(&mut self, no: Option<usize>) {
335        if let Some(n) = no {
336            self.get_mut(n).color = Color::Black;
337        }
338    }
339
340    fn delete_case_1(&mut self, n: usize) {
341        if self.get(n).parent.is_some() {
342            self.delete_case_2(n);
343        }
344    }
345
346    fn delete_case_2(&mut self, n: usize) {
347        let s = self.get(n).get_sibling();
348        if self.get(n).is_sibling_black() {
349            let p = self.get(n).parent.expect("D2 P");
350            self.set_maybe_black(s);
351            self.get_mut(p).color = Color::Red;
352            self.rotate(self.get(n).side(), p);
353        }
354        self.delete_case_3(n);
355    }
356
357    fn delete_case_3(&mut self, n: usize) {
358        let s = self.get(n).get_sibling().expect("D3 S");
359        let p = self.get(n).parent.expect("D3 P");
360        if self.get(n).is_parent_black()
361            && !self.get(s).is_red()
362            && self.get(s).is_child_black(Side::Left)
363            && self.get(s).is_child_black(Side::Right)
364        {
365            self.delete_case_1(p);
366        } else {
367            self.delete_case_4(p);
368        }
369    }
370
371    fn delete_case_4(&mut self, n: usize) {
372        let node = self.get(n);
373        let s = node.get_sibling().expect("D4 S");
374        let p = node.parent.expect("D4 P");
375
376        if !node.is_parent_black()
377            && node.is_sibling_black()
378            && self.get(s).is_child_black(Side::Left)
379            && self.get(s).is_child_black(Side::Right)
380        {
381            self.get_mut(s).color = Color::Red;
382            self.get_mut(p).color = Color::Black;
383        } else {
384            self.delete_case_5(n)
385        }
386    }
387
388    fn delete_case_5(&mut self, n: usize) {
389        let s = self.get(n).get_sibling().expect("D5 S");
390        if !self.get(s).is_red() {
391            if self.get(n).is_child(Side::Left)
392                && self.get(s).is_child_black(Side::Right)
393                && !self.get(s).is_child_black(Side::Left)
394            {
395                let scl = self.get(s).get_child(Side::Left);
396                self.get_mut(s).color = Color::Red;
397                self.set_maybe_black(scl);
398                self.rotate(Side::Right, s);
399            } else if self.get(n).is_child(Side::Right)
400                && self.get(s).is_child_black(Side::Left)
401                && !self.get(s).is_child_black(Side::Right)
402            {
403                let scr = self.get(s).get_child(Side::Right);
404                self.get_mut(s).color = Color::Red;
405                self.set_maybe_black(scr);
406                self.rotate(Side::Left, s);
407            }
408        }
409        self.delete_case_6(n)
410    }
411
412    fn delete_case_6(&mut self, n: usize) {
413        let s = self.get(n).get_sibling().expect("D6 S");
414        let p = self.get(n).parent.expect("D6 P");
415        let pc = self.get(p).color;
416        self.get_mut(s).color = pc;
417        self.get_mut(p).color = Color::Black;
418
419        if self.get(n).is_child(Side::Left) {
420            let scr = self.get(s).get_child(Side::Right);
421            self.set_maybe_black(scr);
422            self.rotate(Side::Left, p);
423        } else {
424            let scl = self.get(s).get_child(Side::Left);
425            self.set_maybe_black(scl);
426            self.rotate(Side::Right, p);
427        }
428    }
429
430    fn fix_del_color_long(&mut self) {
431        let mut t = RBTree::new();
432        let mut v = self.data.borrow_mut().pop();
433        while v.is_some() {
434            let n = v.unwrap();
435            if n.ptr != TREE_END {
436                t.insert(n.value);
437            }
438            v = self.data.borrow_mut().pop();
439        }
440
441        //self.data = t.data;
442        //self.root = t.root;
443        //self.size = t.size;
444        //self.free = t.free;
445        // println!("the new post deleted tree:\n{}", t.to_pretty_string());
446        *self = t;
447        self.size += 1;
448    }
449
450    fn fix_ins_color(&mut self, n: usize) {
451        self.get_mut(n).color = Color::Red;
452        if let Some(p) = self.get(n).parent {
453            if !self.get(p).is_red() {
454                // parent is black
455                // do nothing
456            } else if self.get(n).get_uncle().is_some()
457                && self.get(self.get(n).get_uncle().unwrap()).is_red()
458            {
459                // uncle exists and is red
460                let p = self.get(n).parent.unwrap();
461                let u = self.get(n).get_uncle().unwrap();
462                self.get_mut(p).color = Color::Black;
463                self.get_mut(u).color = Color::Black;
464                self.fix_ins_color(self.get(p).parent.unwrap());
465            } else {
466                // uncle is black
467                self.do_ins_hard_case(n);
468            }
469        } else {
470        }
471        self.get_mut(self.root.unwrap()).color = Color::Black;
472    }
473
474    fn do_ins_hard_case(&mut self, nn: usize) {
475        let mut n = nn;
476        let mut p = self.get(n).parent.unwrap();
477        if self.get(p).is_child(Side::Left) && self.get(n).is_child(Side::Right) {
478            self.rotate(Side::Left, n);
479            n = self.get(n).get_child(Side::Left).unwrap();
480        }
481
482        p = self.get(n).parent.unwrap();
483        if self.get(p).is_child(Side::Right) && self.get(n).is_child(Side::Left) {
484            self.rotate(Side::Right, n);
485            n = self.get(n).get_child(Side::Right).unwrap();
486        }
487
488        self.do_ins_hard_case2(n);
489    }
490
491    fn do_ins_hard_case2(&mut self, n: usize) {
492        let p = self.get(n).parent.unwrap();
493        let g = self.get(p).parent.unwrap();
494
495        self.get_mut(p).color = Color::Black;
496        self.get_mut(g).color = Color::Red;
497        if self.get(p).is_child(Side::Right) {
498            self.rotate(Side::Left, p);
499        } else if self.get(p).is_child(Side::Left) {
500            self.rotate(Side::Right, p);
501        }
502    }
503
504    #[allow(dead_code)]
505    fn get_size_recursive(&self) -> usize {
506        if let Some(root) = self.root {
507            self.get(root).get_size()
508        } else {
509            0
510        }
511    }
512}
513
514#[cfg(test)]
515mod tests {
516    use super::*;
517
518    /** ([P:None C:Black V:50
519        ([P:Some(0) C:Black V:25
520            ([P:Some(1) C:Black V:15
521                ([P:Some(3) C:Black V:0])
522                ([P:Some(3) C:Black V:20])])
523            ([P:Some(1) C:Black V:35
524                ([P:Some(4) C:Black V:30])
525                ([P:Some(4) C:Black V:40])])])
526        ([P:Some(0) C:Black V:75
527            ([P:Some(2) C:Black V:65
528                ([P:Some(5) C:Black V:60])
529                ([P:Some(5) C:Black V:70])])
530            ([P:Some(2) C:Black V:85
531                ([P:Some(6) C:Black V:80])
532                ([P:Some(6) C:Black V:100])])])])
533    */
534    fn make_fake_tree_node_no_balance() -> RBTree<i32> {
535        let mut tree = RBTree::new();
536        for x in vec![50, 25, 75, 15, 35, 65, 85, 0, 20, 30, 40, 60, 70, 80, 100] {
537            tree.insert(x);
538        }
539
540        tree
541    }
542
543    #[test]
544    fn test_tree_print() {
545        let tree = make_fake_tree_node_no_balance();
546        assert_eq!(tree.to_string(), "([P:None C:Black V:50] ([P:Some(0) C:Red V:25] ([P:Some(1) C:Black V:15] ([P:Some(3) C:Red V:0] () ()) ([P:Some(3) C:Red V:20] () ())) ([P:Some(1) C:Black V:35] ([P:Some(4) C:Red V:30] () ()) ([P:Some(4) C:Red V:40] () ()))) ([P:Some(0) C:Red V:75] ([P:Some(2) C:Black V:65] ([P:Some(5) C:Red V:60] () ()) ([P:Some(5) C:Red V:70] () ())) ([P:Some(2) C:Black V:85] ([P:Some(6) C:Red V:80] () ()) ([P:Some(6) C:Red V:100] () ()))))");
547        let tree_empty = RBTree::<i32>::new();
548        assert_eq!(tree_empty.to_string(), "(Empty tree)");
549    }
550
551    #[test]
552    fn test_contains() {
553        let tree = make_fake_tree_node_no_balance();
554        assert!(tree.contains(&100));
555        assert!(tree.contains(&0));
556        assert!(tree.contains(&50));
557        assert!(tree.contains(&25));
558        assert!(tree.contains(&75));
559        assert!(tree.contains(&60));
560        assert!(tree.contains(&40));
561
562        assert!(!tree.contains(&42));
563        assert!(!tree.contains(&99));
564        assert!(!tree.contains(&1));
565    }
566
567    #[test]
568    fn test_insert() {
569        let mut tree = RBTree::<i32>::new();
570        for x in 0..10 {
571            double_size_test(&tree, x as usize);
572            tree.insert(x);
573            double_size_test(&tree, (x + 1) as usize);
574        }
575
576        assert_eq!(tree.to_string(), "([P:None C:Black V:3] ([P:Some(3) C:Black V:1] ([P:Some(1) C:Black V:0] () ()) ([P:Some(1) C:Black V:2] () ())) ([P:Some(3) C:Black V:5] ([P:Some(5) C:Black V:4] () ()) ([P:Some(5) C:Red V:7] ([P:Some(7) C:Black V:6] () ()) ([P:Some(7) C:Black V:8] () ([P:Some(8) C:Red V:9] () ())))))");
577    }
578
579    #[test]
580    fn test_insert_2() {
581        let mut tree = RBTree::<usize>::new();
582        let size = 15;
583        for x in (0..size).rev() {
584            tree.insert(x);
585        }
586
587        assert_eq!(tree.to_string(), "([P:None C:Black V:11] ([P:Some(3) C:Red V:7] ([P:Some(7) C:Black V:5] ([P:Some(9) C:Red V:3] ([P:Some(11) C:Black V:1] ([P:Some(13) C:Red V:0] () ()) ([P:Some(13) C:Red V:2] () ())) ([P:Some(11) C:Black V:4] () ())) ([P:Some(9) C:Black V:6] () ())) ([P:Some(7) C:Black V:9] ([P:Some(5) C:Black V:8] () ()) ([P:Some(5) C:Black V:10] () ()))) ([P:Some(3) C:Black V:13] ([P:Some(1) C:Black V:12] () ()) ([P:Some(1) C:Black V:14] () ())))");
588    }
589
590    fn double_size_test<T: PartialEq + PartialOrd + std::fmt::Debug>(
591        tree: &RBTree<T>,
592        expect: usize,
593    ) {
594        assert_eq!(tree.get_size(), expect);
595        assert_eq!(tree.get_size_recursive(), expect);
596    }
597
598    #[test]
599    fn test_height() {
600        let tree = make_fake_tree_node_no_balance();
601        assert_eq!(tree.get_height(), 4);
602
603        let mut tree2 = RBTree::<i32>::new();
604        for x in 0..10 {
605            tree2.insert(x);
606        }
607        assert_eq!(tree2.get_height(), 5);
608
609        let tree3 = RBTree::<i32>::new();
610        assert_eq!(tree3.get_height(), 0);
611    }
612
613    #[test]
614    fn test_delete_mem() {
615        let mut tree = make_fake_tree_node_no_balance();
616        double_size_test(&tree, 15);
617
618        tree.delete(1);
619        double_size_test(&tree, 15);
620
621        tree.delete(0);
622        double_size_test(&tree, 14);
623        tree.insert(0);
624        double_size_test(&tree, 15);
625
626        tree.delete(50);
627        double_size_test(&tree, 14);
628        tree.insert(50);
629        double_size_test(&tree, 15);
630    }
631
632    #[test]
633    fn test_delete() {
634        let mut tree = RBTree::new();
635        for x in 0..15 {
636            tree.insert(x);
637        }
638        double_size_test(&tree, 15);
639        tree.delete(14);
640        tree.delete(12);
641        tree.delete(13);
642        assert_eq!(tree.to_string(), "([P:None C:Black V:8] ([P:Some(3) C:Black V:4] ([P:Some(7) C:Red V:2] ([P:Some(9) C:Black V:1] ([P:Some(10) C:Red V:0] () ()) ()) ([P:Some(9) C:Black V:3] () ())) ([P:Some(7) C:Red V:6] ([P:Some(5) C:Black V:5] () ()) ([P:Some(5) C:Black V:7] () ()))) ([P:Some(3) C:Black V:10] ([P:Some(1) C:Black V:9] () ()) ([P:Some(1) C:Black V:11] () ())))");
643    }
644
645    #[test]
646    fn find_leaf_count() {
647        let mut tree = RBTree::new();
648        for x in 0..15 {
649            tree.insert(x);
650        }
651        assert_eq!(tree.get_leaf_count(), 8);
652
653        let mut tree = RBTree::new();
654        for x in 0..100 {
655            tree.insert(x);
656        }
657        assert_eq!(tree.get_leaf_count(), 50);
658    }
659
660    #[test]
661    fn is_empty() {
662        let mut tree = RBTree::new();
663        assert!(tree.is_empty());
664        tree.insert(1);
665        assert!(!tree.is_empty());
666        tree.delete(1);
667        assert!(tree.is_empty());
668    }
669}
rust_black_trees/rbtree.rs

rust_black_trees/
rbtree.rs
