test: ensure red-black fails when any criteria is invalid

2 files changed, 62 insertions, 1 deletions

diff --git a/src/03/rb_tree.c b/src/03/rb_tree.c
index 9604d4d..8e082d6 100644
--- a/src/03/rb_tree.c
+++ b/src/03/rb_tree.c
@@ -7,6 +7,31 @@
  * * https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Insertion
  */
 
+static int max(int a, int b) {
+  return a == b ? a : (a > b ? a : b);
+}
+
+/**
+ * Number of black nodes to leaf.
+ *
+ * @param tree The node to traverse down to a leaf.
+ * @return the # of black nodes from the given node to a leaf.
+ */
+static int depth(RBTree *tree) {
+  int total = 1;
+
+  while (tree) {
+    if (tree->colour == black)
+      total += 1;
+    tree = tree->left;
+  }
+  return total;
+}
+
+static bool is_root(RBTree *node) {
+  return node->parent == NULL;
+}
+
 static RBTree *parent_of(RBTree *node) {
   return node ? node->parent : NULL;
 }
@@ -213,5 +238,17 @@ bool rb_equals(RBTree *tree, RBTree *other_tree) {
 }
 
 bool rb_tree_is_valid(RBTree *tree) {
-  return false;
+  if (tree == NULL)
+    return true;
+
+  if (is_root(tree) && tree->colour == red)
+    return false;
+
+  if (tree->colour == red && tree->parent->colour == red)
+    return false;
+
+  if (depth(tree->left) != depth(tree->right))
+      return false;
+
+  return rb_tree_is_valid(tree->left) && rb_tree_is_valid(tree->right);
 }
diff --git a/src/03/rb_tree_test.c b/src/03/rb_tree_test.c
index 663158d..8545326 100644
--- a/src/03/rb_tree_test.c
+++ b/src/03/rb_tree_test.c
@@ -226,6 +226,28 @@ Ensure(is_valid_returns_false_when_root_is_red) {
   assert_that(rb_tree_is_valid(tree), is_equal_to(false));
 }
 
+Ensure(is_valid_returns_false_when_red_node_has_red_child) {
+  RBTree *tree = NULL;
+
+  for (int i = 10; i > 0; i--)
+    tree = rb_tree_insert(tree, i);
+
+  tree->left->colour = red;
+  tree->left->left->colour = red;
+
+  assert_that(rb_tree_is_valid(tree), is_equal_to(false));
+}
+
+Ensure(is_valid_returns_false_when_each_path_to_leaves_does_not_contain_the_same_number_of_black_nodes) {
+  RBTree *tree = NULL;
+
+  for (int i = 10; i > 0; i--)
+    tree = rb_tree_insert(tree, i);
+
+  tree->left->left->colour = black;
+  assert_that(rb_tree_is_valid(tree), is_equal_to(false));
+}
+
 TestSuite *rb_tree_tests() {
   TestSuite *x = create_test_suite();
 
@@ -250,5 +272,7 @@ TestSuite *rb_tree_tests() {
   add_test(x, equals_returns_false_when_root_and_right_subtree_are_not_equal);
 
   add_test(x, is_valid_returns_false_when_root_is_red);
+  add_test(x, is_valid_returns_false_when_red_node_has_red_child);
+  add_test(x, is_valid_returns_false_when_each_path_to_leaves_does_not_contain_the_same_number_of_black_nodes);
   return x;
 }