| /* |
| * Copyright (c) 2019 LK Trusty Authors. All Rights Reserved. |
| * |
| * Permission is hereby granted, free of charge, to any person obtaining |
| * a copy of this software and associated documentation files |
| * (the "Software"), to deal in the Software without restriction, |
| * including without limitation the rights to use, copy, modify, merge, |
| * publish, distribute, sublicense, and/or sell copies of the Software, |
| * and to permit persons to whom the Software is furnished to do so, |
| * subject to the following conditions: |
| * |
| * The above copyright notice and this permission notice shall be |
| * included in all copies or substantial portions of the Software. |
| * |
| * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, |
| * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF |
| * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. |
| * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY |
| * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, |
| * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE |
| * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. |
| */ |
| |
| #include <lib/binary_search_tree.h> |
| |
| /** |
| * bst_node_rank - Internal helper function |
| * @node: Node to check. |
| * |
| * Return: Rank of @node or 0 if @node is %NULL. |
| */ |
| static size_t bst_node_rank(struct bst_node *node) { |
| return node ? node->rank : 0; |
| } |
| |
| /** |
| * bst_is_right_child - Internal helper function |
| * @node: Node to check. |
| * |
| * Return: %true if @node is the right child of @node->parent. %false if |
| * @node->parent is %NULL or if @node is the left child of |
| * @node->parent. |
| */ |
| static bool bst_is_right_child(struct bst_node *node) { |
| DEBUG_ASSERT(node); |
| DEBUG_ASSERT(!node->parent || node->parent->child[0] == node || |
| node->parent->child[1] == node); |
| return node->parent && node->parent->child[1] == node; |
| } |
| |
| /** |
| * bst_parent_ptr - Internal helper function |
| * @root: Tree. |
| * @node: Node in @root. |
| * |
| * Return: Pointer where @node is linked in the tree. If @node is the root node |
| * this is &root->root, otherwise it is the address of child pointer in the |
| * parent node of @node that points to @node. |
| */ |
| static struct bst_node **bst_parent_ptr(struct bst_root *root, |
| struct bst_node *node) { |
| DEBUG_ASSERT(root); |
| DEBUG_ASSERT(node); |
| struct bst_node **parent_ptr = node->parent ? |
| &node->parent->child[bst_is_right_child(node)] : &root->root; |
| DEBUG_ASSERT(*parent_ptr == node); |
| return parent_ptr; |
| } |
| |
| /** |
| * bst_link_node - Internal helper function |
| * @parent: Target node. |
| * @is_right_child: Index of child to set. |
| * @child: New child node. |
| * |
| * Set child node in @parent. If @child is not %NULL, update it to point to |
| * @parent. |
| */ |
| static void bst_link_node(struct bst_node *parent, |
| bool is_right_child, |
| struct bst_node *child) { |
| parent->child[is_right_child] = child; |
| if (child) { |
| child->parent = parent; |
| } |
| } |
| |
| /** |
| * bst_move_node - Internal helper function |
| * @root: Tree. |
| * @old_node: Node to unlink. |
| * @new_node: Node to link where @old_node was. |
| * |
| * Replace node in @root at @old_node with @new_node. |
| */ |
| static void bst_move_node(struct bst_root *root, |
| struct bst_node *old_node, |
| struct bst_node *new_node) { |
| DEBUG_ASSERT(root); |
| DEBUG_ASSERT(old_node); |
| |
| *bst_parent_ptr(root, old_node) = new_node; |
| if (new_node) { |
| new_node->parent = old_node->parent; |
| } |
| old_node->parent = NULL; |
| } |
| |
| /** |
| * bst_rotate - Internal helper function |
| * @root: Tree. |
| * @up: Node to move up. |
| * @down: Node to move down. |
| * @up_was_right_child: %true if @up was the right child of @down. |
| * |
| * Swap nodes @up and @down (pictured for up_was_right_child==false): |
| * |
| * down up |
| * / \ / \ |
| * up C A down |
| * / \ / \ |
| * A B B C |
| * |
| * Caller is responsible for updating the rank of the moved nodes. |
| */ |
| static void bst_rotate(struct bst_root *root, struct bst_node *up, |
| struct bst_node *down, bool up_was_right_child) { |
| DEBUG_ASSERT(down->child[up_was_right_child] == up); |
| struct bst_node *move_subtree = up->child[!up_was_right_child]; |
| bst_move_node(root, down, up); |
| bst_link_node(down, up_was_right_child, move_subtree); |
| bst_link_node(up, !up_was_right_child, down); |
| } |
| |
| /** |
| * bst_rotate_insert - Internal helper function |
| * @root: Tree. |
| * @up1: Node to move up if a single rotate is enough. |
| * @down: Node to move down. |
| * @up_was_right_child: %true if @up1 was the right child of @down. |
| * |
| * Rotate sub-tree (once or twice) after insert and update ranks. |
| */ |
| static void bst_rotate_insert(struct bst_root *root, struct bst_node *up1, |
| struct bst_node *down, bool up_was_right_child) { |
| DEBUG_ASSERT(down->child[up_was_right_child] == up1); |
| DEBUG_ASSERT(up1->rank == down->rank); |
| DEBUG_ASSERT(down->rank >= |
| bst_node_rank(down->child[!up_was_right_child]) + 2); |
| struct bst_node *up2 = up1->child[!up_was_right_child]; |
| if (bst_node_rank(up2) >= down->rank - 1) { |
| DEBUG_ASSERT(bst_node_rank(up2) == down->rank - 1); |
| /* |
| * Swap nodes @up2 and @up1 then @up2 and @down |
| * (pictured for up_was_right_child==false): |
| * |
| * down down up2 |
| * / \ / \ / \ |
| * up1 D up2 D up1 down |
| * / \ / \ / \ / \ |
| * A up2 up1 C A B C D |
| * / \ / \ |
| * B C A B |
| */ |
| up2->rank++; |
| DEBUG_ASSERT(up1->rank == up2->rank); |
| bst_rotate(root, up2, up1, !up_was_right_child); |
| up1->rank--; |
| bst_rotate(root, up2, down, up_was_right_child); |
| down->rank--; |
| } else { |
| /* |
| * Swap nodes @up1 and @down (pictured for up_was_right_child==false): |
| * |
| * down up1 |
| * / \ / \ |
| * up1 C A down |
| * / \ / \ |
| * A B B C |
| */ |
| bst_rotate(root, up1, down, up_was_right_child); |
| down->rank--; |
| } |
| } |
| |
| /** |
| * bst_update_rank_insert - Internal helper function |
| * @root: Tree. |
| * @node: Node to start scan at. |
| * |
| * Promote nodes and/or rotate sub-trees to make @root a valid WAVL tree again. |
| */ |
| void bst_update_rank_insert(struct bst_root *root, struct bst_node *node) { |
| size_t rank; |
| |
| DEBUG_ASSERT(root); |
| DEBUG_ASSERT(node); |
| DEBUG_ASSERT(node->rank == 1); /* Inserted node must have rank 1 */ |
| |
| while (node) { |
| bool is_right_child = bst_is_right_child(node); |
| |
| /* |
| * At this point the rank of @node is 1 greater than it was before the |
| * insert. |
| * |
| * For the tree to be valid, the parent of any node is allowed to a |
| * rank 1 or 2 greater than its child nodes. Assuming the tree was valid |
| * before the insert, the @node->parent currently has the same rank as |
| * @node or it has a rank one grater than the rank of @node. Incremeting |
| * the rank @node->parent to be 2 greater than the rank of @node would |
| * be unnecessary as it could not have had that rank already. Leaving |
| * the rank of @node->parent at the same rank as @node would result in |
| * en invalid tree. That means that the rank of @node->parent should now |
| * be 1 greater than the rank of @node (if that is possible). |
| */ |
| rank = node->rank + 1; |
| node = node->parent; |
| |
| if (!node || node->rank >= rank) { |
| DEBUG_ASSERT(!node || node->rank == rank); |
| /* |
| * Stop updating if we have reached the root, or a node that already |
| * has a rank greater than the node child node we inserted or |
| * updated as the tree is now valid. |
| */ |
| return; |
| } |
| DEBUG_ASSERT(node->rank + 1 == rank); |
| if (bst_node_rank(node->child[!is_right_child]) + 2 < rank) { |
| /* |
| * Rank of @node cannot be incremented. This means it can be moved |
| * down and demoted instead. |
| * |
| * The tree can be rotated as pictured below. (a is @node which |
| * could not be promoted. Numbers are known relative ranks.) |
| * |
| * If rank of c is 2 less than the rank of a (b is inserted or |
| * promoted node): |
| * |
| * a2 b2 |
| * / \ / \ |
| * b2 D0 => A a1 |
| * / \ / \ |
| * A c0 c0 D0 |
| * / \ / \ |
| * B C B C |
| * |
| * If rank of c is 1 less than the rank of a (b is promoted node, c |
| * is inserted or promoted node): |
| * a2 a2 __c2__ |
| * / \ / \ / \ |
| * b2 D0 c2 D0 b1 a1 |
| * / \ => / \ => / \ / \ |
| * A0 c1 b1 C A0 B C D0 |
| * / \ / \ |
| * B C A0 B |
| */ |
| bst_rotate_insert(root, node->child[is_right_child], node, |
| is_right_child); |
| return; |
| } |
| node->rank = rank; |
| } |
| } |
| |
| /** |
| * bst_rotate_delete - Internal helper function |
| * @root: Tree. |
| * @up1: Node to move up if a single rotate is enough. |
| * @down: Node to move down. |
| * @up_was_right_child: %true if @up1 was the right child of @down. |
| * |
| * Rotate sub-tree (once or twice) after delete and update ranks. |
| */ |
| static void bst_rotate_delete(struct bst_root *root, struct bst_node *up1, |
| struct bst_node *down, bool up_was_right_child) { |
| DEBUG_ASSERT(down->child[up_was_right_child] == up1); |
| DEBUG_ASSERT(up1->rank == down->rank - 1); |
| DEBUG_ASSERT(down->rank == |
| bst_node_rank(down->child[!up_was_right_child]) + 3); |
| struct bst_node *up2 = up1->child[!up_was_right_child]; |
| if (bst_node_rank(up1->child[up_was_right_child]) <= down->rank - 3) { |
| DEBUG_ASSERT(bst_node_rank(up2) == down->rank - 2); |
| /* |
| * Swap nodes @up2 and @up1 then @up2 and @down |
| * (pictured for up_was_right_child==false): |
| * |
| * down(0) down up2(0) |
| * / \ / \ / \ |
| * up1(-1) D(-3) up2 D up1(-2) down(-2) |
| * / \ / \ / \ / \ |
| * A(-3) up2(-2) up1 C A(-3) B C D(-3) |
| * / \ / \ |
| * B C A(-3) B |
| */ |
| DEBUG_ASSERT(up1->rank == down->rank - 1); |
| DEBUG_ASSERT(up2->rank == down->rank - 2); |
| bst_rotate(root, up2, up1, !up_was_right_child); |
| bst_rotate(root, up2, down, up_was_right_child); |
| up2->rank += 2; |
| up1->rank--; |
| down->rank -= 2; |
| } else { |
| /* |
| * Swap nodes @up1 and @down (pictured for up_was_right_child==false): |
| * |
| * down(0) up1(0) |
| * / \ / \ |
| * up1(-1) C(-3) A(-2) down(-1) |
| * / \ / \ |
| * A(-2) B(-2/-3) B(-2/-3) C(-3) |
| */ |
| bst_rotate(root, up1, down, up_was_right_child); |
| up1->rank++; |
| down->rank--; |
| if (bst_node_rank(down->child[0]) == down->rank - 2 && |
| bst_node_rank(down->child[1]) == down->rank - 2) { |
| /* Demote down if possible. (Required if down is a leaf node) */ |
| down->rank--; |
| } |
| } |
| } |
| |
| /** |
| * bst_update_rank_delete - Internal helper function |
| * @root: Tree. |
| * @node: Node to start scan at. This is the parent of the node that |
| * was removed from the tree. Note that the removed node will |
| * be a different node than the node passed to bst_delete if |
| * that node had two children. |
| * @is_right_child: %true if the right child of @node was deleted. |
| * |
| * Demote nodes and/or rotate sub-trees to make @root a valid WAVL tree again. |
| */ |
| static void bst_update_rank_delete(struct bst_root *root, struct bst_node *node, |
| bool is_right_child) { |
| DEBUG_ASSERT(root); |
| DEBUG_ASSERT(node); |
| |
| DEBUG_ASSERT(bst_node_rank(node->child[is_right_child]) <= |
| bst_node_rank(node->child[!is_right_child])); |
| |
| while (node) { |
| DEBUG_ASSERT(node->rank > bst_node_rank(node->child[!is_right_child])); |
| DEBUG_ASSERT(node->rank - 1 > |
| bst_node_rank(node->child[is_right_child])); |
| DEBUG_ASSERT(node->rank <= |
| bst_node_rank(node->child[!is_right_child]) + 2); |
| |
| /* |
| * At this point the rank of @node->child[is_right_child] has been |
| * decremented. We may need to also decrement the rank of @node. |
| */ |
| if (!node->child[0] && !node->child[1]) { |
| /* Always demote leaf node (from 2 to 1) */ |
| |
| /* We should not be in this function if the rank is alrady 1 */ |
| DEBUG_ASSERT(node->rank == 2); |
| } else if (node->rank <= |
| bst_node_rank(node->child[is_right_child]) + 2) { |
| /* |
| * If rank of @node does not need to change then we now have a valid |
| * tree. |
| */ |
| return; |
| } |
| DEBUG_ASSERT(node->rank > 1); |
| node->rank--; |
| if (node->rank <= bst_node_rank(node->child[!is_right_child])) { |
| /* We demoted @node, but it is now invalid on the other side */ |
| DEBUG_ASSERT(node->rank == |
| bst_node_rank(node->child[!is_right_child])); |
| if (bst_node_rank(node->child[!is_right_child]->child[0]) == |
| node->rank - 2 && |
| bst_node_rank(node->child[!is_right_child]->child[1]) == |
| node->rank - 2) { |
| /* If the other child can be demoted, demote it and move up */ |
| node->child[!is_right_child]->rank--; |
| } else { |
| /* |
| * If the other child can not be demoted, rotate instead. This |
| * will produce a valid tree without changing the rank of the |
| * node linked at the current spot @node in the tree. |
| * |
| * Undo demotion as current bst_rotate_delete implemention |
| * assumes node rank is unchanged. |
| */ |
| node->rank++; |
| |
| bst_rotate_delete(root, node->child[!is_right_child], node, |
| !is_right_child); |
| return; |
| } |
| } |
| is_right_child = bst_is_right_child(node); |
| node = node->parent; |
| } |
| } |
| |
| /** |
| * bst_find_edge - Internal helper function |
| * @node: Node to start search at. |
| * @edge: Direction if search. |
| * |
| * Return: leftmost (if @edge is %false) or rightmost (if @edge is %true) node |
| * in subtree with @node as root. |
| */ |
| static struct bst_node *bst_find_edge(struct bst_node *node, bool edge) { |
| struct bst_node *saved_node; |
| |
| DEBUG_ASSERT(node); |
| |
| do { |
| saved_node = node; |
| node = node->child[edge]; |
| } while (node); |
| |
| return saved_node; |
| } |
| |
| /** |
| * bst_delete_all_helper - Internal helper function |
| * @root: Tree. |
| * @node: Node to delete (most be the leftmost node in @root). |
| * |
| * Helper function to delete leftmost node in @root, assuming all other nodes |
| * will be deleted next. |
| */ |
| void bst_delete_all_helper(struct bst_root *root, struct bst_node *node) { |
| DEBUG_ASSERT(root); |
| DEBUG_ASSERT(node); |
| DEBUG_ASSERT(!node->child[0]); |
| bst_move_node(root, node, node->child[1]); |
| } |
| |
| void bst_delete(struct bst_root *root, struct bst_node *node) { |
| DEBUG_ASSERT(root); |
| DEBUG_ASSERT(node); |
| |
| struct bst_node *new_child; |
| bool node_is_right_child = bst_is_right_child(node); |
| struct bst_node *update_rank_start = node->parent; |
| bool update_rank_is_right_child = node_is_right_child; |
| if (!node->child[0]) { |
| /* |
| * If @node has no left child, link its right child in its place. (The |
| * right child could be %NULL in this case) |
| */ |
| new_child = node->child[1]; |
| } else if (!node->child[1]) { |
| /* |
| * If @node has no right child, link its left child in its place. |
| */ |
| DEBUG_ASSERT(node->child[0]); |
| new_child = node->child[0]; |
| } else { |
| /* |
| * If @node has both left and right children, delete (from the tree |
| * structure point of view) the left-most node in the right sub-tree or |
| * the right-most node in the left sub-tree instead. Either side would |
| * work. |
| */ |
| struct bst_node *edge_node = bst_find_edge( |
| node->child[!node_is_right_child], node_is_right_child); |
| struct bst_node *edge_child = edge_node->child[!node_is_right_child]; |
| update_rank_start = edge_node->parent; |
| update_rank_is_right_child = bst_is_right_child(edge_node); |
| if (update_rank_start == node) { |
| update_rank_start = edge_node; |
| update_rank_is_right_child = !node_is_right_child; |
| } |
| DEBUG_ASSERT(update_rank_start); |
| bst_move_node(root, edge_node, edge_child); |
| |
| new_child = edge_node; |
| DEBUG_ASSERT(new_child); |
| bst_link_node(new_child, 0, node->child[0]); |
| bst_link_node(new_child, 1, node->child[1]); |
| new_child->rank = node->rank; |
| } |
| bst_move_node(root, node, new_child); |
| node->rank = 0; |
| if (update_rank_start) { |
| bst_update_rank_delete(root, update_rank_start, update_rank_is_right_child); |
| } |
| } |
| |
| /** |
| * bst_prev_next - Internal helper function |
| * @root: Tree. |
| * @node: Node to move from. |
| * @dir_next: Directon to move. |
| * |
| * Return: If @node is %NULL and @dir_next is %false, right-most node in @root. |
| * If @node is %NULL and @dir_next is %true, left-most node in @root. |
| * If @node is not %NULL and @dir_next is %false, right-most node to the |
| * left of @node. |
| * If @node is not %NULL and @dir_next is %true, left-most node to the |
| * right of @node. |
| * %NULL if the node described above does not exist. |
| */ |
| static struct bst_node *bst_prev_next(const struct bst_root *root, |
| struct bst_node *node, |
| bool dir_next) { |
| DEBUG_ASSERT(root); |
| |
| struct bst_node *next_child = node ? node->child[dir_next] : root->root; |
| |
| if (!node && !next_child) { |
| return NULL; /* Empty tree */ |
| } |
| |
| /* |
| * Comments below assume @dir_next is %true. For the @dir_next is %false |
| * case, swap left and right. |
| */ |
| if (next_child) { |
| /* There is a right child, return the left-most node in that subtree */ |
| return bst_find_edge(next_child, !dir_next); |
| } else { |
| /* No right child, next node is the first right parent */ |
| struct bst_node *next_parent = node; |
| while (bst_is_right_child(next_parent) == dir_next) { |
| next_parent = next_parent->parent; |
| if (!next_parent) { |
| return NULL; |
| } |
| } |
| return next_parent->parent; |
| } |
| } |
| |
| struct bst_node *bst_prev(struct bst_root *root, struct bst_node *node) { |
| return bst_prev_next(root, node, false); |
| } |
| |
| struct bst_node *bst_next(const struct bst_root *root, struct bst_node *node) { |
| return bst_prev_next(root, node, true); |
| } |
