Weak AVL

wavl.cpp

//===-- Implementation header for weak AVL tree -----------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
// Weak AVL tree implementation based on the algorithm described in:
// 1. https://maskray.me/blog/2025-12-14-weak-avl-tree
// 2. https://reviews.freebsd.org/D25480
// 3. https://ics.uci.edu/~goodrich/teach/cs165/notes/WeakAVLTrees.pdf
//===----------------------------------------------------------------------===//

#ifndef LLVM_LIBC_SRC___SUPPORT_WEAK_AVL_H
#define LLVM_LIBC_SRC___SUPPORT_WEAK_AVL_H

#include "hdr/stdint_proxy.h"
#include "src/__support/CPP/bit.h"
#include "src/__support/CPP/new.h"
#include "src/__support/CPP/utility/move.h"
#include "src/__support/libc_assert.h"
#include "src/__support/macros/attributes.h"
#include "src/__support/macros/config.h"

namespace LIBC_NAMESPACE_DECL {

template <typename T> class WeakAVLNode {
  // Data
  T data;

  // Packs the parent pointer with 2 flag bits in the low bits. Bit 0 indicates
  // whether the left child has rank difference 2; bit 1 indicates whether the
  // right child has rank difference 2. A cleared bit means rank difference 1.
  // All rank differences are 1 or 2, and every leaf has rank 0.
  uintptr_t parent_and_flags;

  // Children pointers
  WeakAVLNode *children[2];

  // Constants
  static LIBC_INLINE_VAR constexpr uintptr_t FLAGS_MASK = 0b11;
  static LIBC_INLINE_VAR constexpr uintptr_t LEFT_FLAG_BIT = 0b01;
  static LIBC_INLINE_VAR constexpr uintptr_t RIGHT_FLAG_BIT = 0b10;

  // Auxiliary methods for accessing fields
  LIBC_INLINE WeakAVLNode *parent() const {
    return cpp::bit_cast<WeakAVLNode *>(parent_and_flags & ~FLAGS_MASK);
  }
  LIBC_INLINE uintptr_t flags() const { return parent_and_flags & FLAGS_MASK; }
  LIBC_INLINE void set_parent(WeakAVLNode *p) {
    parent_and_flags = cpp::bit_cast<uintptr_t>(p) | flags();
  }
  LIBC_INLINE bool has_rank_diff_2(bool is_right) const {
    return flags() & (is_right ? RIGHT_FLAG_BIT : LEFT_FLAG_BIT);
  }
  LIBC_INLINE void toggle_rank_diff_2(bool is_right) {
    parent_and_flags ^= (is_right ? RIGHT_FLAG_BIT : LEFT_FLAG_BIT);
  }
  LIBC_INLINE void clear_flags() { parent_and_flags &= ~FLAGS_MASK; }
  LIBC_INLINE void set_flags(uintptr_t flags) {
    parent_and_flags = (parent_and_flags & ~FLAGS_MASK) | (flags & FLAGS_MASK);
  }
  LIBC_INLINE bool is_leaf() {
    return (children[0] == nullptr) && (children[1] == nullptr);
  }
  LIBC_INLINE WeakAVLNode(T data)
      : data(cpp::move(data)), parent_and_flags(0), children{nullptr, nullptr} {
  }

  LIBC_INLINE static WeakAVLNode *create(T value) {
    AllocChecker ac;
    WeakAVLNode *res = ::new (ac) WeakAVLNode(value);
    if (ac)
      return res;
    return nullptr;
  }

  // Unlink a node from tree. The corresponding flag is not update. The node is
  // not deleted and its pointers are not cleared.
  // FixupSite is the lowest surviving node from which rank/flag invariants may
  // be violated.
  struct FixupSite {
    WeakAVLNode *parent;
    bool is_right;
  };
  LIBC_INLINE static FixupSite unlink(WeakAVLNode *&root, WeakAVLNode *node) {
    bool has_left = node->children[0] != nullptr;
    bool has_right = node->children[1] != nullptr;

    // Case 0: no children
    if (!has_left && !has_right) {
      if (!node->parent()) {
        root = nullptr;
        return {nullptr, false};
      }
      FixupSite site = {node->parent(), node->parent()->children[1] == node};
      site.parent->children[site.is_right] = nullptr;
      return site;
    }

    // Case 1: one child
    if (has_left != has_right) {
      WeakAVLNode *child = node->children[has_right];
      if (!node->parent()) {
        root = child;
        child->set_parent(nullptr);
        return {nullptr, false};
      }
      FixupSite site = {node->parent(), node->parent()->children[1] == node};
      site.parent->children[site.is_right] = child;
      child->set_parent(site.parent);
      return site;
    }

    // Case 2: two children: replace by successor (leftmost in right subtree)
    WeakAVLNode *succ = node->children[1];
    while (succ->children[0])
      succ = succ->children[0];

    WeakAVLNode *succ_parent = succ->parent();
    bool succ_was_right =
        succ_parent->children[1] == succ; // true only if succ_parent==node
    WeakAVLNode *succ_rchild = succ->children[1];

    // 1) Splice successor out of its old position (flags intentionally
    // unchanged)
    FixupSite site = {succ_parent, succ_was_right};
    succ_parent->children[succ_was_right] = succ_rchild;
    if (succ_rchild)
      succ_rchild->set_parent(succ_parent);

    // 2) Transplant successor into node's position
    succ->set_parent(node->parent());
    succ->set_flags(
        node->flags()); // key-swap emulation: successor takes node's rank/flags

    succ->children[0] = node->children[0];
    succ->children[1] = node->children[1];
    if (succ->children[0])
      succ->children[0]->set_parent(succ);
    if (succ->children[1])
      succ->children[1]->set_parent(succ);

    if (succ->parent()) {
      bool node_was_right = succ->parent()->children[1] == node;
      succ->parent()->children[node_was_right] = succ;
    } else {
      root = succ;
    }

    // 3) If the physical removal was under `node`, fixup parent must be the
    //    successor (since `node` is deleted and successor now occupies that
    //    spot).
    if (site.parent == node)
      site.parent = succ;

    return site;
  }

public:
  LIBC_INLINE const WeakAVLNode *get_left() const { return children[0]; }
  LIBC_INLINE const WeakAVLNode *get_right() const { return children[1]; }
  LIBC_INLINE const T &get_data() const { return data; }
  LIBC_INLINE bool is_rank_diff_2(bool is_right) const {
    return has_rank_diff_2(is_right);
  }

  // Destroy the subtree rooted at node
  LIBC_INLINE static void destroy(WeakAVLNode *node) {
    if (!node)
      return;
    destroy(node->children[0]);
    destroy(node->children[1]);
    ::delete node;
  }
  // Rotate the subtree rooted at node in the given direction.
  //
  // Illustration for is_right = true (Left Rotation):
  //
  //          (Node)                       (Pivot)
  //          /    \                       /     \
  //         A   (Pivot)       =>       (Node)    C
  //             /     \                /    \
  //            B       C              A      B
  //
  LIBC_INLINE static WeakAVLNode *rotate(WeakAVLNode *&root, WeakAVLNode *node,
                                         bool is_right) {
    WeakAVLNode *pivot = node->children[is_right];
    // Handover pivot's child
    WeakAVLNode *grandchild = pivot->children[!is_right];
    node->children[is_right] = grandchild;
    if (grandchild)
      grandchild->set_parent(node);
    pivot->set_parent(node->parent());
    // Pivot becomes the new root of the subtree
    if (!node->parent())
      root = pivot;
    else {
      bool node_is_right = node->parent()->children[1] == node;
      node->parent()->children[node_is_right] = pivot;
    }
    pivot->children[!is_right] = node;
    node->set_parent(pivot);
    return pivot;
  }

  // Find data in the subtree rooted at root. If not found, returns nullptr.
  // `Compare` returns integer values for ternary comparison.
  template <typename Compare>
  LIBC_INLINE static WeakAVLNode *find(WeakAVLNode *root, T data,
                                       Compare &&comp) {
    WeakAVLNode *cursor = root;
    while (cursor != nullptr) {
      int comp_result = comp(cursor->data, data);
      if (comp_result == 0)
        return cursor; // Node found
      bool is_right = comp_result < 0;
      cursor = cursor->children[is_right];
    }
    return nullptr; // Node not found
  }
  // Insert data into the subtree rooted at root.
  // Returns the node if insertion is successful or the node exists in
  // the tree.
  // Returns nullptr if memory allocation fails.
  // `Compare` returns integer values for ternary comparison.
  template <typename Compare>
  LIBC_INLINE static WeakAVLNode *find_or_insert(WeakAVLNode *&root, T data,
                                                 Compare &&comp) {
    WeakAVLNode *parent = nullptr, *cursor = root;
    bool is_right = false;
    while (cursor != nullptr) {
      parent = cursor;
      int comp_result = comp(parent->data, data);
      if (comp_result == 0)
        return parent; // Node already exists
      is_right = comp_result < 0;
      cursor = cursor->children[is_right];
    }
    WeakAVLNode *allocated = create(cpp::move(data));
    if (!allocated)
      return nullptr;
    WeakAVLNode *node = allocated;
    node->set_parent(parent);

    // Case 0: inserting into an empty tree
    if (!parent) {
      root = node; // Tree was empty
      return node;
    }

    parent->children[is_right] = node;
    // Rebalance process
    while (parent) {
      is_right = (parent->children[1] == node);
      // Case 1: parent does not need to be promoted as node is lowering
      // than the parent by 2 ranks.
      //      (P)                       (P)
      //     /  \                      /  \
      //    2    1           =>       1    1
      //   /      \                  /      \
      // (N)       (*)             (N)       (*)
      if (parent->has_rank_diff_2(is_right)) {
        parent->toggle_rank_diff_2(is_right);
        break;
      }

      bool sibling_has_rank_diff_2 = parent->has_rank_diff_2(!is_right);
      // Case 2: node's sibling has rank-difference 1.
      // Promoting parent will fix the conflict of the trinodes but we may need
      // to continue on parent.
      //
      //         (GP)                       (GP)
      //          |         Promote          |   x - 1
      //          | x        ----->         (P)
      //      0   |         /           1  /   \
      // (N) --- (P)    ----             (N)    \ 2
      //            \  1                         \
      //             (S)                          (S)
      if (!sibling_has_rank_diff_2) {
        parent->toggle_rank_diff_2(!is_right);
        node = parent;
        parent = node->parent();
        continue;
      }

      LIBC_ASSERT((node->flags() != 0b11) &&
                  "there should be no 2-2 node along the insertion fixup path");

      LIBC_ASSERT((node == allocated || node->flags() != 0) &&
                  "Internal node must have a child with rank-difference 2, "
                  "otherwise it should have already been handled.");

      // Case 3: node's sibling has rank-difference 2. And node has a 1-node
      // along the same direction. We can do a single rotation to fix the
      // trinode.
      //                   (GP)                            (GP)
      //               0    |   X      Rotate               |
      //         (N) ----- (P)           =>                (N)
      //     1  /   \  2      \  2                      1  /  \ 1
      //      (C1)   \         \                        (C1)   (P)
      //             (C2)       (S)                         1 /  \ 1
      //                                                    (C2)  (S)
      if (node->has_rank_diff_2(!is_right)) {
        WeakAVLNode *new_subroot = rotate(root, parent, is_right);
        new_subroot->clear_flags();
        parent->clear_flags();
        break;
      }
      // Case 4: node's sibling has rank-difference 2. And node has a 1-node
      // along the opposite direction. We need a double rotation to fix the
      // trinode.
      //                   (GP)                            (GP)
      //               0    |   X      Zig-Zag              |      X
      //         (N) ----- (P)           =>                (C1)
      //     2  /   \  1      \  2                      1  /  \ 1
      //       /    (C1)       \                        (N)    (P)
      //     (C2) L /  \ R      (S)                  1 / \ L R / \ 1
      //          (A)  (B)                           (C2) (A)(B) (S)
      // (mirrored)
      //         (GP)                                      (GP)
      //        X | 0                Zig-Zag                |      X
      //         (P) ----- (N)           =>                (C1)
      //    2  /         1 / \ 2                        1  /  \ 1
      //      /         (C1)  \                         (P)    (N)
      //    (S)       L /  \ R (C2)                   1 / \ L R / \ 1
      //              (A)  (B)                        (S)(A)  (B)(C2)
      WeakAVLNode *subroot1 = rotate(root, node, !is_right); // First rotation
      [[maybe_unused]] WeakAVLNode *subroot2 =
          rotate(root, parent, is_right); // Second rotation
      LIBC_ASSERT(subroot1 == subroot2 &&
                  "Subroots after double rotation should be the same");
      uintptr_t flags = subroot1->flags();
      node->clear_flags();
      parent->clear_flags();
      subroot1->clear_flags();
      // Select destinations
      WeakAVLNode *dst_left = is_right ? parent : node;
      WeakAVLNode *dst_right = is_right ? node : parent;
      // Masked toggles
      if (flags & LEFT_FLAG_BIT)
        dst_left->toggle_rank_diff_2(true);

      if (flags & RIGHT_FLAG_BIT)
        dst_right->toggle_rank_diff_2(false);
      break;
    }
    return allocated;
  }

  // Erase the node from the tree rooted at root.
  LIBC_INLINE static void erase(WeakAVLNode *&root, WeakAVLNode *node) {
    // Unlink the node from the tree
    auto [cursor, is_right] = unlink(root, node);
    ::delete node;
    while (cursor) {
      // Case 0. cursor previously had rank-difference 1 on the side of the
      // deleted node. We can simply update the rank-difference and stop.
      // Notice that this step may create 2-2 nodes, thus deviate from "strong"
      // AVL tree.
      //
      //          (C)                 (C)
      //       X /   \ 1     =>    X /   \
      //       (*)   (D)           (*)    \ 2
      //                                   (D)
      if (!cursor->has_rank_diff_2(is_right)) {
        cursor->toggle_rank_diff_2(is_right);
        // If we created a 2-2 leaf, we must demote it and continue.
        if (cursor->flags() == 0b11 && cursor->is_leaf()) {
          cursor->clear_flags();
          if (cursor->parent())
            is_right = (cursor->parent()->children[1] == cursor);
          cursor = cursor->parent();
          continue;
        }
        break;
      }

      // Case 1. cursor previously had rank-difference 2 on the side of the
      // deleted node. Now it has rank-difference 3, which violates the
      // weak-AVL property. We found that we have a sibling with rank-difference
      // 2, so we can demote cursor and continue upwards.
      //
      //          (P)                 (P)
      //           |   X               |   (X + 1)
      //          (C)                  |
      //         /   \      =>        (C)
      //     2  /     \            1  / \
      //      (*)      \ 3         (*)   \ 2
      //               (D)                (D)
      if (cursor->has_rank_diff_2(!is_right)) {
        cursor->toggle_rank_diff_2(!is_right);
        if (cursor->parent())
          is_right = (cursor->parent()->children[1] == cursor);
        cursor = cursor->parent();
        continue;
      }

      // Case 2. continue from Case 1; but the sibling has rank-difference 1.
      // However, we found that the sibling is a 2-2 node. We demote both
      // sibling and cursor, and continue upwards.
      //
      //          (P)                 (P)
      //           |   X               |   (X + 1)
      //          (C)                  |
      //      1  /   \      =>        (C)
      //       (S)    \            1  / \
      //      /  \     \ 3         (S)   \ 2
      //   2 /    \ 2   (D)     1 /  \ 1  (D)
      //   (*)    (*)           (*)  (*)
      WeakAVLNode *sibling = cursor->children[!is_right];
      LIBC_ASSERT(sibling && "rank-difference 1 sibling cannot be empty");
      if (sibling->flags() == 0b11) {
        sibling->clear_flags();
        if (cursor->parent())
          is_right = (cursor->parent()->children[1] == cursor);
        cursor = cursor->parent();
        continue;
      }

      // Case 3. continue from Case 2; but the sibling cannot be demoted.
      // Sibling has a node T along the same direction with rank-difference 1.
      //
      //          (P)                             (P)
      //           | X                             | X
      //          (C)                             (S)
      //      1  /   \          Rotate         2 /   \ 1
      //       (S)    \           =>            /    (C)
      //    1  / \ Y   \ 3                    (T)   Y / \ 2
      //    (T)   \     (D)                        (*)   \
      //           (*)                                    (D)
      bool sibling_is_right = !is_right;
      if (!sibling->has_rank_diff_2(sibling_is_right)) {
        WeakAVLNode *new_subroot = rotate(root, cursor, sibling_is_right);
        LIBC_ASSERT(new_subroot == sibling &&
                    "sibling should become the subtree root");
        // Update flags
        bool sibling_alter_child_has_rank_diff_2 =
            new_subroot->has_rank_diff_2(!sibling_is_right);
        new_subroot->clear_flags();
        new_subroot->toggle_rank_diff_2(sibling_is_right);

        // Cursor only needs to be updated if it become a 2-2 node
        if (sibling_alter_child_has_rank_diff_2) {
          // Demote a 2-2 cursor if it is a leaf
          if (cursor->is_leaf()) {
            cursor->clear_flags();
            new_subroot->toggle_rank_diff_2(!sibling_is_right);
            LIBC_ASSERT(new_subroot->flags() == 0b11 &&
                        "sibling should become a 2-2 node.");
          } else {
            cursor->toggle_rank_diff_2(sibling_is_right);
            LIBC_ASSERT(cursor->flags() == 0b11 &&
                        "cursor should become a 2-2 node.");
          }
        }
        break;
      }
      // Case 4. continue from Case 3; but rank-difference 1 child T of sibling
      // is on the opposite direction.
      //
      //             (P)                                     (P)
      //              | X                                     | X
      //             (C)               Zig-Zag               (T)
      //          1 /   \                =>                  / \
      //          (S)    \                                2 /   \ 2
      //         /   \ 1  \ 3                             (S)   (C)
      //    2   /     (T)  (D)                         1 / Y \ / Z \ 1
      //      (*)   Y /  \ Z                            (*) (A)(B) (D)
      //            (A)  (B)
      WeakAVLNode *target_child = rotate(root, sibling, !sibling_is_right);
      uintptr_t flags = target_child->flags();
      WeakAVLNode *new_subroot = rotate(root, cursor, sibling_is_right);
      LIBC_ASSERT(new_subroot == target_child &&
                  "target_child should become the subtree root");
      // Set flags
      target_child->set_flags(0b11);
      cursor->clear_flags();
      sibling->clear_flags();
      // Select destinations
      WeakAVLNode *dst_left = sibling_is_right ? cursor : sibling;
      WeakAVLNode *dst_right = sibling_is_right ? sibling : cursor;
      // Masked toggles
      if (flags & LEFT_FLAG_BIT)
        dst_left->toggle_rank_diff_2(true);
      if (flags & RIGHT_FLAG_BIT)
        dst_right->toggle_rank_diff_2(false);
      break;
    }
  }
};

} // namespace LIBC_NAMESPACE_DECL

#endif // LLVM_LIBC_SRC___SUPPORT_WEAK_AVL_H