1#include "structures/tests/test_internal.h"
2
3TEST_GROUP_DECLARE(avl);
4
5struct test_avl_node {
6 int key;
7 struct avl_tree_node node;
8};
9
10static int test_avl_cmp(const struct avl_tree_node *a,
11 const struct avl_tree_node *b) {
12 int ka = avl_entry(a, struct test_avl_node, node)->key;
13 int kb = avl_entry(b, struct test_avl_node, node)->key;
14 return (ka > kb) - (ka < kb);
15}
16
17static int test_avl_cmp_key(const struct avl_tree_node *a, const void *key) {
18 int ka = avl_entry(a, struct test_avl_node, node)->key;
19 int kk = *(const int *) key;
20 return (ka > kk) - (ka < kk);
21}
22
23static const struct avl_tree_node_ops test_avl_ops = {
24 .cmp = test_avl_cmp,
25 .cmp_key = test_avl_cmp_key,
26};
27
28static int node_height(struct avl_tree_node *n) {
29 if (!n)
30 return 0;
31 int lh = node_height(n: n->left);
32 int rh = node_height(n: n->right);
33 return (lh > rh ? lh : rh) + 1;
34}
35
36static bool verify_avl_invariants(struct avl_tree_node *n) {
37 if (!n)
38 return true;
39
40 int lh = node_height(n: n->left);
41 int rh = node_height(n: n->right);
42 int diff = lh - rh;
43 if (diff < -1 || diff > 1)
44 return false;
45
46 if (n->height != (lh > rh ? lh : rh) + 1)
47 return false;
48
49 if (n->left && n->left->parent != n)
50 return false;
51 if (n->right && n->right->parent != n)
52 return false;
53
54 return verify_avl_invariants(n: n->left) && verify_avl_invariants(n: n->right);
55}
56
57TEST_DECLARE_UNIT(avl, rotations_and_balance) {
58 struct avl_tree tree;
59 avl_tree_init(tree: &tree, ops: &test_avl_ops);
60
61 struct test_avl_node nodes[7];
62 int insert_keys[7] = {30, 20, 40, 10, 25, 35, 50};
63
64 for (int i = 0; i < 7; i++) {
65 nodes[i].key = insert_keys[i];
66 avl_tree_insert(tree: &tree, node: &nodes[i].node);
67 TEST_ASSERT(verify_avl_invariants(tree.root));
68 }
69
70 /* Verify search */
71 for (int i = 0; i < 7; i++) {
72 struct avl_tree_node *found = avl_tree_find(tree: &tree, key: &insert_keys[i]);
73 TEST_ASSERT_NONNULL(found);
74 TEST_ASSERT_EQ(avl_entry(found, struct test_avl_node, node)->key,
75 insert_keys[i]);
76 }
77
78 int missing = 999;
79 TEST_ASSERT_NULL(avl_tree_find(&tree, &missing));
80
81 /* In-order traversal must be ascending */
82 struct avl_tree_node *cur = avl_tree_first(tree: &tree);
83 int prev_key = -1;
84 int count = 0;
85 while (cur) {
86 int k = avl_entry(cur, struct test_avl_node, node)->key;
87 TEST_ASSERT_GT_S(k, prev_key);
88 prev_key = k;
89 count++;
90 cur = avl_tree_next(node: cur);
91 }
92 TEST_ASSERT_EQ(count, 7);
93
94 /* removal with successor transplant */
95 for (int i = 0; i < 7; i++) {
96 avl_tree_remove(tree: &tree, node: &nodes[i].node);
97 TEST_ASSERT(verify_avl_invariants(tree.root));
98 TEST_ASSERT_NULL(avl_tree_find(&tree, &insert_keys[i]));
99 }
100
101 TEST_ASSERT(avl_tree_empty(&tree));
102
103 return TEST_SUCCESS;
104}
105