3 files changed, 42 insertions, 9 deletions

diff --git a/src/03/rb_tree.c b/src/03/rb_tree.c
index 60b5c1a..cb576b4 100644
--- a/src/03/rb_tree.c
+++ b/src/03/rb_tree.c
@@ -7,11 +7,11 @@
  * * https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Insertion
  */
 
-RBTree *parent_of(RBTree *node) {
+static RBTree *parent_of(RBTree *node) {
   return node ? node->parent : NULL;
 }
 
-RBTree *root_of(RBTree *node) {
+static RBTree *root_of(RBTree *node) {
   RBTree *current = node;
   RBTree *next = parent_of(current);
 
@@ -22,11 +22,11 @@ RBTree *root_of(RBTree *node) {
   return current;
 }
 
-RBTree *grand_parent_of(RBTree *node) {
+static RBTree *grand_parent_of(RBTree *node) {
   return parent_of(parent_of(node));
 }
 
-RBTree *sibling_of(RBTree *node) {
+static RBTree *sibling_of(RBTree *node) {
   RBTree *parent = parent_of(node);
 
   if (!parent)
@@ -35,7 +35,7 @@ RBTree *sibling_of(RBTree *node) {
   return node == parent->left ? parent->right : parent->left;
 }
 
-RBTree *pibling_of(RBTree *node) {
+static RBTree *pibling_of(RBTree *node) {
   return sibling_of(parent_of(node));
 }
 
@@ -83,8 +83,9 @@ static void repair_from(RBTree *tree) {
   RBTree *parent = parent_of(tree);
   RBTree *pibling = pibling_of(tree);
 
-  if (parent == NULL || parent->colour == black)
+  if (parent == NULL || parent->colour == black) {
     return;
+  }
 
   if (pibling && pibling->colour == red) {
     parent->colour = black;
@@ -94,6 +95,8 @@ static void repair_from(RBTree *tree) {
   } else {
     RBTree *grand_parent = grand_parent_of(tree);
 
+    if (!grand_parent)
+      return;
     if (tree == parent->right && parent == grand_parent->left) {
       rb_rotate_left(parent);
     } else if (tree == parent->left && parent == grand_parent->right) {
@@ -104,10 +107,12 @@ static void repair_from(RBTree *tree) {
     parent = parent_of(tree);
     grand_parent = grand_parent_of(tree);
 
-    if (tree == parent->left)
+    if (tree == parent->left) {
       rb_rotate_right(grand_parent);
-    else
+    }
+    else {
       rb_rotate_left(grand_parent);
+    }
     parent->colour = black;
     grand_parent->colour = red;
   }
@@ -183,6 +188,17 @@ void rb_tree_inspect(RBTree *tree) {
   print_tree(tree, 0);
 }
 
+int rb_tree_size(RBTree *tree) {
+  int total = 0;
+  if (tree == NULL)
+    return total;
+  if (tree->left)
+    total += rb_tree_size(tree->left);
+  if (tree->right)
+    total += rb_tree_size(tree->right);
+  return total + 1;
+}
+
 bool rb_equals(RBTree *tree, RBTree *other_tree) {
   if (!tree || !other_tree)
     return tree == other_tree;
diff --git a/src/03/rb_tree.h b/src/03/rb_tree.h
index 9d4f0e4..c43bd61 100644
--- a/src/03/rb_tree.h
+++ b/src/03/rb_tree.h
@@ -16,5 +16,6 @@ typedef struct rb_node {
 RBTree *rb_tree_initialize(int value);
 RBTree *rb_tree_initialize_with(int value, enum Colour colour);
 RBTree *rb_tree_insert(RBTree *tree, int value);
-void rb_tree_inspect(RBTree *tree);
 bool rb_equals(RBTree *tree, RBTree *other_tree);
+int rb_tree_size(RBTree *tree);
+void rb_tree_inspect(RBTree *tree);
diff --git a/src/03/rb_tree_test.c b/src/03/rb_tree_test.c
index 5ffff22..0193250 100644
--- a/src/03/rb_tree_test.c
+++ b/src/03/rb_tree_test.c
@@ -139,6 +139,21 @@ Ensure(rb_tree_insert_repaints_the_new_node) {
   assert_that(tree->right->colour, is_equal_to(black));
 }
 
+Ensure(rb_tree_insert_handles_large_trees) {
+  RBTree *tree = NULL;
+  int n = 100;
+
+  for (int i = n; i > 0; i--)
+    tree = rb_tree_insert(tree, i);
+  rb_tree_inspect(tree);
+
+  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(rb_tree_size(tree), is_equal_to(n));
+}
+
 Ensure(equals_returns_false_when_tree_is_NULL) {
   assert_that(rb_equals(NULL, rb_tree_initialize(10)), is_equal_to(false));
 }
@@ -216,6 +231,7 @@ TestSuite *rb_tree_tests() {
   add_test(x, rb_tree_insert_performs_a_right_rotation);
   add_test(x, rb_tree_insert_performs_a_left_rotation);
   add_test(x, rb_tree_insert_repaints_the_new_node);
+  add_test(x, rb_tree_insert_handles_large_trees);
 
   add_test(x, equals_returns_false_when_tree_is_NULL);
   add_test(x, equals_returns_false_when_other_tree_is_NULL);