bstree: Add balancing code

Add self balancing code. Balance tree after insert and delete according
to WAVL rules.

Bug: 141329802
Change-Id: Ic8d43b3b4b173c7ecdb0c48a2c58de0a28d501c4

3 files changed, 740 insertions, 0 deletions

diff --git a/lib/binary_search_tree/binary_search_tree.c b/lib/binary_search_tree/binary_search_tree.c
index 144200b9..9edeac4c 100644
--- a/lib/binary_search_tree/binary_search_tree.c
+++ b/lib/binary_search_tree/binary_search_tree.c
@@ -24,6 +24,16 @@
 #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.
  *
@@ -97,6 +107,298 @@ static void bst_move_node(struct bst_root *root,
 }
 
 /**
+ * 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.
@@ -138,6 +440,8 @@ void bst_delete(struct bst_root *root, struct bst_node *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
@@ -160,15 +464,26 @@ void bst_delete(struct bst_root *root, struct bst_node *node) {
         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);
+    }
 }
 
 /**
diff --git a/lib/binary_search_tree/hosttest/binary_search_tree_test.cpp b/lib/binary_search_tree/hosttest/binary_search_tree_test.cpp
index c30a90ca..4e68a2e7 100644
--- a/lib/binary_search_tree/hosttest/binary_search_tree_test.cpp
+++ b/lib/binary_search_tree/hosttest/binary_search_tree_test.cpp
@@ -112,9 +112,14 @@ static void bst_test_check_node(struct bst_root *root, struct bst_node *node) {
     if (node->parent) {
         ASSERT_NE(root->root, node) << node << root;
         ASSERT_EQ(node->parent->child[bst_is_right_child(node)], node) << node << root;
+        ASSERT_GE(node->parent->rank, node->rank + 1) << node << root;
+        ASSERT_LE(node->parent->rank, node->rank + 2) << node << root;
     } else {
         ASSERT_EQ(root->root, node) << node << root;;
     }
+    if (!node->child[0] && !node->child[1]) {
+        ASSERT_EQ(node->rank, 1U) << node << root;;
+    }
 }
 
 static void print_rep(char ch, size_t count) {
@@ -181,6 +186,89 @@ static void bst_test_check_tree_valid(struct bst_root *root) {
     }
 }
 
+static void bst_test_check_sub_tree(struct bst_node *subtree,
+                                    struct bst_node nodes[],
+                                    struct bst_node *parent,
+                                    const char **expected_tree,
+                                    int is_right_child);
+
+static void bst_test_check_child(struct bst_node *subtree,
+                                 struct bst_node nodes[],
+                                 const char **expected_tree,
+                                 int child) {
+    if (BstTest::GetParam()) {
+        child = !child;
+    }
+    while (isspace(**expected_tree)) {
+        (*expected_tree)++;
+    }
+
+    if (**expected_tree == '(') {
+        (*expected_tree)++;
+        bst_test_check_sub_tree(subtree->child[child], nodes, subtree,
+                                expected_tree, child);
+        if (::testing::Test::HasFatalFailure()) {
+            return;
+        }
+        ASSERT_EQ(**expected_tree, ')');
+        (*expected_tree)++;
+    } else {
+        EXPECT_EQ(subtree->child[child], nullptr);
+    }
+
+    while (isspace(**expected_tree)) {
+        (*expected_tree)++;
+    }
+}
+
+static void bst_test_check_sub_tree(struct bst_node *subtree,
+                                    struct bst_node nodes[],
+                                    struct bst_node *parent,
+                                    const char **expected_tree,
+                                    int is_right_child) {
+    size_t index;
+    size_t rank;
+
+    if (!parent && **expected_tree == '\0') {
+        return;
+    }
+    ASSERT_NE(subtree, nullptr);
+
+    bst_test_check_child(subtree, nodes, expected_tree, 0);
+    if (::testing::Test::HasFatalFailure()) {
+        return;
+    }
+
+    index = strtoul(*expected_tree, (char **)expected_tree, 0);
+    ASSERT_EQ(**expected_tree, 'r');
+    (*expected_tree)++;
+    rank = strtoul(*expected_tree, (char **)expected_tree, 0);
+
+    EXPECT_EQ(subtree, &nodes[index]);
+    EXPECT_EQ(subtree->rank, rank);
+    EXPECT_EQ(subtree->parent, parent);
+    EXPECT_EQ(bst_is_right_child(subtree), is_right_child);
+
+    bst_test_check_child(subtree, nodes, expected_tree, 1);
+    if (::testing::Test::HasFatalFailure()) {
+        return;
+    }
+}
+
+/*
+ * Check if tree has expected structure and ranks. The expected tree is
+ * described by a string, "(left child) <node index>r<rank> (right child)". For
+ * example: "((0r1) 1r2 (2r1)) 3r3 ((4r1) 5r2)".
+ */
+static void _bst_test_check_tree(struct bst_root *root, struct bst_node nodes[],
+                                 const char *expected_tree) {
+    bst_test_check_sub_tree(root->root, nodes, NULL, &expected_tree, 0);
+    EXPECT_EQ(*expected_tree, '\0');
+    if(::testing::Test::HasFailure()) {
+        bst_test_print_tree(root, nodes);
+    }
+}
+
 static void bst_test_check_array(struct bst_root *root,
                                      struct bst_node nodes[],
                                      size_t index[],
@@ -216,6 +304,12 @@ static void bst_test_check_array(struct bst_root *root,
     if (HasFatalFailure()) return; \
 } while(0)
 
+#define bst_test_check_tree(root, node, treestr) do { \
+    SCOPED_TRACE("bst_test_check_tree"); \
+    _bst_test_check_tree(root, nodes, treestr); \
+    if (HasFatalFailure()) return; \
+} while(0)
+
 static void bst_test_insert_func(struct bst_root *root, struct bst_node *node) {
     ASSERT_EQ(node->rank, 0U);
     bst_insert(root, node, bst_test_compare);
@@ -239,6 +333,16 @@ static void bst_test_insert_func(struct bst_root *root, struct bst_node *node) {
     bst_test_check(root, nodes, items); \
 } while(0)
 
+#define bst_test_insert_check_tree(root, nodes, insert, treestr) do { \
+    bst_test_insert(root, &nodes[insert]); \
+    bst_test_check_tree(root, nodes, treestr); \
+} while(0)
+
+#define bst_test_delete_check_tree(root, nodes, insert, treestr) do { \
+    bst_test_delete(root, &nodes[insert]); \
+    bst_test_check_tree(root, nodes, treestr); \
+} while(0)
+
 static void bst_test_delete_func(struct bst_root *root, struct bst_node *node) {
     bst_delete(root, node);
     bst_test_check_tree_valid(root);
@@ -297,6 +401,9 @@ TEST_P(BstTest, InsertAscending) {
     }
     bst_test_check_array(&root, nodes, NULL, countof(nodes), NULL, NULL);
     EXPECT_GE(bst_depth(&root), 4U); /* Minimum depth for a binary tree */
+    EXPECT_LT(bst_depth(&root), 15U); /* We should have a tree, not a list */
+    EXPECT_LE(bst_depth(&root), 8U); /* RB tree should have depth <= 8 */
+    EXPECT_LE(bst_depth(&root), 6U); /* WAVL tree should have depth <= 6 */
 }
 
 TEST_P(BstTest, InsertBalanced) {
@@ -492,6 +599,320 @@ TEST_P(BstTest, ForEveryEntryDelete) {
     EXPECT_EQ(bst_next(&root, NULL), nullptr);
 }
 
+/*
+ * WAVL specific tests.
+ * Name describe non-mirrored test (/0).
+ * Swap left/right for mirrrored test run (/1).
+ */
+TEST_P(BstTest, WAVLPromoteRootInsertLeft) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 1] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 1, "1r1");
+    bst_test_insert_check_tree(&root, nodes, 0, "(0r1) 1r2");
+}
+
+TEST_P(BstTest, WAVLPromote2InsertLeft) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 3] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 2, "2r1");
+    bst_test_insert_check_tree(&root, nodes, 1, "(1r1) 2r2");
+    bst_test_insert_check_tree(&root, nodes, 3, "(1r1) 2r2 (3r1)");
+    bst_test_insert_check_tree(&root, nodes, 0, "((0r1) 1r2) 2r3 (3r1)");
+}
+
+TEST_P(BstTest, WAVLRootRotateRightNoChildNodesInsertLeft) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 2] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 2, "2r1");
+    bst_test_insert_check_tree(&root, nodes, 1, "(1r1) 2r2");
+
+    /*
+     *       2r2*         1r2
+     *      /            /   \
+     *    1r2      =>  0r1   2r1
+     *   /
+     * 0r1
+     */
+    bst_test_insert_check_tree(&root, nodes, 0, "(0r1) 1r2 (2r1)");
+}
+
+TEST_P(BstTest, WAVLRootRotateRightNoChildNodesInsertRight) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 2] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 2, "2r1");
+    bst_test_insert_check_tree(&root, nodes, 0, "(0r1) 2r2");
+
+    /*
+     *       ___2r2*         1r2
+     *      /               /   \
+     *    0r2         =>  0r1   2r1
+     *       \
+     *       1r1
+     */
+    bst_test_insert_check_tree(&root, nodes, 1, "(0r1) 1r2 (2r1)");
+}
+
+
+TEST_P(BstTest, WAVLRootRotateRightWithChildNodesInsertLeft) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 5] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 4, "4r1");
+    bst_test_insert_check_tree(&root, nodes, 2, "(2r1) 4r2");
+    bst_test_insert_check_tree(&root, nodes, 5, "(2r1) 4r2 (5r1)");
+    bst_test_insert_check_tree(&root, nodes, 1, "((1r1) 2r2) 4r3 (5r1)");
+    bst_test_insert_check_tree(&root, nodes, 3, "((1r1) 2r2 (3r1)) 4r3 (5r1)");
+
+    /*
+     *          ___4r3*              2r3___
+     *         /      \             /      \
+     *       2r3      5r1  =>     1r2      4r2
+     *      /   \                /        /   \
+     *    1r2   3r1            0r1      3r1   5r1
+     *   /
+     * 0r1
+     */
+    bst_test_insert_check_tree(&root, nodes, 0, "((0r1) 1r2) 2r3 ((3r1) 4r2 (5r1))");
+}
+
+TEST_P(BstTest, WAVLNonRootRotateRightWithChildNodesInsertLeft) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 8] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 6, "6r1");
+    bst_test_insert_check_tree(&root, nodes, 4, "(4r1) 6r2");
+    bst_test_insert_check_tree(&root, nodes, 7, "(4r1) 6r2 (7r1)");
+    bst_test_insert_check_tree(&root, nodes, 2, "((2r1) 4r2) 6r3 (7r1)");
+    bst_test_insert_check_tree(&root, nodes, 8, "((2r1) 4r2) 6r3 (7r2 (8r1))");
+    bst_test_insert_check_tree(&root, nodes, 5, "((2r1) 4r2 (5r1)) 6r3 (7r2 (8r1))");
+    bst_test_insert_check_tree(&root, nodes, 1, "(((1r1) 2r2) 4r3 (5r1)) 6r4 (7r2 (8r1))");
+    bst_test_insert_check_tree(&root, nodes, 3, "(((1r1) 2r2 (3r1)) 4r3 (5r1)) 6r4 (7r2 (8r1))");
+
+    /*
+     *                ___6r4...             _________6r4...
+     *               /                     /
+     *          ___4r3*                  2r3___
+     *         /      \                 /      \
+     *       2r3      5r1      =>     1r2      4r1
+     *      /   \                    /        /   \
+     *    1r2   3r1                0r1      3r1   5r1
+     *   /
+     * 0r1
+     */
+    bst_test_insert_check_tree(&root, nodes, 0, "(((0r1) 1r2) 2r3 ((3r1) 4r2 (5r1))) 6r4 (7r2 (8r1))");
+}
+
+TEST_P(BstTest, WAVLRotateLeftRightSimpleInsertRight) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 5] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 4, "4r1");
+    bst_test_insert_check_tree(&root, nodes, 1, "(1r1) 4r2");
+    bst_test_insert_check_tree(&root, nodes, 5, "(1r1) 4r2 (5r1)");
+    bst_test_insert_check_tree(&root, nodes, 0, "((0r1) 1r2) 4r3 (5r1)");
+    bst_test_insert_check_tree(&root, nodes, 2, "((0r1) 1r2 (2r1)) 4r3 (5r1)");
+
+    /*
+     *       ______4r3                 ___4r3               2r3___
+     *      /         \               /      \             /      \
+     *    1r3         5r1           2r3      5r1         1r2      4r2
+     *   /   \             =>      /   \           =>   /        /   \
+     * 0r1   2r2                 1r2   3r1            0r1      3r1   5r1
+     *          \               /
+     *          3r1           0r1
+     */
+    bst_test_insert_check_tree(&root, nodes, 3, "((0r1) 1r2) 2r3 ((3r1) 4r2 (5r1))");
+}
+
+TEST_P(BstTest, WAVLRotateRightLeftSimpleInsertLeft) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 5] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 1, "1r1");
+    bst_test_insert_check_tree(&root, nodes, 0, "(0r1) 1r2");
+    bst_test_insert_check_tree(&root, nodes, 4, "(0r1) 1r2 (4r1)");
+    bst_test_insert_check_tree(&root, nodes, 3, "(0r1) 1r3 ((3r1) 4r2)");
+    bst_test_insert_check_tree(&root, nodes, 5, "(0r1) 1r3 ((3r1) 4r2 (5r1))");
+
+    /*
+     *    1r3______               1r3___                     ___3r3
+     *   /         \             /      \                   /      \
+     * 0r1         4r3         0r1      3r3               1r2      4r2
+     *            /   \    =>          /   \       =>    /   \        \
+     *          3r2   5r1            2r1   4r2         0r1   2r1      5r1
+     *         /                              \
+     *       2r1                              5r1
+     */
+    bst_test_insert_check_tree(&root, nodes, 2, "((0r1) 1r2 (2r1)) 3r3 (4r2 (5r1))");
+}
+
+TEST_P(BstTest, WAVLDemoteLeaf) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 1] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 1, "1r1");
+    bst_test_insert_check_tree(&root, nodes, 0, "(0r1) 1r2");
+    /*
+     *    1r2      1r2*      1r1
+     *   /     =>        =>
+     * 0r1
+     */
+    bst_test_delete_check_tree(&root, nodes, 0, "1r1");
+}
+
+TEST_P(BstTest, WAVLDemoteNonLeaf) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 3] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 1, "1r1");
+    bst_test_insert_check_tree(&root, nodes, 0, "(0r1) 1r2");
+    bst_test_insert_check_tree(&root, nodes, 2, "(0r1) 1r2 (2r1)");
+    bst_test_insert_check_tree(&root, nodes, 3, "(0r1) 1r3 (2r2 (3r1))");
+    bst_test_delete_check_tree(&root, nodes, 3, "(0r1) 1r3 (2r1)");
+    /*
+     *    1r3         1r3         1r2
+     *   /   \    =>     \    =>     \
+     * 0r1   2r1         2r1         2r1
+     */
+    bst_test_delete_check_tree(&root, nodes, 0, "1r2 (2r1)");
+}
+
+TEST_P(BstTest, WAVLDoubleDemote) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 6] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 2, "2r1");
+    bst_test_insert_check_tree(&root, nodes, 1, "(1r1) 2r2");
+    bst_test_insert_check_tree(&root, nodes, 4, "(1r1) 2r2 (4r1)");
+    bst_test_insert_check_tree(&root, nodes, 0, "((0r1) 1r2) 2r3 (4r1)");
+    bst_test_insert_check_tree(&root, nodes, 3, "((0r1) 1r2) 2r3 ((3r1) 4r2)");
+    bst_test_insert_check_tree(&root, nodes, 5, "((0r1) 1r2) 2r3 ((3r1) 4r2 (5r1))");
+    bst_test_insert_check_tree(&root, nodes, 6, "((0r1) 1r2) 2r4 ((3r1) 4r3 (5r2 (6r1)))");
+    bst_test_delete_check_tree(&root, nodes, 6, "((0r1) 1r2) 2r4 ((3r1) 4r3 (5r1))");
+    /*
+     *       2r4___            2r4___               2r3___
+     *      /      \    =>    /      \       =>    /      \
+     *    1r2      4r3      1r1      4r3         1r1      4r2
+     *   /        /   \             /   \                /   \
+     * 0r1      3r1   5r1         3r1   5r1            3r1   5r1
+     */
+    bst_test_delete_check_tree(&root, nodes, 0, "(1r1) 2r3 ((3r1) 4r2 (5r1))");
+}
+
+TEST_P(BstTest, WAVLRotateNoChildrenAfterDelete) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 3] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 1, "1r1");
+    bst_test_insert_check_tree(&root, nodes, 0, "(0r1) 1r2");
+    bst_test_insert_check_tree(&root, nodes, 2, "(0r1) 1r2 (2r1)");
+    bst_test_insert_check_tree(&root, nodes, 3, "(0r1) 1r3 (2r2 (3r1))");
+    /*
+     *    1r3            1r3               2r3
+     *   /   \              \             /   \
+     * 0r1   2r2     =>     2r2     =>  1r1   3r1
+     *          \              \
+     *          3r1            3r1
+     *
+     * (2r2 would also be a valid wavl tree after this rotate, but that does
+     * not work in the general case where 2 could have started with a left
+     * child)
+     */
+    bst_test_delete_check_tree(&root, nodes, 0, "(1r1) 2r3 (3r1)");
+}
+
+TEST_P(BstTest, WAVLRotateWithChildren1AfterDelete) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 6] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 2, "2r1");
+    bst_test_insert_check_tree(&root, nodes, 1, "(1r1) 2r2");
+    bst_test_insert_check_tree(&root, nodes, 4, "(1r1) 2r2 (4r1)");
+    bst_test_insert_check_tree(&root, nodes, 0, "((0r1) 1r2) 2r3 (4r1)");
+    bst_test_insert_check_tree(&root, nodes, 3, "((0r1) 1r2) 2r3 ((3r1) 4r2)");
+    bst_test_insert_check_tree(&root, nodes, 5, "((0r1) 1r2) 2r3 ((3r1) 4r2 (5r1))");
+    bst_test_insert_check_tree(&root, nodes, 6, "((0r1) 1r2) 2r4 ((3r1) 4r3 (5r2 (6r1)))");
+    /*
+     *       2r4___                 2r4___                     ___4r4
+     *      /      \               /      \                   /      \
+     *    1r2      4r3           1r1      4r3               2r2      5r2
+     *   /        /   \      =>          /   \       =>    /   \        \
+     * 0r1      3r1   5r2              3r1   5r2         1r1   3r1      6r1
+     *                   \                      \
+     *                   6r1                    6r1
+     *
+     * (4r3/2r2 would also be a valid wavl tree after this rotate, but that does
+     * not work in the general case where 3r1 could be 3r2. 2r3 would also be
+     * valid, but since we have to demote it when it is a leaf node, the current
+     * implementation demotes it when possible)
+     */
+    bst_test_delete_check_tree(&root, nodes, 0, "((1r1) 2r2 (3r1)) 4r4 (5r2 (6r1))");
+}
+
+TEST_P(BstTest, WAVLRotateWithChildren2AfterDelete) {
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 7] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 2, "2r1");
+    bst_test_insert_check_tree(&root, nodes, 1, "(1r1) 2r2");
+    bst_test_insert_check_tree(&root, nodes, 5, "(1r1) 2r2 (5r1)");
+    bst_test_insert_check_tree(&root, nodes, 0, "((0r1) 1r2) 2r3 (5r1)");
+    bst_test_insert_check_tree(&root, nodes, 4, "((0r1) 1r2) 2r3 ((4r1) 5r2)");
+    bst_test_insert_check_tree(&root, nodes, 6, "((0r1) 1r2) 2r3 ((4r1) 5r2 (6r1))");
+    bst_test_insert_check_tree(&root, nodes, 3, "((0r1) 1r2) 2r4 (((3r1) 4r2) 5r3 (6r1))");
+    bst_test_insert_check_tree(&root, nodes, 7, "((0r1) 1r2) 2r4 (((3r1) 4r2) 5r3 (6r2 (7r1)))");
+    /*
+     *       2r4_____               2r4______                   ______5r4
+     *      /        \             /         \                 /         \
+     *    1r2         5r3         1r1         5r3             2r3___      6r2
+     *   /           /   \     =>            /   \      =>   /      \        \
+     * 0r1         4r2   6r2               4r2   6r2       1r1      4r2      7r1
+     *            /         \             /         \              /
+     *          3r1         7r1         3r1         7r1          3r1
+     */
+    bst_test_delete_check_tree(&root, nodes, 0, "((1r1) 2r3 ((3r1) 4r2)) 5r4 (6r2 (7r1))");
+}
+
+TEST_P(BstTest, WAVLDoubleRotateWithChildren2AfterDelete) {
+        /*
+         * 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
+         */
+    struct bst_root root = BST_ROOT_INITIAL_VALUE;
+    struct bst_node nodes[] = {[0 ... 7] = BST_NODE_INITIAL_VALUE};
+
+    bst_test_insert_check_tree(&root, nodes, 2, "2r1");
+    bst_test_insert_check_tree(&root, nodes, 1, "(1r1) 2r2");
+    bst_test_insert_check_tree(&root, nodes, 6, "(1r1) 2r2 (6r1)");
+    bst_test_insert_check_tree(&root, nodes, 0, "((0r1) 1r2) 2r3 (6r1)");
+    bst_test_insert_check_tree(&root, nodes, 4, "((0r1) 1r2) 2r3 ((4r1) 6r2)");
+    bst_test_insert_check_tree(&root, nodes, 7, "((0r1) 1r2) 2r3 ((4r1) 6r2 (7r1))");
+    bst_test_insert_check_tree(&root, nodes, 3, "((0r1) 1r2) 2r4 (((3r1) 4r2) 6r3 (7r1))");
+    bst_test_insert_check_tree(&root, nodes, 5, "((0r1) 1r2) 2r4 (((3r1) 4r2 (5r1)) 6r3 (7r1))");
+    /*
+     *    2r4_________             2r4___                     ___4r4___
+     *   /            \           /      \                   /         \
+     * 1r1         ___6r3       1r1      4r3___            2r2         6r2
+     *            /      \   =>         /      \     =>   /   \       /   \
+     *          4r2      7r1          3r1      6r2      1r1   3r1   5r1   7r1
+     *         /   \                          /   \
+     *       3r1   5r1                      5r1   7r1
+     */
+    bst_test_delete_check_tree(&root, nodes, 0, "((1r1) 2r2 (3r1)) 4r4 ((5r1) 6r2 (7r1))");
+}
+
 TEST_P(BstTest, RandomInsert) {
     struct bst_root root = BST_ROOT_INITIAL_VALUE;
     struct bst_node nodes[] = {[0 ... 999] = BST_NODE_INITIAL_VALUE};
diff --git a/lib/binary_search_tree/include/lib/binary_search_tree.h b/lib/binary_search_tree/include/lib/binary_search_tree.h
index e6bfd931..fc85e849 100644
--- a/lib/binary_search_tree/include/lib/binary_search_tree.h
+++ b/lib/binary_search_tree/include/lib/binary_search_tree.h
@@ -124,6 +124,9 @@ static inline struct bst_node *bst_search(const struct bst_root *root,
     containerof_null_safe(bst_search(root, &(item)->member, compare), type, \
                           member)
 
+/* Internal helper. Don't call directly */
+void bst_update_rank_insert(struct bst_root *root, struct bst_node *node);
+
 /**
  * bst_insert - Insert node in tree.
  * @root:       Tree.
@@ -154,6 +157,7 @@ static inline bool bst_insert(struct bst_root *root, struct bst_node *node,
             node->child[0] = NULL;
             node->child[1] = NULL;
             *parent_ptr = node;
+            bst_update_rank_insert(root, node);
             return true;
         }
         diff = compare(tree_node, node);