diff --git a/src/tree/vieter_tree.c b/src/tree/vieter_tree.c
index 792ad30..bd0200e 100644
--- a/src/tree/vieter_tree.c
+++ b/src/tree/vieter_tree.c
@@ -16,7 +16,7 @@ vieter_tree_error vieter_tree_insert(vieter_tree *tree, uint64_t key,
     tree->root = vieter_tree_node_init();
     tree->root->key = key;
     tree->root->data = data;
-    vieter_tree_node_set(tree->root, vieter_node_black);
+    vieter_tree_node_set(tree->root, vieter_tree_node_black);
     tree->size = 1;
 
     return vieter_tree_ok;
@@ -99,9 +99,9 @@ end:
   free(tree);
 }
 
-int vieter_tree_validate(vieter_tree *tree) {
+bool vieter_tree_validate(vieter_tree *tree) {
   if (tree->size == 0) {
-    return 0;
+    return true;
   }
 
   // DFS to get expected black nodes
@@ -109,12 +109,13 @@ int vieter_tree_validate(vieter_tree *tree) {
   vieter_tree_node *node = tree->root;
 
   while (node != NULL) {
-    if (vieter_tree_node_get(node, vieter_node_black)) {
+    if (vieter_tree_node_get(node, vieter_tree_node_black)) {
       expected_black_nodes++;
     }
 
     node = node->left;
   }
 
-  return vieter_tree_node_validate(tree->root, 0, expected_black_nodes);
+  return vieter_tree_node_get(tree->root, vieter_tree_node_black) &&
+         vieter_tree_node_validate(tree->root, 0, expected_black_nodes);
 }
diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index 70c52e0..d9fc66e 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -1,41 +1,31 @@
 #include "vieter_tree_balancing.h"
 
-int vieter_tree_node_validate(vieter_tree_node *node,
-                              uint64_t passed_black_nodes,
-                              uint64_t expected_black_nodes) {
-  if (vieter_tree_node_get(node, vieter_node_black)) {
+bool vieter_tree_node_validate(vieter_tree_node *node,
+                               uint64_t passed_black_nodes,
+                               uint64_t expected_black_nodes) {
+  // Path to NULL children should contain the same amount of black nodes
+  if (node == NULL) {
+    return passed_black_nodes == expected_black_nodes;
+  }
+
+  if (vieter_tree_node_get(node, vieter_tree_node_black)) {
     passed_black_nodes++;
-  } else {
-    // A red node should have black child nodes
-    if ((node->left != NULL &&
-         !vieter_tree_node_get(node->left, vieter_node_black)) ||
-        (node->right != NULL &&
-         !vieter_tree_node_get(node->right, vieter_node_black))) {
-      return 1;
-    }
   }
 
-  // All paths to a NULL child should have the same amount of black nodes
-  if ((node->left == NULL || node->right == NULL) &&
-      passed_black_nodes != expected_black_nodes) {
-    return 1;
-  }
+  // Either the node itself is black, or its children are both either NULL or
+  // black.
+  bool correctly_colored_children =
+      vieter_tree_node_get(node, vieter_tree_node_black) ||
+      ((node->left == NULL ||
+        vieter_tree_node_get(node->left, vieter_tree_node_black)) &&
+       (node->right == NULL ||
+        vieter_tree_node_get(node->right, vieter_tree_node_black)));
 
-  if (node->left != NULL) {
-    if (vieter_tree_node_validate(node->left, passed_black_nodes,
-                                  expected_black_nodes) != 0) {
-      return 1;
-    }
-  }
-
-  if (node->right != NULL) {
-    if (vieter_tree_node_validate(node->right, passed_black_nodes,
-                                  expected_black_nodes) != 0) {
-      return 1;
-    }
-  }
-
-  return 0;
+  return correctly_colored_children &&
+         vieter_tree_node_validate(node->left, passed_black_nodes,
+                                   expected_black_nodes) &&
+         vieter_tree_node_validate(node->right, passed_black_nodes,
+                                   expected_black_nodes);
 }
 
 vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
@@ -106,9 +96,9 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
     children[3] = node->right;
   }
 
-  vieter_tree_node_set(left, vieter_node_black);
-  vieter_tree_node_set(right, vieter_node_black);
-  vieter_tree_node_unset(root, vieter_node_black);
+  vieter_tree_node_set(left, vieter_tree_node_black);
+  vieter_tree_node_set(right, vieter_tree_node_black);
+  vieter_tree_node_unset(root, vieter_tree_node_black);
 
   left->left = children[0];
   left->right = children[1];
@@ -137,14 +127,16 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
 }
 
 void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
-  while (!vieter_tree_node_get(node, vieter_node_black) &&
+  while (!vieter_tree_node_get(node, vieter_tree_node_black) &&
          node->parent != NULL &&
-         !vieter_tree_node_get(node->parent, vieter_node_black)) {
+         !vieter_tree_node_get(node->parent, vieter_tree_node_black)) {
     node = vieter_tree_node_balance(node);
   }
 
   // The root must always be black
   if (node->parent == NULL) {
-    vieter_tree_node_set(node, vieter_node_black);
+    vieter_tree_node_set(node, vieter_tree_node_black);
   }
 }
+
+void vieter_tree_node_remove_balanced(vieter_tree_node *node) {}
diff --git a/src/tree/vieter_tree_balancing.h b/src/tree/vieter_tree_balancing.h
index 4492588..007af34 100644
--- a/src/tree/vieter_tree_balancing.h
+++ b/src/tree/vieter_tree_balancing.h
@@ -3,6 +3,17 @@
 
 #include "vieter_tree_node.h"
 
+/*
+ * Ensure the tree remains a valid red-black tree after having inserting the
+ * node.
+ */
 void vieter_tree_node_balance_after_insert(vieter_tree_node *node);
 
+/*
+ * Remove the given node, ensuring the tree remains a valid red-black tree.
+ *
+ * @param node node to remove. This node should have at least one NULL child.
+ */
+void vieter_tree_node_remove_balanced(vieter_tree_node *node);
+
 #endif
diff --git a/src/tree/vieter_tree_internal.h b/src/tree/vieter_tree_internal.h
index dc4ee81..670d45b 100644
--- a/src/tree/vieter_tree_internal.h
+++ b/src/tree/vieter_tree_internal.h
@@ -1,8 +1,10 @@
 #include "vieter_tree.h"
 
+#include <stdbool.h>
+
 /*
  * Inspect whether the tree is still a valid red-black-tree.
  *
- * @return 0 if correct, 1 if not.
+ * @return true if valid, false otherwise.
  */
-int vieter_tree_validate(vieter_tree *tree);
+bool vieter_tree_validate(vieter_tree *tree);
diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c
index 40407a3..b7e82a8 100644
--- a/src/tree/vieter_tree_node.c
+++ b/src/tree/vieter_tree_node.c
@@ -153,14 +153,15 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
   return vieter_tree_ok;
 }
 
-void vieter_tree_node_set(vieter_tree_node *node, vieter_node_flag flag) {
+void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag) {
   node->flags |= flag;
 }
 
-void vieter_tree_node_unset(vieter_tree_node *node, vieter_node_flag flag) {
+void vieter_tree_node_unset(vieter_tree_node *node,
+                            vieter_tree_node_flag flag) {
   node->flags &= ~flag;
 }
 
-bool vieter_tree_node_get(vieter_tree_node *node, vieter_node_flag flag) {
+bool vieter_tree_node_get(vieter_tree_node *node, vieter_tree_node_flag flag) {
   return (node->flags & flag) != 0;
 }
diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h
index 1c88159..5e74280 100644
--- a/src/tree/vieter_tree_node.h
+++ b/src/tree/vieter_tree_node.h
@@ -5,11 +5,9 @@
 #include <stdint.h>
 #include <stdbool.h>
 
-#define VIETER_NODE_BLACK ((uint8_t)1)
-
-typedef enum vieter_node_flag {
-    vieter_node_black = ((uint8_t)1) << 0
-} vieter_node_flag;
+typedef enum vieter_tree_node_flag {
+    vieter_tree_node_black = ((uint8_t)1) << 0
+} vieter_tree_node_flag;
 
 typedef struct vieter_tree_node {
     uint64_t key;
@@ -26,20 +24,24 @@ void vieter_tree_node_free(vieter_tree_node *node);
 
 vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key, void *data);
 
+/*
+ * Return the node representing the requested value.
+ */
 vieter_tree_error vieter_tree_node_search_node(vieter_tree_node **out, vieter_tree_node *root, uint64_t key);
+
 vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root, uint64_t key);
 
 vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, uint64_t key);
 
-void vieter_tree_node_replace(vieter_tree_node *to_replace, vieter_tree_node *replacement);
+void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace, vieter_tree_node *replacement);
 
-void vieter_tree_node_set(vieter_tree_node *node, vieter_node_flag flag);
-void vieter_tree_node_unset(vieter_tree_node *node, vieter_node_flag flag);
-bool vieter_tree_node_get(vieter_tree_node *node, vieter_node_flag flag);
+void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag);
+void vieter_tree_node_unset(vieter_tree_node *node, vieter_tree_node_flag flag);
+bool vieter_tree_node_get(vieter_tree_node *node, vieter_tree_node_flag flag);
 
 void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
                                 vieter_tree_node *child);
 
-int vieter_tree_node_validate(vieter_tree_node *node, uint64_t passed_black_nodes, uint64_t expected_black_nodes);
+bool vieter_tree_node_validate(vieter_tree_node *node, uint64_t passed_black_nodes, uint64_t expected_black_nodes);
 
 #endif
diff --git a/test/tree/test_tree.c b/test/tree/test_tree.c
index c875137..257ec66 100644
--- a/test/tree/test_tree.c
+++ b/test/tree/test_tree.c
@@ -18,7 +18,7 @@ void test_insert() {
     for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_insert(tree, i, NULL) == vieter_tree_ok);
         TEST_SIZE(tree, i + 1);
-        TEST_CHECK(vieter_tree_validate(tree) == 0);
+        TEST_CHECK(vieter_tree_validate(tree));
     }
 
     void *out;
@@ -37,7 +37,7 @@ void test_remove() {
 
     for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_insert(tree, i, NULL) == vieter_tree_ok);
-        TEST_CHECK(vieter_tree_validate(tree) == 0);
+        TEST_CHECK(vieter_tree_validate(tree));
     }
 
     void *out;
@@ -45,7 +45,7 @@ void test_remove() {
     for (uint64_t i = 0; i < 25; i++) {
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_ok);
         TEST_CHECK(vieter_tree_remove(&out, tree, i) == vieter_tree_ok);
-        TEST_CHECK(vieter_tree_validate(tree) == 0);
+        TEST_CHECK(vieter_tree_validate(tree));
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_not_present);
         TEST_SIZE(tree, 250 - i - 1);
     }