1#include "structures/tests/test_internal.h"
2
3TEST_GROUP_DEFINE(rbt);
4
5struct test_rbt_node {
6 int key;
7 struct rbt_node node;
8};
9
10static size_t test_rbt_get_data(struct rbt_node *node) {
11 return (size_t) rbt_entry(node, struct test_rbt_node, node)->key;
12}
13
14static int32_t test_rbt_compare(const struct rbt_node *a,
15 const struct rbt_node *b) {
16 int a_key = rbt_entry(a, struct test_rbt_node, node)->key;
17 int b_key = rbt_entry(b, struct test_rbt_node, node)->key;
18 return (a_key > b_key) - (a_key < b_key);
19}
20
21static int rbt_black_height(const struct rbt_node *node) {
22 if (!node)
23 return 1;
24
25 if (rbt_color(node) == TREE_NODE_RED &&
26 ((node->left && rbt_color(node: node->left) == TREE_NODE_RED) ||
27 (node->right && rbt_color(node: node->right) == TREE_NODE_RED)))
28 return -1;
29
30 int left_height = rbt_black_height(node: node->left);
31 int right_height = rbt_black_height(node: node->right);
32 if (left_height < 0 || left_height != right_height)
33 return -1;
34
35 return left_height + (rbt_color(node) == TREE_NODE_BLACK);
36}
37
38static bool verify_rbt_invariants(const struct rbt_node *node, const int *min,
39 const int *max) {
40 if (!node)
41 return true;
42
43 int key = rbt_entry(node, struct test_rbt_node, node)->key;
44 if ((min && key <= *min) || (max && key >= *max))
45 return false;
46 if ((node->left && rbt_parent(node: node->left) != node) ||
47 (node->right && rbt_parent(node: node->right) != node))
48 return false;
49
50 return verify_rbt_invariants(node: node->left, min, max: &key) &&
51 verify_rbt_invariants(node: node->right, min: &key, max);
52}
53
54TEST_DEFINE_UNIT(rbt, delete_black_leaf_rebalances) {
55 struct rbt tree;
56 struct test_rbt_node nodes[] = {
57 {.key = 8}, {.key = 75}, {.key = 80}, {.key = 31}};
58 int expected[] = {8, 31, 75};
59
60 rbt_init(t: &tree, get_data: test_rbt_get_data, compare: test_rbt_compare);
61 for (size_t i = 0; i < TEST_ARRAY_LEN(nodes); i++) {
62 rbt_init_node(n: &nodes[i].node);
63 rbt_insert(tree: &tree, new_node: &nodes[i].node);
64 }
65
66 rbt_delete(tree: &tree, z: &nodes[2].node);
67
68 TEST_ASSERT_NONNULL(tree.root);
69 TEST_ASSERT_EQ(rbt_color(tree.root), TREE_NODE_BLACK);
70 TEST_ASSERT_GE(rbt_black_height(tree.root), 0);
71 TEST_ASSERT(verify_rbt_invariants(tree.root, NULL, NULL));
72
73 size_t index = 0;
74 struct rbt_node *node;
75 rbt_for_each(node, &tree) {
76 TEST_ASSERT_LT(index, TEST_ARRAY_LEN(expected));
77 TEST_ASSERT_EQ(rbt_entry(node, struct test_rbt_node, node)->key,
78 expected[index++]);
79 }
80 TEST_ASSERT_EQ(index, TEST_ARRAY_LEN(expected));
81 TEST_ASSERT_NULL(rbt_search(&tree, 80));
82
83 return TEST_SUCCESS;
84}
85