fix: root of red/black tree is black

2 files changed, 8 insertions, 5 deletions

diff --git a/src/03/rb_tree.c b/src/03/rb_tree.c
index cb576b4..1b74d90 100644
--- a/src/03/rb_tree.c
+++ b/src/03/rb_tree.c
@@ -90,7 +90,9 @@ static void repair_from(RBTree *tree) {
   if (pibling && pibling->colour == red) {
     parent->colour = black;
     pibling->colour = black;
-    grand_parent_of(tree)->colour = red;
+    RBTree *grand_parent = grand_parent_of(tree);
+    if (grand_parent->parent)
+      grand_parent->colour = red;
     repair_from(grand_parent_of(tree));
   } else {
     RBTree *grand_parent = grand_parent_of(tree);
@@ -114,7 +116,8 @@ static void repair_from(RBTree *tree) {
       rb_rotate_left(grand_parent);
     }
     parent->colour = black;
-    grand_parent->colour = red;
+    if (grand_parent->parent)
+      grand_parent->colour = red;
   }
 }
 
diff --git a/src/03/rb_tree_test.c b/src/03/rb_tree_test.c
index 9b78dc0..4f12f66 100644
--- a/src/03/rb_tree_test.c
+++ b/src/03/rb_tree_test.c
@@ -113,7 +113,7 @@ Ensure(rb_tree_insert_performs_a_left_rotation) {
 
 Ensure(rb_tree_insert_repaints_the_new_node) {
 /*
-     (20:b)                   (20:r)
+     (20:b)                   (20:b)
      /    \                  /     \
   (10:r)  (30:r)  -->   (10:b)    (30:b)
   /                     /
@@ -127,7 +127,7 @@ Ensure(rb_tree_insert_repaints_the_new_node) {
 
   assert_that(tree, is_not_equal_to(NULL));
   assert_that(tree->value, is_equal_to(20));
-  assert_that(tree->colour, is_equal_to(red));
+  assert_that(tree->colour, is_equal_to(black));
 
   assert_that(tree->left->value, is_equal_to(10));
   assert_that(tree->left->colour, is_equal_to(black));
@@ -148,7 +148,7 @@ Ensure(rb_tree_insert_handles_large_trees) {
 
   assert_that(tree, is_not_equal_to(NULL));
   assert_that(tree->value, is_equal_to(69));
-  assert_that(tree->colour, is_equal_to(red));
+  assert_that(tree->colour, is_equal_to(black));
 
   assert_that(rb_tree_size(tree), is_equal_to(n));
 }