1#include <console/panic.h>
2#include <math/min_max.h>
3#include <structures/avl.h>
4
5static inline int avl_height(const struct avl_tree_node *node) {
6 return node ? node->height : 0;
7}
8
9static inline int avl_balance_factor(const struct avl_tree_node *node) {
10 return node ? avl_height(node: node->left) - avl_height(node: node->right) : 0;
11}
12
13static inline void avl_update_height(struct avl_tree_node *node) {
14 int lh = avl_height(node: node->left);
15 int rh = avl_height(node: node->right);
16 node->height = 1 + MAX(lh, rh);
17}
18
19static void avl_check_cycle(struct avl_tree_node *node, const char *caller) {
20 struct avl_tree_node *tortoise = node;
21 struct avl_tree_node *hare = node;
22
23 while (hare && hare->left != NULL) {
24 tortoise = tortoise ? tortoise->left : NULL;
25 hare = hare->left->left;
26
27 if (tortoise == NULL || hare == NULL)
28 break;
29
30 if (tortoise == hare) {
31 panic(CORRUPTION,
32 "Cycle detected in %s (left spine): tortoise == hare at "
33 "node=%p",
34 caller, (void *) tortoise);
35 break;
36 }
37 }
38
39 tortoise = node;
40 hare = node;
41
42 while (hare && hare->right != NULL) {
43 tortoise = tortoise ? tortoise->right : NULL;
44 hare = hare->right->right;
45
46 if (tortoise == NULL || hare == NULL)
47 break;
48
49 if (tortoise == hare) {
50 panic(CORRUPTION,
51 "Cycle detected in %s (right spine): tortoise == hare at "
52 "node=%p",
53 caller, (void *) tortoise);
54 break;
55 }
56 }
57}
58
59/* rotations
60 *
61 * left_rotate(x):
62 * x y
63 * \ / \
64 * y => x c
65 * / \ \
66 * b c b
67 */
68static void avl_left_rotate(struct avl_tree *tree, struct avl_tree_node *x) {
69 struct avl_tree_node *y = x->right;
70 x->right = y->left;
71 if (y->left)
72 y->left->parent = x;
73
74 y->parent = x->parent;
75 if (!x->parent)
76 tree->root = y;
77 else if (x == x->parent->left)
78 x->parent->left = y;
79 else
80 x->parent->right = y;
81
82 y->left = x;
83 x->parent = y;
84
85 avl_update_height(node: x);
86 avl_update_height(node: y);
87}
88
89static void avl_right_rotate(struct avl_tree *tree, struct avl_tree_node *y) {
90 struct avl_tree_node *x = y->left;
91 y->left = x->right;
92 if (x->right)
93 x->right->parent = y;
94
95 x->parent = y->parent;
96 if (!y->parent)
97 tree->root = x;
98 else if (y == y->parent->left)
99 y->parent->left = x;
100 else
101 y->parent->right = x;
102
103 x->right = y;
104 y->parent = x;
105
106 avl_update_height(node: y);
107 avl_update_height(node: x);
108}
109
110/* Walk up from `node` rebalancing as needed. */
111static void avl_rebalance(struct avl_tree *tree, struct avl_tree_node *node) {
112 while (node) {
113 struct avl_tree_node *parent = node->parent;
114 avl_update_height(node);
115 int bf = avl_balance_factor(node);
116
117 if (bf > 1) {
118 if (avl_balance_factor(node: node->left) < 0)
119 avl_left_rotate(tree, x: node->left);
120 avl_right_rotate(tree, y: node);
121 } else if (bf < -1) {
122 if (avl_balance_factor(node: node->right) > 0)
123 avl_right_rotate(tree, y: node->right);
124 avl_left_rotate(tree, x: node);
125 }
126
127 node = parent;
128 }
129}
130
131void avl_tree_init(struct avl_tree *tree, const struct avl_tree_node_ops *ops) {
132 tree->root = NULL;
133 tree->ops = ops;
134}
135
136void avl_tree_insert(struct avl_tree *tree, struct avl_tree_node *new_node) {
137 avl_check_cycle(node: tree->root, caller: __func__);
138 new_node->left = new_node->right = NULL;
139 new_node->height = 1;
140
141 if (!tree->root) {
142 new_node->parent = NULL;
143 tree->root = new_node;
144 return;
145 }
146
147 struct avl_tree_node *current = tree->root;
148 struct avl_tree_node *parent = NULL;
149 int last_cmp = 0;
150 while (current) {
151 parent = current;
152 /* ties go right: cmp == 0 => take right branch */
153 last_cmp = tree->ops->cmp(new_node, current);
154 current = (last_cmp < 0) ? current->left : current->right;
155 }
156
157 new_node->parent = parent;
158 if (last_cmp < 0)
159 parent->left = new_node;
160 else
161 parent->right = new_node;
162
163 avl_rebalance(tree, node: parent);
164}
165
166static void avl_transplant(struct avl_tree *tree, struct avl_tree_node *u,
167 struct avl_tree_node *v) {
168 if (!u->parent)
169 tree->root = v;
170 else if (u == u->parent->left)
171 u->parent->left = v;
172 else
173 u->parent->right = v;
174 if (v)
175 v->parent = u->parent;
176}
177
178static struct avl_tree_node *avl_min_node(struct avl_tree_node *node) {
179 while (node && node->left)
180 node = node->left;
181 return node;
182}
183
184void avl_tree_remove(struct avl_tree *tree, struct avl_tree_node *node) {
185 avl_check_cycle(node: tree->root, caller: __func__);
186 struct avl_tree_node *rebalance_start;
187
188 if (!node->left) {
189 rebalance_start = node->parent;
190 avl_transplant(tree, u: node, v: node->right);
191 } else if (!node->right) {
192 rebalance_start = node->parent;
193 avl_transplant(tree, u: node, v: node->left);
194 } else {
195 /* splice in the in-order successor. */
196 struct avl_tree_node *succ = avl_min_node(node: node->right);
197
198 if (succ->parent == node) {
199 rebalance_start = succ;
200 } else {
201 rebalance_start = succ->parent;
202 avl_transplant(tree, u: succ, v: succ->right);
203 succ->right = node->right;
204 succ->right->parent = succ;
205 }
206
207 avl_transplant(tree, u: node, v: succ);
208 succ->left = node->left;
209 succ->left->parent = succ;
210 }
211
212 if (rebalance_start)
213 avl_rebalance(tree, node: rebalance_start);
214}
215
216struct avl_tree_node *avl_tree_find(const struct avl_tree *tree,
217 const void *key) {
218 avl_check_cycle(node: tree->root, caller: __func__);
219 struct avl_tree_node *node = tree->root;
220 while (node) {
221 int c = tree->ops->cmp_key(node, key);
222 if (c == 0)
223 return node;
224 node = (c > 0) ? node->left : node->right;
225 }
226 return NULL;
227}
228
229struct avl_tree_node *avl_tree_first(const struct avl_tree *tree) {
230 avl_check_cycle(node: tree->root, caller: __func__);
231 return avl_min_node(node: tree->root);
232}
233
234struct avl_tree_node *avl_tree_last(const struct avl_tree *tree) {
235 avl_check_cycle(node: tree->root, caller: __func__);
236 struct avl_tree_node *node = tree->root;
237 while (node && node->right)
238 node = node->right;
239 return node;
240}
241
242struct avl_tree_node *avl_tree_next(const struct avl_tree_node *node) {
243 if (!node)
244 return NULL;
245 avl_check_cycle(node: (struct avl_tree_node *) node, caller: __func__);
246 if (node->right)
247 return avl_min_node(node: node->right);
248
249 /* climb until we come up from a left child */
250 const struct avl_tree_node *p = node->parent;
251 while (p && node == p->right) {
252 node = p;
253 p = p->parent;
254 }
255 return (struct avl_tree_node *) p;
256}
257
258struct avl_tree_node *avl_tree_prev(const struct avl_tree_node *node) {
259 if (!node)
260 return NULL;
261 avl_check_cycle(node: (struct avl_tree_node *) node, caller: __func__);
262 if (node->left) {
263 struct avl_tree_node *n = node->left;
264 while (n->right)
265 n = n->right;
266 return n;
267 }
268
269 const struct avl_tree_node *p = node->parent;
270 while (p && node == p->left) {
271 node = p;
272 p = p->parent;
273 }
274 return (struct avl_tree_node *) p;
275}
276
277static int avl_validate_node(const struct avl_tree *tree,
278 const struct avl_tree_node *node,
279 const struct avl_tree_node *parent,
280 int *height_out) {
281 if (!node) {
282 *height_out = 0;
283 return 1;
284 }
285
286 if (node->parent != parent) {
287 panic(CORRUPTION, "AVL parent link broken");
288 return 0;
289 }
290
291 if (node->left && tree->ops->cmp(node->left, node) > 0) {
292 panic(CORRUPTION, "AVL ordering violation: left child > parent");
293 return 0;
294 }
295 if (node->right && tree->ops->cmp(node->right, node) < 0) {
296 panic(CORRUPTION, "AVL ordering violation: right child < parent");
297 return 0;
298 }
299
300 int lh = 0, rh = 0;
301 if (!avl_validate_node(tree, node: node->left, parent: node, height_out: &lh))
302 return 0;
303 if (!avl_validate_node(tree, node: node->right, parent: node, height_out: &rh))
304 return 0;
305
306 int expected = 1 + MAX(lh, rh);
307 if (node->height != expected) {
308 panic(CORRUPTION, "AVL height mismatch: stored %d, actual %d",
309 node->height, expected);
310 return 0;
311 }
312
313 int bf = lh - rh;
314 if (bf < -1 || bf > 1) {
315 panic(CORRUPTION, "AVL balance violation: bf = %d", bf);
316 return 0;
317 }
318
319 *height_out = expected;
320 return 1;
321}
322
323int avl_tree_validate(const struct avl_tree *tree) {
324 avl_check_cycle(node: tree->root, caller: __func__);
325 int h = 0;
326 return avl_validate_node(tree, node: tree->root, NULL, height_out: &h);
327}
328