fix a bunch of issues in RedBlackTree#insert! and RedBlackTree#delete! algorithms

Author: joshuay03

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## [Unreleased]
 
+- Fix a bunch of issues in `RedBlackTree#insert!` and `RedBlackTree#delete!` algorithms
 - Fix `RedBlackTree::LeafNode`s being marked red
 - Handle comparison with `RedBlackTree::LeafNode` in subclasses of `RedBlackTreeNode`
 - Add `RedBlackTree#include?`

lib/red-black-tree.rb

@@ -91,8 +91,6 @@ def insert! node, target_parent = nil, direction = nil
       raise ArgumentError, "Target parent already has #{direction} child" if (child = target_parent[direction]) && child.valid?
     end
 
-    opp_direction = opposite_direction direction if direction
-
     node.parent = nil
     node.left = LeafNode.new
     node.left.parent = node
@@ -114,17 +112,16 @@ def insert! node, target_parent = nil, direction = nil
           node.parent.parent.red!
           node = node.parent.parent
         else
-          if node.position == direction
-            node = node.parent
-            rotate_sub_tree! node, opp_direction
+          opp_direction = node.opposite_position
+          if node.parent.position == opp_direction
+            rotate_sub_tree! node.parent, opp_direction
+            node = node[opp_direction]
           end
 
+          opp_direction = node.opposite_position
+          rotate_sub_tree! node.parent.parent, opp_direction
           node.parent.black!
-
-          if node.parent.parent
-            node.parent.parent.red!
-            rotate_sub_tree! node.parent.parent, direction
-          end
+          node.parent[opp_direction].red!
         end
 
         @root.black!
@@ -146,6 +143,8 @@ def insert! node, target_parent = nil, direction = nil
   def delete! node
     raise ArgumentError, "cannot delete leaf node" if node.instance_of? LeafNode
 
+    original_node = node
+
     if node.children_are_valid?
       successor = node.left
       successor = successor.left until successor.left.leaf?
@@ -166,59 +165,54 @@ def delete! node
       if is_root? node
         @root = nil
       elsif node.red?
-        leaf = LeafNode.new
-        node.swap_position_with! leaf
+        node.swap_position_with! LeafNode.new
       else
-        direction = node.position
-        opp_direction = opposite_direction direction
-
         loop do
-          if node.sibling.valid? && node.sibling.red?
+          if node.sibling && node.sibling.valid? && node.sibling.red?
             node.parent.red!
             node.sibling.black!
-            rotate_sub_tree! node.parent, direction
-
-            next
+            rotate_sub_tree! node.parent, node.position
           end
 
           if node.distant_nephew && node.distant_nephew.valid? && node.distant_nephew.red?
-            unless node.sibling.leaf?
-              case node.parent.colour
-              when Node::RED then node.sibling.red!
-              when Node::BLACK then node.sibling.black!
-              end
+            case node.parent.colour
+            when Node::RED then node.sibling.red!
+            when Node::BLACK then node.sibling.black!
             end
             node.parent.black!
             node.distant_nephew.black!
-            rotate_sub_tree! node.parent, direction
+            rotate_sub_tree! node.parent, node.position
 
             break
           end
 
           if node.close_nephew && node.close_nephew.valid? && node.close_nephew.red?
             node.sibling.red! unless node.sibling.leaf?
             node.close_nephew.black!
-            rotate_sub_tree! node.sibling, opp_direction
+            rotate_sub_tree! node.sibling, node.opposite_position
 
             next
           end
 
-          if node.parent.red?
+          if node.parent && node.parent.red?
             node.sibling.red! unless node.sibling.leaf?
             node.parent.black!
 
             break
           end
 
-          break
+          if node.sibling && !node.sibling.leaf?
+            node.sibling.red!
+          end
+
+          break unless node = node.parent
         end
 
-        leaf = LeafNode.new
-        node.swap_position_with! leaf
+        original_node.swap_position_with! LeafNode.new
       end
     end
 
-    node.validate_free!
+    original_node.validate_free!
 
     decrement_size!
     update_left_most_node!

lib/red_black_tree/node.rb

@@ -72,9 +72,9 @@ def single_child_is_leaf? = children.any?(&:leaf?) && !children_are_leaves?
       def position
         return unless @parent
 
-        case self
-        when @parent.left then LEFT
-        when @parent.right then RIGHT
+        case self.object_id
+        when @parent.left.object_id then LEFT
+        when @parent.right.object_id then RIGHT
         else raise StructuralError, "Disowned by parent"
         end
       end
@@ -147,13 +147,13 @@ def swap_position_with! other_node
         self_position = position
         other_position = other_node.position
 
-        if other_node.parent == self
+        if other_node.parent.object_id == self.object_id
           self[other_position] = other_node[other_position]
           other_node[other_position] = self
 
           other_node.parent = @parent
           @parent = other_node
-        elsif other_node == @parent
+        elsif other_node.object_id == @parent.object_id
           other_node[self_position] = self[self_position]
           self[self_position] = other_node
 

test/test_red_black_tree.rb

@@ -829,7 +829,7 @@ def test_single_node_tree
     end
   end
 
-  class TestInheritance < Minitest::Test
+  class TestIntegration < Minitest::Test
     Work = Struct.new :min_latency, keyword_init: true
 
     class WorkNode < RedBlackTree::Node
@@ -838,17 +838,15 @@ def <=> other
       end
     end
 
-    def test_multiple_nodes
-      tree = RedBlackTree.new
-      tree << WorkNode.new(Work.new min_latency: 0.8)
-      tree << WorkNode.new(Work.new min_latency: 1.2)
-      tree << WorkNode.new(Work.new min_latency: 0.8)
-      tree << WorkNode.new(Work.new min_latency: 0.4)
-
-      expected = [0.4, 0.8, 0.8, 1.2]
-      until tree.empty?
-        work = tree.shift
-        assert_equal work.data.min_latency, expected.shift
+    def test_ordering
+      10.times do
+        tree = RedBlackTree.new
+        expected = 250.times.map do
+          rand(10).tap do |value|
+            tree << WorkNode.new(Work.new min_latency: value)
+          end
+        end.sort
+        assert_equal expected.shift, tree.shift.data.min_latency until tree.empty?
       end
     end
   end