1 /*
2 * Copyright (c) 2019 LK Trusty Authors. All Rights Reserved.
3 *
4 * Permission is hereby granted, free of charge, to any person obtaining
5 * a copy of this software and associated documentation files
6 * (the "Software"), to deal in the Software without restriction,
7 * including without limitation the rights to use, copy, modify, merge,
8 * publish, distribute, sublicense, and/or sell copies of the Software,
9 * and to permit persons to whom the Software is furnished to do so,
10 * subject to the following conditions:
11 *
12 * The above copyright notice and this permission notice shall be
13 * included in all copies or substantial portions of the Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
16 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
17 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
18 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
19 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
20 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
21 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
22 */
23
24 #include <lib/binary_search_tree.h>
25
26 /**
27 * bst_node_rank - Internal helper function
28 * @node: Node to check.
29 *
30 * Return: Rank of @node or 0 if @node is %NULL.
31 */
bst_node_rank(struct bst_node * node)32 static size_t bst_node_rank(struct bst_node *node) {
33 return node ? node->rank : 0;
34 }
35
36 /**
37 * bst_is_right_child - Internal helper function
38 * @node: Node to check.
39 *
40 * Return: %true if @node is the right child of @node->parent. %false if
41 * @node->parent is %NULL or if @node is the left child of
42 * @node->parent.
43 */
bst_is_right_child(struct bst_node * node)44 static bool bst_is_right_child(struct bst_node *node) {
45 DEBUG_ASSERT(node);
46 DEBUG_ASSERT(!node->parent || node->parent->child[0] == node ||
47 node->parent->child[1] == node);
48 return node->parent && node->parent->child[1] == node;
49 }
50
51 /**
52 * bst_parent_ptr - Internal helper function
53 * @root: Tree.
54 * @node: Node in @root.
55 *
56 * Return: Pointer where @node is linked in the tree. If @node is the root node
57 * this is &root->root, otherwise it is the address of child pointer in the
58 * parent node of @node that points to @node.
59 */
bst_parent_ptr(struct bst_root * root,struct bst_node * node)60 static struct bst_node **bst_parent_ptr(struct bst_root *root,
61 struct bst_node *node) {
62 DEBUG_ASSERT(root);
63 DEBUG_ASSERT(node);
64 struct bst_node **parent_ptr = node->parent ?
65 &node->parent->child[bst_is_right_child(node)] : &root->root;
66 DEBUG_ASSERT(*parent_ptr == node);
67 return parent_ptr;
68 }
69
70 /**
71 * bst_link_node - Internal helper function
72 * @parent: Target node.
73 * @is_right_child: Index of child to set.
74 * @child: New child node.
75 *
76 * Set child node in @parent. If @child is not %NULL, update it to point to
77 * @parent.
78 */
bst_link_node(struct bst_node * parent,bool is_right_child,struct bst_node * child)79 static void bst_link_node(struct bst_node *parent,
80 bool is_right_child,
81 struct bst_node *child) {
82 parent->child[is_right_child] = child;
83 if (child) {
84 child->parent = parent;
85 }
86 }
87
88 /**
89 * bst_move_node - Internal helper function
90 * @root: Tree.
91 * @old_node: Node to unlink.
92 * @new_node: Node to link where @old_node was.
93 *
94 * Replace node in @root at @old_node with @new_node.
95 */
bst_move_node(struct bst_root * root,struct bst_node * old_node,struct bst_node * new_node)96 static void bst_move_node(struct bst_root *root,
97 struct bst_node *old_node,
98 struct bst_node *new_node) {
99 DEBUG_ASSERT(root);
100 DEBUG_ASSERT(old_node);
101
102 *bst_parent_ptr(root, old_node) = new_node;
103 if (new_node) {
104 new_node->parent = old_node->parent;
105 }
106 old_node->parent = NULL;
107 }
108
109 /**
110 * bst_rotate - Internal helper function
111 * @root: Tree.
112 * @up: Node to move up.
113 * @down: Node to move down.
114 * @up_was_right_child: %true if @up was the right child of @down.
115 *
116 * Swap nodes @up and @down (pictured for up_was_right_child==false):
117 *
118 * down up
119 * / \ / \
120 * up C A down
121 * / \ / \
122 * A B B C
123 *
124 * Caller is responsible for updating the rank of the moved nodes.
125 */
bst_rotate(struct bst_root * root,struct bst_node * up,struct bst_node * down,bool up_was_right_child)126 static void bst_rotate(struct bst_root *root, struct bst_node *up,
127 struct bst_node *down, bool up_was_right_child) {
128 DEBUG_ASSERT(down->child[up_was_right_child] == up);
129 struct bst_node *move_subtree = up->child[!up_was_right_child];
130 bst_move_node(root, down, up);
131 bst_link_node(down, up_was_right_child, move_subtree);
132 bst_link_node(up, !up_was_right_child, down);
133 }
134
135 /**
136 * bst_rotate_insert - Internal helper function
137 * @root: Tree.
138 * @up1: Node to move up if a single rotate is enough.
139 * @down: Node to move down.
140 * @up_was_right_child: %true if @up1 was the right child of @down.
141 *
142 * Rotate sub-tree (once or twice) after insert and update ranks.
143 */
bst_rotate_insert(struct bst_root * root,struct bst_node * up1,struct bst_node * down,bool up_was_right_child)144 static void bst_rotate_insert(struct bst_root *root, struct bst_node *up1,
145 struct bst_node *down, bool up_was_right_child) {
146 DEBUG_ASSERT(down->child[up_was_right_child] == up1);
147 DEBUG_ASSERT(up1->rank == down->rank);
148 DEBUG_ASSERT(down->rank >=
149 bst_node_rank(down->child[!up_was_right_child]) + 2);
150 struct bst_node *up2 = up1->child[!up_was_right_child];
151 if (bst_node_rank(up2) >= down->rank - 1) {
152 DEBUG_ASSERT(bst_node_rank(up2) == down->rank - 1);
153 /*
154 * Swap nodes @up2 and @up1 then @up2 and @down
155 * (pictured for up_was_right_child==false):
156 *
157 * down down up2
158 * / \ / \ / \
159 * up1 D up2 D up1 down
160 * / \ / \ / \ / \
161 * A up2 up1 C A B C D
162 * / \ / \
163 * B C A B
164 */
165 up2->rank++;
166 DEBUG_ASSERT(up1->rank == up2->rank);
167 bst_rotate(root, up2, up1, !up_was_right_child);
168 up1->rank--;
169 bst_rotate(root, up2, down, up_was_right_child);
170 down->rank--;
171 } else {
172 /*
173 * Swap nodes @up1 and @down (pictured for up_was_right_child==false):
174 *
175 * down up1
176 * / \ / \
177 * up1 C A down
178 * / \ / \
179 * A B B C
180 */
181 bst_rotate(root, up1, down, up_was_right_child);
182 down->rank--;
183 }
184 }
185
186 /**
187 * bst_update_rank_insert - Internal helper function
188 * @root: Tree.
189 * @node: Node to start scan at.
190 *
191 * Promote nodes and/or rotate sub-trees to make @root a valid WAVL tree again.
192 */
bst_update_rank_insert(struct bst_root * root,struct bst_node * node)193 void bst_update_rank_insert(struct bst_root *root, struct bst_node *node) {
194 size_t rank;
195
196 DEBUG_ASSERT(root);
197 DEBUG_ASSERT(node);
198 DEBUG_ASSERT(node->rank == 1); /* Inserted node must have rank 1 */
199
200 while (node) {
201 bool is_right_child = bst_is_right_child(node);
202
203 /*
204 * At this point the rank of @node is 1 greater than it was before the
205 * insert.
206 *
207 * For the tree to be valid, the parent of any node is allowed to a
208 * rank 1 or 2 greater than its child nodes. Assuming the tree was valid
209 * before the insert, the @node->parent currently has the same rank as
210 * @node or it has a rank one grater than the rank of @node. Incremeting
211 * the rank @node->parent to be 2 greater than the rank of @node would
212 * be unnecessary as it could not have had that rank already. Leaving
213 * the rank of @node->parent at the same rank as @node would result in
214 * en invalid tree. That means that the rank of @node->parent should now
215 * be 1 greater than the rank of @node (if that is possible).
216 */
217 rank = node->rank + 1;
218 node = node->parent;
219
220 if (!node || node->rank >= rank) {
221 DEBUG_ASSERT(!node || node->rank == rank);
222 /*
223 * Stop updating if we have reached the root, or a node that already
224 * has a rank greater than the node child node we inserted or
225 * updated as the tree is now valid.
226 */
227 return;
228 }
229 DEBUG_ASSERT(node->rank + 1 == rank);
230 if (bst_node_rank(node->child[!is_right_child]) + 2 < rank) {
231 /*
232 * Rank of @node cannot be incremented. This means it can be moved
233 * down and demoted instead.
234 *
235 * The tree can be rotated as pictured below. (a is @node which
236 * could not be promoted. Numbers are known relative ranks.)
237 *
238 * If rank of c is 2 less than the rank of a (b is inserted or
239 * promoted node):
240 *
241 * a2 b2
242 * / \ / \
243 * b2 D0 => A a1
244 * / \ / \
245 * A c0 c0 D0
246 * / \ / \
247 * B C B C
248 *
249 * If rank of c is 1 less than the rank of a (b is promoted node, c
250 * is inserted or promoted node):
251 * a2 a2 __c2__
252 * / \ / \ / \
253 * b2 D0 c2 D0 b1 a1
254 * / \ => / \ => / \ / \
255 * A0 c1 b1 C A0 B C D0
256 * / \ / \
257 * B C A0 B
258 */
259 bst_rotate_insert(root, node->child[is_right_child], node,
260 is_right_child);
261 return;
262 }
263 node->rank = rank;
264 }
265 }
266
267 /**
268 * bst_rotate_delete - Internal helper function
269 * @root: Tree.
270 * @up1: Node to move up if a single rotate is enough.
271 * @down: Node to move down.
272 * @up_was_right_child: %true if @up1 was the right child of @down.
273 *
274 * Rotate sub-tree (once or twice) after delete and update ranks.
275 */
bst_rotate_delete(struct bst_root * root,struct bst_node * up1,struct bst_node * down,bool up_was_right_child)276 static void bst_rotate_delete(struct bst_root *root, struct bst_node *up1,
277 struct bst_node *down, bool up_was_right_child) {
278 DEBUG_ASSERT(down->child[up_was_right_child] == up1);
279 DEBUG_ASSERT(up1->rank == down->rank - 1);
280 DEBUG_ASSERT(down->rank ==
281 bst_node_rank(down->child[!up_was_right_child]) + 3);
282 struct bst_node *up2 = up1->child[!up_was_right_child];
283 if (bst_node_rank(up1->child[up_was_right_child]) <= down->rank - 3) {
284 DEBUG_ASSERT(bst_node_rank(up2) == down->rank - 2);
285 /*
286 * Swap nodes @up2 and @up1 then @up2 and @down
287 * (pictured for up_was_right_child==false):
288 *
289 * down(0) down up2(0)
290 * / \ / \ / \
291 * up1(-1) D(-3) up2 D up1(-2) down(-2)
292 * / \ / \ / \ / \
293 * A(-3) up2(-2) up1 C A(-3) B C D(-3)
294 * / \ / \
295 * B C A(-3) B
296 */
297 DEBUG_ASSERT(up1->rank == down->rank - 1);
298 DEBUG_ASSERT(up2->rank == down->rank - 2);
299 bst_rotate(root, up2, up1, !up_was_right_child);
300 bst_rotate(root, up2, down, up_was_right_child);
301 up2->rank += 2;
302 up1->rank--;
303 down->rank -= 2;
304 } else {
305 /*
306 * Swap nodes @up1 and @down (pictured for up_was_right_child==false):
307 *
308 * down(0) up1(0)
309 * / \ / \
310 * up1(-1) C(-3) A(-2) down(-1)
311 * / \ / \
312 * A(-2) B(-2/-3) B(-2/-3) C(-3)
313 */
314 bst_rotate(root, up1, down, up_was_right_child);
315 up1->rank++;
316 down->rank--;
317 if (bst_node_rank(down->child[0]) == down->rank - 2 &&
318 bst_node_rank(down->child[1]) == down->rank - 2) {
319 /* Demote down if possible. (Required if down is a leaf node) */
320 down->rank--;
321 }
322 }
323 }
324
325 /**
326 * bst_update_rank_delete - Internal helper function
327 * @root: Tree.
328 * @node: Node to start scan at. This is the parent of the node that
329 * was removed from the tree. Note that the removed node will
330 * be a different node than the node passed to bst_delete if
331 * that node had two children.
332 * @is_right_child: %true if the right child of @node was deleted.
333 *
334 * Demote nodes and/or rotate sub-trees to make @root a valid WAVL tree again.
335 */
bst_update_rank_delete(struct bst_root * root,struct bst_node * node,bool is_right_child)336 static void bst_update_rank_delete(struct bst_root *root, struct bst_node *node,
337 bool is_right_child) {
338 DEBUG_ASSERT(root);
339 DEBUG_ASSERT(node);
340
341 DEBUG_ASSERT(bst_node_rank(node->child[is_right_child]) <=
342 bst_node_rank(node->child[!is_right_child]));
343
344 while (node) {
345 DEBUG_ASSERT(node->rank > bst_node_rank(node->child[!is_right_child]));
346 DEBUG_ASSERT(node->rank - 1 >
347 bst_node_rank(node->child[is_right_child]));
348 DEBUG_ASSERT(node->rank <=
349 bst_node_rank(node->child[!is_right_child]) + 2);
350
351 /*
352 * At this point the rank of @node->child[is_right_child] has been
353 * decremented. We may need to also decrement the rank of @node.
354 */
355 if (!node->child[0] && !node->child[1]) {
356 /* Always demote leaf node (from 2 to 1) */
357
358 /* We should not be in this function if the rank is alrady 1 */
359 DEBUG_ASSERT(node->rank == 2);
360 } else if (node->rank <=
361 bst_node_rank(node->child[is_right_child]) + 2) {
362 /*
363 * If rank of @node does not need to change then we now have a valid
364 * tree.
365 */
366 return;
367 }
368 DEBUG_ASSERT(node->rank > 1);
369 node->rank--;
370 if (node->rank <= bst_node_rank(node->child[!is_right_child])) {
371 /* We demoted @node, but it is now invalid on the other side */
372 DEBUG_ASSERT(node->rank ==
373 bst_node_rank(node->child[!is_right_child]));
374 if (bst_node_rank(node->child[!is_right_child]->child[0]) ==
375 node->rank - 2 &&
376 bst_node_rank(node->child[!is_right_child]->child[1]) ==
377 node->rank - 2) {
378 /* If the other child can be demoted, demote it and move up */
379 node->child[!is_right_child]->rank--;
380 } else {
381 /*
382 * If the other child can not be demoted, rotate instead. This
383 * will produce a valid tree without changing the rank of the
384 * node linked at the current spot @node in the tree.
385 *
386 * Undo demotion as current bst_rotate_delete implemention
387 * assumes node rank is unchanged.
388 */
389 node->rank++;
390
391 bst_rotate_delete(root, node->child[!is_right_child], node,
392 !is_right_child);
393 return;
394 }
395 }
396 is_right_child = bst_is_right_child(node);
397 node = node->parent;
398 }
399 }
400
401 /**
402 * bst_find_edge - Internal helper function
403 * @node: Node to start search at.
404 * @edge: Direction if search.
405 *
406 * Return: leftmost (if @edge is %false) or rightmost (if @edge is %true) node
407 * in subtree with @node as root.
408 */
bst_find_edge(struct bst_node * node,bool edge)409 static struct bst_node *bst_find_edge(struct bst_node *node, bool edge) {
410 struct bst_node *saved_node;
411
412 DEBUG_ASSERT(node);
413
414 do {
415 saved_node = node;
416 node = node->child[edge];
417 } while (node);
418
419 return saved_node;
420 }
421
422 /**
423 * bst_delete_all_helper - Internal helper function
424 * @root: Tree.
425 * @node: Node to delete (most be the leftmost node in @root).
426 *
427 * Helper function to delete leftmost node in @root, assuming all other nodes
428 * will be deleted next.
429 */
bst_delete_all_helper(struct bst_root * root,struct bst_node * node)430 void bst_delete_all_helper(struct bst_root *root, struct bst_node *node) {
431 DEBUG_ASSERT(root);
432 DEBUG_ASSERT(node);
433 DEBUG_ASSERT(!node->child[0]);
434 bst_move_node(root, node, node->child[1]);
435 }
436
bst_delete(struct bst_root * root,struct bst_node * node)437 void bst_delete(struct bst_root *root, struct bst_node *node) {
438 DEBUG_ASSERT(root);
439 DEBUG_ASSERT(node);
440
441 struct bst_node *new_child;
442 bool node_is_right_child = bst_is_right_child(node);
443 struct bst_node *update_rank_start = node->parent;
444 bool update_rank_is_right_child = node_is_right_child;
445 if (!node->child[0]) {
446 /*
447 * If @node has no left child, link its right child in its place. (The
448 * right child could be %NULL in this case)
449 */
450 new_child = node->child[1];
451 } else if (!node->child[1]) {
452 /*
453 * If @node has no right child, link its left child in its place.
454 */
455 DEBUG_ASSERT(node->child[0]);
456 new_child = node->child[0];
457 } else {
458 /*
459 * If @node has both left and right children, delete (from the tree
460 * structure point of view) the left-most node in the right sub-tree or
461 * the right-most node in the left sub-tree instead. Either side would
462 * work.
463 */
464 struct bst_node *edge_node = bst_find_edge(
465 node->child[!node_is_right_child], node_is_right_child);
466 struct bst_node *edge_child = edge_node->child[!node_is_right_child];
467 update_rank_start = edge_node->parent;
468 update_rank_is_right_child = bst_is_right_child(edge_node);
469 if (update_rank_start == node) {
470 update_rank_start = edge_node;
471 update_rank_is_right_child = !node_is_right_child;
472 }
473 DEBUG_ASSERT(update_rank_start);
474 bst_move_node(root, edge_node, edge_child);
475
476 new_child = edge_node;
477 DEBUG_ASSERT(new_child);
478 bst_link_node(new_child, 0, node->child[0]);
479 bst_link_node(new_child, 1, node->child[1]);
480 new_child->rank = node->rank;
481 }
482 bst_move_node(root, node, new_child);
483 node->rank = 0;
484 if (update_rank_start) {
485 bst_update_rank_delete(root, update_rank_start, update_rank_is_right_child);
486 }
487 }
488
489 /**
490 * bst_prev_next - Internal helper function
491 * @root: Tree.
492 * @node: Node to move from.
493 * @dir_next: Directon to move.
494 *
495 * Return: If @node is %NULL and @dir_next is %false, right-most node in @root.
496 * If @node is %NULL and @dir_next is %true, left-most node in @root.
497 * If @node is not %NULL and @dir_next is %false, right-most node to the
498 * left of @node.
499 * If @node is not %NULL and @dir_next is %true, left-most node to the
500 * right of @node.
501 * %NULL if the node described above does not exist.
502 */
bst_prev_next(const struct bst_root * root,struct bst_node * node,bool dir_next)503 static struct bst_node *bst_prev_next(const struct bst_root *root,
504 struct bst_node *node,
505 bool dir_next) {
506 DEBUG_ASSERT(root);
507
508 struct bst_node *next_child = node ? node->child[dir_next] : root->root;
509
510 if (!node && !next_child) {
511 return NULL; /* Empty tree */
512 }
513
514 /*
515 * Comments below assume @dir_next is %true. For the @dir_next is %false
516 * case, swap left and right.
517 */
518 if (next_child) {
519 /* There is a right child, return the left-most node in that subtree */
520 return bst_find_edge(next_child, !dir_next);
521 } else {
522 /* No right child, next node is the first right parent */
523 struct bst_node *next_parent = node;
524 while (bst_is_right_child(next_parent) == dir_next) {
525 next_parent = next_parent->parent;
526 if (!next_parent) {
527 return NULL;
528 }
529 }
530 return next_parent->parent;
531 }
532 }
533
bst_prev(struct bst_root * root,struct bst_node * node)534 struct bst_node *bst_prev(struct bst_root *root, struct bst_node *node) {
535 return bst_prev_next(root, node, false);
536 }
537
bst_next(const struct bst_root * root,struct bst_node * node)538 struct bst_node *bst_next(const struct bst_root *root, struct bst_node *node) {
539 return bst_prev_next(root, node, true);
540 }
541