From 714be51160ab7849affcac7c555327cf53acb926 Mon Sep 17 00:00:00 2001 From: Britt Cyr Date: Thu, 10 Oct 2024 16:24:11 -0400 Subject: [PATCH] Remove tail recursion in red black tree (#174) * remove tail recursion in red black tree * Unwrap more * do not recurse --- lib/src/red_black_tree.rs | 61 +++++++++++++++++++++++++-------------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/lib/src/red_black_tree.rs b/lib/src/red_black_tree.rs index 7c8686e93..4bad57d76 100644 --- a/lib/src/red_black_tree.rs +++ b/lib/src/red_black_tree.rs @@ -1022,8 +1022,17 @@ impl<'a, V: Payload> HyperTreeWriteOperations<'a, V> for RedBlackTree<'a, V> { if self.max_index != NIL && *get_helper::>(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::() @@ -1080,7 +1089,15 @@ impl<'a, V: Payload> HyperTreeWriteOperations<'a, V> for RedBlackTree<'a, V> { let child_index: DataIndex = self.get_child_index::(index); let parent_index: DataIndex = self.get_parent_index::(index); self.update_parent_child::(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; + } + } } } @@ -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::(current_index, parent_index); @@ -1134,7 +1155,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { self.set_color::(parent_index, sibling_color); self.set_color::(sibling_index, parent_color); self.rotate_right::(parent_index); - return; + return (NIL, NIL); } // ii left right if self.get_color::(sibling_right_child_index) == Color::Red @@ -1145,7 +1166,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { self.set_color::(sibling_index, Color::Black); self.rotate_left::(sibling_index); self.rotate_right::(parent_index); - return; + return (NIL, NIL); } // iii right right if self.get_color::(sibling_right_child_index) == Color::Red @@ -1155,7 +1176,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { self.set_color::(parent_index, sibling_color); self.set_color::(sibling_index, parent_color); self.rotate_left::(parent_index); - return; + return (NIL, NIL); } // iv right left if self.get_color::(sibling_left_child_index) == Color::Red @@ -1166,7 +1187,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { self.set_color::(sibling_index, Color::Black); self.rotate_right::(sibling_index); self.rotate_left::(parent_index); - return; + return (NIL, NIL); } unreachable!(); } @@ -1176,12 +1197,10 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { if sibling_color == Color::Black { self.set_color::(sibling_index, Color::Red); if parent_color == Color::Black { - // Recurse on the parent - self.remove_fix(parent_index, self.get_parent_index::(parent_index)); - return; + return (parent_index, self.get_parent_index::(parent_index)); } else { self.set_color::(parent_index, Color::Black); - return; + return (NIL, NIL); } } @@ -1191,13 +1210,14 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { self.rotate_right::(parent_index); self.set_color::(parent_index, Color::Red); self.set_color::(sibling_index, Color::Black); - self.remove_fix(current_index, parent_index); + return (current_index, parent_index); } else if self.is_right_child::(sibling_index) { self.rotate_left::(parent_index); self.set_color::(parent_index, Color::Red); self.set_color::(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 @@ -1258,10 +1278,10 @@ 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::(index_to_fix, Color::Black); - return; + return NIL; } // Check the color of the parent. If it is black, then nothing left to do. @@ -1269,7 +1289,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { let parent_color: Color = self.get_color::(parent_index); if parent_color == Color::Black { - return; + return NIL; } let grandparent_index: DataIndex = self.get_parent_index::(parent_index); @@ -1296,9 +1316,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { self.set_color::(uncle_index, Color::Black); self.set_color::(grandparent_index, Color::Red); - // Recurse - self.insert_fix(grandparent_index); - return; + return grandparent_index; } let grandparent_color: Color = self.get_color::(grandparent_index); @@ -1310,7 +1328,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { if grandparent_index == NIL && parent_color == Color::Red { self.set_color::(parent_index, Color::Black); - return; + return NIL; } // Case II: Uncle is black, left left @@ -1340,6 +1358,7 @@ impl<'a, V: Payload> RedBlackTree<'a, V> { self.set_color::(index_to_fix, grandparent_color); self.set_color::(grandparent_index, index_to_fix_color); } + NIL } }