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