Skip to content

Commit

Permalink
Remove tail recursion in red black tree (#174)
Browse files Browse the repository at this point in the history 
* remove tail recursion in red black tree

* Unwrap more

* do not recurse
  • Loading branch information
@brittcyr
brittcyr authored Oct 10, 2024
1 parent 143cd73 commit 714be51
Showing 1 changed file with 40 additions and 21 deletions.
61 changes: 40 additions & 21 deletions lib/src/red_black_tree.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1022,8 +1022,17 @@ impl<'a, V: Payload> HyperTreeWriteOperations<'a, V> for RedBlackTree<'a, V> {
if self.max_index != NIL && *get_helper::<RBNode<V>>(self.data, self.max_index) < new_node {
self.max_index = index;
}

self.insert_node_no_fix(new_node, index);
self.insert_fix(index);

// Avoid recursion by doing a loop here.
let mut node_to_fix: DataIndex = index;
loop {
node_to_fix = self.insert_fix(node_to_fix);
if node_to_fix == NIL {
break;
}
}

#[cfg(test)]
self.verify_rb_tree::<V>()
Expand Down Expand Up @@ -1080,7 +1089,15 @@ impl<'a, V: Payload> HyperTreeWriteOperations<'a, V> for RedBlackTree<'a, V> {
let child_index: DataIndex = self.get_child_index::<V>(index);
let parent_index: DataIndex = self.get_parent_index::<V>(index);
self.update_parent_child::<V>(index);
self.remove_fix(child_index, parent_index);

// Avoid recursion by doing a loop here.
let mut nodes_to_fix: (DataIndex, DataIndex) = (child_index, parent_index);
loop {
nodes_to_fix = self.remove_fix(nodes_to_fix.0, nodes_to_fix.1);
if nodes_to_fix.0 == NIL && nodes_to_fix.1 == NIL {
break;
}
}
}
}

Expand All @@ -1105,13 +1122,17 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.remove_by_index(index);
}

fn remove_fix(&mut self, current_index: DataIndex, parent_index: DataIndex) {
fn remove_fix(
&mut self,
current_index: DataIndex,
parent_index: DataIndex,
) -> (DataIndex, DataIndex) {
// Current is double black. It could be NIL if we just deleted a leaf,
// so we need the parent to know where in the tree we are.

// If we get to the root, then we are done.
if self.root_index == current_index {
return;
return (NIL, NIL);
}

let sibling_index: DataIndex = self.get_sibling_index::<V>(current_index, parent_index);
Expand All @@ -1134,7 +1155,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.set_color::<V>(parent_index, sibling_color);
self.set_color::<V>(sibling_index, parent_color);
self.rotate_right::<V>(parent_index);
return;
return (NIL, NIL);
}
// ii left right
if self.get_color::<V>(sibling_right_child_index) == Color::Red
Expand All @@ -1145,7 +1166,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.set_color::<V>(sibling_index, Color::Black);
self.rotate_left::<V>(sibling_index);
self.rotate_right::<V>(parent_index);
return;
return (NIL, NIL);
}
// iii right right
if self.get_color::<V>(sibling_right_child_index) == Color::Red
Expand All @@ -1155,7 +1176,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.set_color::<V>(parent_index, sibling_color);
self.set_color::<V>(sibling_index, parent_color);
self.rotate_left::<V>(parent_index);
return;
return (NIL, NIL);
}
// iv right left
if self.get_color::<V>(sibling_left_child_index) == Color::Red
Expand All @@ -1166,7 +1187,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.set_color::<V>(sibling_index, Color::Black);
self.rotate_right::<V>(sibling_index);
self.rotate_left::<V>(parent_index);
return;
return (NIL, NIL);
}
unreachable!();
}
Expand All @@ -1176,12 +1197,10 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
if sibling_color == Color::Black {
self.set_color::<V>(sibling_index, Color::Red);
if parent_color == Color::Black {
// Recurse on the parent
self.remove_fix(parent_index, self.get_parent_index::<V>(parent_index));
return;
return (parent_index, self.get_parent_index::<V>(parent_index));
} else {
self.set_color::<V>(parent_index, Color::Black);
return;
return (NIL, NIL);
}
}

Expand All @@ -1191,13 +1210,14 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.rotate_right::<V>(parent_index);
self.set_color::<V>(parent_index, Color::Red);
self.set_color::<V>(sibling_index, Color::Black);
self.remove_fix(current_index, parent_index);
return (current_index, parent_index);
} else if self.is_right_child::<V>(sibling_index) {
self.rotate_left::<V>(parent_index);
self.set_color::<V>(parent_index, Color::Red);
self.set_color::<V>(sibling_index, Color::Black);
self.remove_fix(current_index, parent_index);
return (current_index, parent_index);
}
return (NIL, NIL);
}

/// Insert a node into the subtree without fixing. This node could be a leaf
Expand Down Expand Up @@ -1258,18 +1278,18 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
}
}

fn insert_fix(&mut self, index_to_fix: DataIndex) {
fn insert_fix(&mut self, index_to_fix: DataIndex) -> DataIndex {
if self.root_index == index_to_fix {
self.set_color::<V>(index_to_fix, Color::Black);
return;
return NIL;
}

// Check the color of the parent. If it is black, then nothing left to do.
let parent_index: DataIndex = self.get_parent_index::<V>(index_to_fix);
let parent_color: Color = self.get_color::<V>(parent_index);

if parent_color == Color::Black {
return;
return NIL;
}

let grandparent_index: DataIndex = self.get_parent_index::<V>(parent_index);
Expand All @@ -1296,9 +1316,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.set_color::<V>(uncle_index, Color::Black);
self.set_color::<V>(grandparent_index, Color::Red);

// Recurse
self.insert_fix(grandparent_index);
return;
return grandparent_index;
}

let grandparent_color: Color = self.get_color::<V>(grandparent_index);
Expand All @@ -1310,7 +1328,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {

if grandparent_index == NIL && parent_color == Color::Red {
self.set_color::<V>(parent_index, Color::Black);
return;
return NIL;
}

// Case II: Uncle is black, left left
Expand Down Expand Up @@ -1340,6 +1358,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> {
self.set_color::<V>(index_to_fix, grandparent_color);
self.set_color::<V>(grandparent_index, index_to_fix_color);
}
NIL
}
}

Expand Down

0 comments on commit 714be51

Please sign in to comment.