[libc-commits] [libc] [libc][tsearch] add weak AVL tree for tsearch implementation (PR #172411)

[Michael Jones via libc-commits]

================
@@ -0,0 +1,583 @@
+//===-- 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
+//
+//===----------------------------------------------------------------------===//
+
+#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 {
+
+// A general self-balancing binary search tree where the node pointer can
+// be used as stable handles to the stored values.
+//
+// The self-balancing strategy is the Weak AVL (WAVL) tree, based on the
+// following foundational references:
+// 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
+// 4. https://dl.acm.org/doi/10.1145/2689412 (Rank-Balanced Trees)
+//
+// WAVL trees belong to the rank-balanced binary search tree framework
+// (reference 4), alongside AVL and Red-Black trees.
+//
+// Key Properties of WAVL Trees:
+// 1. Relationship to Red-Black Trees: A WAVL tree can always be colored as a
+//    Red-Black tree.
+// 2. Relationship to AVL Trees: An AVL tree meets all the requirements of a
+//    WAVL tree. Insertion-only WAVL trees maintain the same structure as AVL
+//    trees.
+//
+// Rank-Based Balancing:
+// In rank-balanced trees, each node is assigned a rank (conceptually similar
+// to height). The rank difference between a parent and its child is
+// strictly enforced to be either **1** or **2**.
+//
+// - **AVL Trees:** Rank is equivalent to height. The strict condition is that
+//   there are no 2-2 nodes (a parent with rank difference 2 to both children).
+// - **WAVL Trees:** The no 2-2 node rule is relaxed for internal nodes during
+//   the deletion fixup process, making WAVL trees less strictly balanced than
+//   AVL trees but easier to maintain than Red-Black trees.
+//
+// Balancing Mechanics (Promotion/Demotion):
+// - **Null nodes** are considered to have rank -1.
+// - **External/leaf nodes** have rank 0.
+// - **Insertion:** Inserting a node may create a situation where a parent and
+//   child have the same rank (difference 0). This is fixed by **promoting**
+//   the rank of the parent and propagating the fix upwards using at most two
+//   rotations (trinode fixup).
+// - **Deletion:** Deleting a node may result in a parent being 3 ranks higher
+//   than a child (difference 3). This is fixed by **demoting** the parent's
+//   rank and propagating the fix upwards.
+//
+// Implementation Detail:
+// The rank is **implicitly** maintained. We never store the full rank. Instead,
+// a 2-bit tag is used on each node to record the rank difference to each child:
+// - Bit cleared (0) -> Rank difference is **1**.
+// - Bit set (1)     -> Rank difference is **2**.
+template <typename T> class WeakAVLNode {
+  // Data
+  T data;
+
+  // Parent pointer
+  WeakAVLNode *parent;
+
+  // Children pointers
+  WeakAVLNode *children[2];
+
+  // Flags
+  unsigned char left_rank_diff_2 : 1;
+  unsigned char right_rank_diff_2 : 1;
+
+  LIBC_INLINE bool is_leaf() {
+    return (children[0] == nullptr) && (children[1] == nullptr);
+  }
+
+  LIBC_INLINE void toggle_rank_diff_2(bool is_right) {
+    if (is_right)
+      right_rank_diff_2 ^= 1;
+    else
+      left_rank_diff_2 ^= 1;
+  }
+
+  LIBC_INLINE bool both_flags_set() const {
+    return left_rank_diff_2 && right_rank_diff_2;
+  }
+
+  LIBC_INLINE bool any_flag_set() const {
+    return left_rank_diff_2 || right_rank_diff_2;
+  }
+
+  LIBC_INLINE void clear_flags() {
+    left_rank_diff_2 = 0;
+    right_rank_diff_2 = 0;
+  }
+
+  LIBC_INLINE void set_both_flags() {
+    left_rank_diff_2 = 1;
+    right_rank_diff_2 = 1;
+  }
+
+  LIBC_INLINE WeakAVLNode(T data)
+      : data(cpp::move(data)), parent(nullptr), children{nullptr, nullptr},
+        left_rank_diff_2(0), right_rank_diff_2(0) {}
+
+  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 updated. 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.
+  // Our tree requires value to stay in their node to maintain stable addresses.
+  // This complicates the unlink operation as the successor transplanting needs
+  // to update all the pointers and flags.
+  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->parent = nullptr;
+        return {nullptr, false};
+      }
+      FixupSite site = {node->parent, node->parent->children[1] == node};
+      site.parent->children[site.is_right] = child;
+      child->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;
+    // succ and node may be adjacent to each other, so we
+    // still need to check the exact direction of the successor.
+    bool succ_was_right = succ_parent->children[1] == succ;
+    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->parent = succ_parent;
+
+    // 2) Transplant successor into node's position
+    succ->parent = node->parent;
+    succ->left_rank_diff_2 = node->left_rank_diff_2;
+    succ->right_rank_diff_2 = node->right_rank_diff_2;
+
+    succ->children[0] = node->children[0];
+    succ->children[1] = node->children[1];
+    if (succ->children[0])
+      succ->children[0]->parent = succ;
+    if (succ->children[1])
+      succ->children[1]->parent = succ;
+
+    if (succ->parent) {
+      bool node_was_right = succ->parent->children[1] == node;
+      succ->parent->children[node_was_right] = succ;
+    } else
+      root = succ;
----------------
michaelrj-google wrote:

given that the `if` uses braces the `else` should for symmetry

https://github.com/llvm/llvm-project/pull/172411