Ifaz Kabir - Symmetry in Red-Black Trees

This post assumes familiarity with the Rc<RefCell<...>> type in Rust. The Red-Black Tree parts aren’t specific to Rust, but some of pointer dereferencing is a bit awkward because we’re using Rc<RefCell<...>>.

I’ve been taking another look at Rust. I read the Rust Book, and then have been implementing some data structures which led me to implement Red-Black Trees. Going through the Red-Black Tree implementation in CLRS, I was very unsatisfied with the repeated code with the symmetric cases. In the different cases, there’s a bunch of code first figuring out which child of a parent a given node is and then acting accordingly. This felt counterproductive to me since we could have precomputed this while creating the tree. The pre-computation uses some additional space, but Red-Black Trees already use some additional space for the Color, so it feels harmless. If there’s some extra space left in the Node struct after storing Color, we might not even need extra space.

We start with defining the Node type.

type BareTree<K, T> = Rc<RefCell<Node<K, T>>>;
type Tree<K, T> = Option<BareTree<K, T>>;

pub struct Node<K: Ord, T> {
    color: Color,
    key: K,
    value: T,
    parent: Tree<K, T>,
    child_of_parent: Option<RBSide>,
    left: Tree<K, T>,
    right: Tree<K, T>,
}

Here, the only addition compared to a standard implementation is the child_of_parent field. This field is used to store whether the node is a left child or a right child.

We’re going to need some primitive operations on Nodes: set_parent, set_child, take_child, and get_child.

impl<K, T> Node<K, T> {
    fn get_child(&self, child: RBSide) -> Tree<K, T> {
        match child {
            RBSide::Left => self.left.clone(),
            RBSide::Right => self.right.clone(),
        }
    }

    fn take_child(&mut self, child: RBSide) -> Tree<K, T> {
        match child {
            RBSide::Left => self.left.take(),
            RBSide::Right => self.right.take(),
        }
    }

    fn set_parent(&mut self, child_of_parent: RBSide, parent: BareTree<K, T>) {
        self.parent = Some(parent);
        self.child_of_parent = Some(child_of_parent);
    }

    fn set_child(&mut self, side: RBSide, child: Tree<K, T>) {
        match side {
            RBSide::Left => self.left = child,
            RBSide::Right => self.right = child,
        }
    }
}

They all take an RBSide parameter. This is what allows us to share the code between the different symmetric cases in Red-Black Tree. Instead of rewriting code for left and right cases, we implicitly act on the correct case using RBSide. The set_parent and set_child are shallow operations, setting the child does not automatically set the parent of that child. The take_child sets the child field to None when fetching the child.

The code for adding a node to a tree is similar to regular Red-Black Trees, but still benefits from code sharing between the symmetric cases. We simply get which direction to go down, and then string together take_child, a recursive call, a set_parent, and a set_child. A difference from the regular adding for Red-Black Trees, the set_parent calls set the child_of_parent field as we are recurse down the tree.

pub struct RBTree<K: Ord, T> {
    root: Tree<K, T>,
}

impl<K: Ord + Copy, T: Clone> RBTree<K, T> {
    pub fn add(&mut self, key: K, value: T) {
        let root = self.root.take();
        let (_new_rooted, new_node) = self.add_r(root, key, value);
        self.root = self.fix_tree(new_node)
    }

    fn insert_side(&self, subroot_key: K, insert_key: K) -> RBSide {
        if subroot_key <= insert_key {
            RBSide::Right
        } else {
            RBSide::Left
        }
    }

    fn add_r(&mut self, node: Tree<K, T>, key: K, value: T) -> (BareTree<K, T>, BareTree<K, T>) {
        match node {
            None => {
                let new = Node::new(key, value);
                (new.clone(), new)
            }
            Some(n) => {
                let current_key = n.borrow().key;
                let insert_side = self.insert_side(current_key, key);

                // Insert into insert_side
                let old_child_subtree = n.borrow_mut().take_child(insert_side);
                let (new_child_subtree, inserted_node) = self.add_r(old_child_subtree, key, value);

                new_child_subtree
                    .borrow_mut()
                    .set_parent(insert_side, n.clone());
                n.borrow_mut()
                    .set_child(insert_side, Some(new_child_subtree));
                (n, inserted_node)
            }
        }
    }
}

After an insert, to restore the red-black tree properties, we must fix recursively balance any red-red edges or the color the root red if it’s black. This is the job of the RBTree::fix_tree function, which relies on two helper functions, RBTree::uncle and RBTree::rotate.

RBTree::uncle not only gets the uncle node, but also which side of the grand parent the uncle lies on. Because we precomputed which side of the grandparent node we start from, getting to the uncle node is vastly simplified. In the code below, most of the lines are just there for dereferencing through &Rc<RefCell<...>>.

    // should only be called nodes where the parent is red, i.e. grand parent exists
    fn uncle(&self, node: &BareTree<K, T>) -> (RBSide, Tree<K, T>) {
        let node_ref = node.borrow(); 
        let parent = node_ref.parent.as_ref().unwrap();
        let parent = parent.borrow();
        let other_side = parent.child_of_parent.unwrap().other();
        let grand_parent = parent.parent.as_ref().unwrap();
        let uncle = grand_parent.borrow().get_child(other_side);
        (other_side, uncle)
    }

RBTree::rotate rotates the subtree at parent and assumes that there’s a non-None child on child_side. We don’t have to write separate code for left and right rotations, and just use the single function for both kinds of rotation, just passing the direction where the child node is.

    fn rotate(&self, child_side: RBSide, parent: BareTree<K, T>) {
        let mut parent_mut = RefCell::borrow_mut(&parent);
        let other_side = child_side.other();

        let child_rc = parent_mut.get_child(child_side).unwrap();
        {
            // scope for borrowing child_rc
            let mut child = RefCell::borrow_mut(&child_rc);

            let grand_parent = parent_mut.parent.take();
            let child_of_grand_parent = parent_mut.child_of_parent.take();
            if let Some((gp, c)) = grand_parent.as_ref().zip(child_of_grand_parent) {
                gp.borrow_mut().set_child(c, Some(child_rc.clone()));
            }

            let grand_child = child.take_child(other_side);
            if let Some(gc) = grand_child.as_ref() {
                gc.borrow_mut().set_parent(child_side, parent.clone())
            }
            parent_mut.set_child(child_side, grand_child);

            child.parent = grand_parent;
            child.child_of_parent = child_of_grand_parent;
            child.set_child(child_side.other(), Some(parent.clone()));
        }
        parent_mut.set_parent(other_side, child_rc);
    }

We finally come to to RBTree::fix_tree. The only violations are that the root is red or there are red-red edges. The cases for red-red edges are as follows:

The uncle node is Red. In this case we color the grandparent Red, color the uncle and parent Black, and recurse on the grandparent.
The uncle node is Black. This case boils down to two cases.
1. The parent is on one side of the grandparent, but the node is on a different side of the parent. For this case, we rotate to reduce to case b.
2. Node-parent is same as parent-grandparent. The grandparent is Black. We rotate on the grandparent to make a subtree rooted at parent, with the current node and grandparent node as its children. The color of the parent node is changed from Red to Black, the color of the grandparent is changed from Black to Red, and the color of the current node stays Red.

    fn fix_tree(&mut self, inserted: BareTree<K, T>) -> Tree<K, T> {
        let mut not_root = inserted.borrow().parent.is_some();
        let root: BareTree<K, T> = if not_root {
            let mut n: BareTree<K, T> = Rc::clone(&inserted);
            let mut parent_is_red = self.parent_color(&inserted) == Color::Red;
            // n is red
            while parent_is_red && not_root {
                // parent_is_red implies grand_parent exists and is black
                let (uncle_side, uncle) = self.uncle(&n);
                let mut parent = n.borrow().parent.as_ref().unwrap().clone();
                if uncle.is_some() && uncle.as_ref().unwrap().borrow().color == Color::Red {
                    // uncle red
                    let uncle = uncle.unwrap();
                    parent.borrow_mut().color = Color::Black;
                    uncle.borrow_mut().color = Color::Black;
                    let grand_parent = parent.borrow().parent.as_ref().unwrap().clone();
                    grand_parent.borrow_mut().color = Color::Red;
                    n = grand_parent;
                } else {
                    // uncle is black
                    let parent_side = uncle_side.other();
                    let node_side = parent.borrow().child_of_parent.unwrap();
                    if node_side != parent_side {
                        // rotate to make parent of child of node
                        self.rotate(node_side, parent.clone());
                        {
                            let temp = parent;
                            parent = n;
                            n = temp;
                        }
                    };
                    // parent (currently red) will replace black grand_parent
                    parent.borrow_mut().color = Color::Black;
                    let grand_parent = RefCell::borrow(&parent).parent.clone().unwrap();
                    grand_parent.borrow_mut().color = Color::Red;
                    self.rotate(parent_side, grand_parent);
                }
                not_root = n.borrow().parent.is_some();
                parent_is_red = not_root && { self.parent_color(&n) == Color::Red }
            }
            while not_root {
                let temp = n.borrow().parent.clone().unwrap();
                not_root = temp.borrow().parent.is_some();
                n = temp;
            }
            n
        } else {
            inserted
        };
        RefCell::borrow_mut(&root).color = Color::Black;
        Some(root)
    }

    // only called on non-root (red) nodes, so parent can be unwrapped
    fn parent_color(&self, node: &BareTree<K, T>) -> Color {
        let node_ref = node.borrow();
        let parent = node_ref.parent.as_ref().unwrap();
        let parent_ref = parent.borrow();
        parent_ref.color
    }

Notice that in the above implementation, there is no repeated code for the symmetric cases. For the sake of completion, the Color and RBSide types are as follows.

enum Color {
    Red,
    Black,
}

enum RBSide {
    Left,
    Right,
}

impl RBSide {
    fn other(self) -> Self {
        match self {
            RBSide::Left => RBSide::Right,
            RBSide::Right => RBSide::Left,
        }
    }
}

The full code for this post is available here.