3
\$\begingroup\$

Although I have done a few toy projects in Rust, I haven't really done anything related to tedious memory management tasks. So I set out to create a database with the help of the project 500 lines or less. I am using Red black trees for now. I would love to hear your opinions and critic on this partial Red-black tree which only supports insertion.
My main concern is how to improve the code to make it more idiomatic. I have plans to write a tutorial after I am done with the project and I don't want others to set off in the wrong direction.

Code:

use std::cell::RefCell;
use std::rc::{Rc, Weak};

struct TreeNode<T: PartialOrd> {
    datum: Option<T>,
    is_red: RefCell<bool>,
    left: RefCell<Option<Rc<TreeNode<T>>>>,
    right: RefCell<Option<Rc<TreeNode<T>>>>,
    parent: RefCell<Option<Weak<TreeNode<T>>>>,
}

impl<T> TreeNode<T> where T: PartialOrd {
    fn new() -> Self {
        TreeNode {
            datum: None,
            is_red: RefCell::new(false),
            left: RefCell::new(None),
            right: RefCell::new(None),
            parent: RefCell::new(None)
        }
    }

    fn from(t: T, is_red: bool) -> Self {
        let left = RefCell::new(None);
        let right = RefCell::new(None);
        let parent = RefCell::new(None);
        TreeNode {
            datum: Some(t),
            is_red: RefCell::new(is_red),
            left,
            right,
            parent
        }
    }

    fn parent_is_red(&self) -> bool {
        if let Some(ref p) = *self.parent.borrow() {
            if let Some(ref p_upgrade) = p.upgrade() {
                *p_upgrade.is_red.borrow()
            } else {
                false
            }
        } else {
            false
        }
    }

    fn height(&self) -> usize {
        let left_height = match *self.left.borrow() {
            Some(ref left) => left.height(),
            None => 0,
        };
        let right_height = match *self.right.borrow() {
            Some(ref right) => right.height(),
            None => 0,
        };
        if left_height > right_height { left_height + 1 } else { right_height + 1 }
    }

    fn inorder_scan<F>(&self, f: Rc<F>) where F:Fn(&T) {
        if let Some(ref left) = *self.left.borrow() {
            left.inorder_scan(Rc::clone(&f));
        }
        Rc::clone(&f)(self.datum.as_ref().unwrap());
        if let Some(ref right) = *self.right.borrow() {
            right.inorder_scan(Rc::clone(&f));
        }
    }
}

pub struct RBTree<T: PartialOrd> {
    root: Rc<TreeNode<T>>,
    count: usize
}

impl<T> RBTree<T> where T:PartialOrd {
    pub fn new() -> Self {
        RBTree{root: Rc::new(TreeNode::new()), count: 0 }
    }

    fn rotate_right(&mut self, node: Rc<TreeNode<T>>) {
        // make sure node.left is not Nil
        if let None = *node.left.borrow() {
            return;
        }
        // take out node.left, leaving node.left with None
        let left = Rc::clone(node.left.borrow_mut().take().as_ref().unwrap());
        // put left.right at node.left
        *node.left.borrow_mut() = left.right.borrow_mut().take();
        // if the above taken node is not empty, set it's parent to node
        if let Some(ref node_left) = *node.left.borrow() {
            *node_left.parent.borrow_mut() = Some(Rc::downgrade(&Rc::clone(&node)));
        }

        // if node's parent is not empty, and node is parent's right child
        // set left as parent's right, else set it as parent's left, and set it's parent to node's parent
        // if node's parent is empty, set left as root, and it's parent as None
        if let Some(ref parent) = *node.parent.borrow() {
            let parent_ref = Rc::clone(parent.upgrade().as_ref().unwrap());
            let parent_right = parent_ref.right.borrow();

            match *parent_right {
                Some(ref inner_parent_right) if &*Rc::clone(parent_right.as_ref().unwrap()) as *const _ == &*node as *const _ => {
                    *inner_parent_right.right.borrow_mut() = Some(Rc::clone(&left));
                },
                _ => {
                    *parent_ref.left.borrow_mut() = Some(Rc::clone(&left));
                }
            };

            *left.parent.borrow_mut() = Some(Weak::clone(parent));
        } else {
            self.root = Rc::clone(&left);
            *left.parent.borrow_mut() = None;
        }
        *left.right.borrow_mut() = Some(Rc::clone(&node));
        *node.parent.borrow_mut() = Some(Rc::downgrade(&left));
    }

    // opposite version of rotate_right
    fn rotate_left(&mut self, node: Rc<TreeNode<T>>) {
        if let None = *node.right.borrow() {
            return;
        }
        let right = Rc::clone(node.right.borrow_mut().take().as_ref().unwrap());
        *node.right.borrow_mut() = right.left.borrow_mut().take();
        if let Some(ref node_right) = *node.right.borrow() {
            *node_right.parent.borrow_mut() = Some(Rc::downgrade(&Rc::clone(&node)));
        }

        if let Some(ref parent) = *node.parent.borrow() {
            let parent_ref = Rc::clone(parent.upgrade().as_ref().unwrap());
            let parent_left = parent_ref.left.borrow();

            match *parent_left {
                Some(ref _p_left) if &*Rc::clone(parent_left.as_ref().unwrap()) as *const _ == &*node as *const _ => {
                    *_p_left.right.borrow_mut() = Some(Rc::clone(&right));
                },
                _ => {
                    *parent_ref.right.borrow_mut() = Some(Rc::clone(&right));
                }
            };

            if parent_left.is_some() &&
                &*Rc::clone(parent_left.as_ref().unwrap()) as *const TreeNode<T> == &*node as *const TreeNode<T> {
                drop(parent_left);
                *parent_ref.left.borrow_mut() = Some(Rc::clone(&right));
            } else {
                drop(parent_left);
                *parent_ref.right.borrow_mut() = Some(Rc::clone(&right));
            }
            *right.parent.borrow_mut() = Some(Weak::clone(parent));
        } else {
            self.root = Rc::clone(&right);
            *right.parent.borrow_mut() = None;
        }
        *right.left.borrow_mut() = Some(Rc::clone(&node));
        *node.parent.borrow_mut() = Some(Rc::downgrade(&right));
    }

    // insert t at the leaf and fix the tree if there is any violation
    pub fn insert(&mut self, t: T) {
        if self.count != 0 {
            let mut parent = None;
            let mut current = Some(Rc::clone(&self.root));
            // while current is not None
            while current.is_some() {
                parent = current;
                let parent_unwrapped = parent.as_ref().unwrap();
                if parent_unwrapped.datum.as_ref().unwrap() > &t {
                    match *parent_unwrapped.left.borrow() {
                        Some(ref left) => current = Some(Rc::clone(&left)),
                        None => current = None,
                    };
                } else {
                    match *parent_unwrapped.right.borrow() {
                        Some(ref right) => current = Some(Rc::clone(&right)),
                        None => current = None,
                    };
                }
            }

            // parent can be safely unwrapped because it will not be None
            // at this stage, current is None because we're at the leaf now
            // if parent's datum is less than t, then t must be inserted at the left
            // otherwise, t must be inserted at the left
            let node: Rc<TreeNode<T>>;
            let parent_unwrapped = parent.as_ref().unwrap();
            if parent_unwrapped.datum.as_ref().unwrap() > &t {
                node = Rc::new(TreeNode::from(t, true));
                *node.parent.borrow_mut() = Some(Rc::downgrade(parent_unwrapped));
                *parent_unwrapped.left.borrow_mut() = Some(Rc::clone(&node));
            } else {
                node = Rc::new(TreeNode::from(t, true));
                *node.parent.borrow_mut() = Some(Rc::downgrade(parent_unwrapped));
                *parent_unwrapped.right.borrow_mut() = Some(Rc::clone(&node));
            }
            &self.fix_tree(node);
        } else {
            // root must be black
            self.root = Rc::new(TreeNode::from(t, false));
        }
        self.count += 1;
    }

    // restore violation of red-black properties
    fn fix_tree(&mut self, node: Rc<TreeNode<T>>) {
        let mut current = Rc::clone(&node);

        while current.parent_is_red() {
            // grand parent exists because parent is red, and rb tree's rule 1 states root is black
            // so there's at least one node at the root above current's parent
            // the unwrap here is safe for both parent and grandparent
            let parent = Rc::clone(&current.parent.borrow().as_ref().unwrap().upgrade().unwrap());
            let grandparent = Rc::clone(&parent.parent.borrow().as_ref().unwrap().upgrade().unwrap());

            // parent is the left child of grandparent, i.e current's uncle is grandparent's right child
            if grandparent.left.borrow().is_some() &&
                &**grandparent.left.borrow().as_ref().unwrap() as *const TreeNode<T> == &*parent as *const TreeNode<T> {
                if let Some(uncle) = grandparent.right.borrow().as_ref() {
                    // uncle exists and is red
                    // paints uncle and parent black, grandparent red
                    // points current to grandparent
                    if *uncle.is_red.borrow() {
                        *uncle.is_red.borrow_mut() = false;
                        *parent.is_red.borrow_mut() = false;
                        *grandparent.is_red.borrow_mut() = true;
                        current = Rc::clone(&grandparent);
                        continue;
                    }
                }
                // uncle is black; either it points to None, or the color is black
                // if current is parent's right child, left rotate it first
                if parent.right.borrow().is_some() &&
                    &**parent.right.borrow().as_ref().unwrap() as *const TreeNode<T> == &*current as *const TreeNode<T> {
                    current = Rc::clone(&parent);
                    self.rotate_left(Rc::clone(&current));
                }

                *parent.is_red.borrow_mut() = false;
                *grandparent.is_red.borrow_mut() = true;
                &self.rotate_right(Rc::clone(&grandparent));
            } else {
                // do the opposite of the above case because parent is the right child of the grandfather
                if let Some(uncle) = grandparent.left.borrow().as_ref() {
                    if *uncle.is_red.borrow() {
                        *uncle.is_red.borrow_mut() = false;
                        *parent.is_red.borrow_mut() = false;
                        *grandparent.is_red.borrow_mut() = true;
                        current = Rc::clone(&grandparent);
                        continue;
                    }
                }

                if parent.left.borrow().is_some() &&
                    &**parent.left.borrow().as_ref().unwrap() as *const TreeNode<T> == &*current as *const TreeNode<T> {
                    current = Rc::clone(&parent);
                    self.rotate_right(Rc::clone(&current));
                }

                *parent.is_red.borrow_mut() = false;
                *grandparent.is_red.borrow_mut() = true;
                &self.rotate_left(Rc::clone(&grandparent));
            }
        }
        *self.root.is_red.borrow_mut() = false;
    }

    pub fn height(&self) -> usize {
        if self.count == 0 {
            0
        } else {
            self.root.height()
        }
    }

    pub fn inorder_scan<F>(&self, f: F) where F: Fn(&T) {
        if self.count > 0 {
            let rc = Rc::new(f);
            Rc::clone(&self.root).inorder_scan(rc);
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn test_insertion_height_and_count() {
        for i in 0..20 {
            let vec: Vec<u32> = (0..i).collect();
            let mut tree = RBTree::new();
            for item in vec {
                tree.insert(item)
            }
            let log = 2.0 * (i as f32 + 1.0).log(2.0);
            assert!(tree.height() as f32 <= log);
            assert_eq!(i as usize, tree.count);
            tree.inorder_scan(|x| print!("{} ", x));
            println!();
        }
    }
}
\$\endgroup\$
  • \$\begingroup\$ I don't have time to do a full review right now, but Clippy gives some helpful suggestions when run on your code. Chief among them is that you don't need to be taking f or node by value into your methods - a reference would suffice. Also, stuff like Rc::clone(&x) can be replaced with x.clone(). \$\endgroup\$ – Joe Clay Apr 11 '18 at 12:10

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.