1 files changed, 385 insertions, 0 deletions
diff --git a/src/util/sparse_bitset.h b/src/util/sparse_bitset.h
new file mode 100644
index 00000000000..61bae25be59
--- /dev/null
+++ b/src/util/sparse_bitset.h
@@ -0,0 +1,385 @@
+/*
+ * Copyright © 2025 Valve Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#ifndef _UTIL_SPARSE_BITSET_H
+#define _UTIL_SPARSE_BITSET_H
+
+#include "rb_tree.h"
+#include "bitset.h"
+#include "ralloc.h"
+
+#define U_SPARSE_BITSET_LOG2_BITS_PER_NODE 10
+#define U_SPARSE_BITSET_BITS_PER_NODE      (1u << U_SPARSE_BITSET_LOG2_BITS_PER_NODE)
+#define U_SPARSE_BITSET_BIT_INDEX_MASK     (U_SPARSE_BITSET_BITS_PER_NODE - 1)
+#define U_SPARSE_BITSET_OFFSET_MASK        (~U_SPARSE_BITSET_BIT_INDEX_MASK)
+
+/* Sets under this # of bits use small-set optimization */
+#define U_SPARSE_BITSET_SMALL_SET_THRESHOLD 0
+
+struct u_sparse_bitset_node {
+   struct rb_node node;
+
+   /* The first bit covered by this node */
+   unsigned offset;
+   BITSET_DECLARE(vals, U_SPARSE_BITSET_BITS_PER_NODE);
+};
+
+#define to_u_sparse_bitset_node(rb_node)                                       \
+   rb_node_data(struct u_sparse_bitset_node, rb_node, node)
+
+/*
+ * sparse_bitset wraps around a rb_tree of bitset nodes.
+ *
+ * Using sparse_bitset over a regular bitset is advantageous when you have a
+ * large number of potentially-set bits, but expect most of them to be zero
+ * (with the set bits mostly being within small, scattered regions).
+ *
+ * By default, bits are assumed to be unset. Areas that have set bits are
+ * represented by nodes in the rb_tree. One node represents a fixed-size bit
+ * range (internally with a non-sparse bitset of U_SPARSE_BITSET_BITS_PER_NODE).
+ */
+struct u_sparse_bitset {
+   union {
+      /* Large set - rb_tree of bitset nodes */
+      struct {
+         void *mem_ctx;
+         struct rb_tree tree;
+      };
+
+      /* Small set optimization - bitset on heap */
+      BITSET_WORD *vals;
+   };
+
+   /* Capacity of a small set, or 0 to indicate a large set */
+   unsigned capacity;
+};
+
+static inline bool
+_u_sparse_bitset_is_small(const struct u_sparse_bitset *s)
+{
+   return s->capacity != 0;
+}
+
+static inline void
+u_sparse_bitset_init(struct u_sparse_bitset *s, unsigned capacity,
+                     void *mem_ctx)
+{
+   if (capacity && capacity < U_SPARSE_BITSET_SMALL_SET_THRESHOLD) {
+      s->vals = BITSET_RZALLOC(mem_ctx, capacity);
+      s->capacity = capacity;
+   } else {
+      rb_tree_init(&s->tree);
+      s->mem_ctx = mem_ctx;
+      s->capacity = 0;
+   }
+}
+
+static inline int
+_u_sparse_bitset_node_comparator(const struct rb_node *a, unsigned offset_b)
+{
+   unsigned offset_a = to_u_sparse_bitset_node(a)->offset;
+   return (offset_a < offset_b) - (offset_a > offset_b);
+}
+
+static inline int
+_u_sparse_bitset_node_compare(const struct rb_node *a, const struct rb_node *b)
+{
+   unsigned offs_b = to_u_sparse_bitset_node(b)->offset;
+   return _u_sparse_bitset_node_comparator(a, offs_b);
+}
+
+static inline int
+_u_sparse_bitset_node_search(const struct rb_node *node, const void *key)
+{
+   return _u_sparse_bitset_node_comparator(node, (unsigned)(uintptr_t)key);
+}
+
+static inline struct u_sparse_bitset_node *
+_u_sparse_bitset_get_node(struct u_sparse_bitset *s, unsigned offset)
+{
+   assert(!_u_sparse_bitset_is_small(s));
+
+   struct rb_node *node = rb_tree_search(&s->tree, (void *)(uintptr_t)offset,
+                                         _u_sparse_bitset_node_search);
+   if (!node)
+      return NULL;
+
+   return rb_node_data(struct u_sparse_bitset_node, node, node);
+}
+
+static inline struct u_sparse_bitset_node *
+_u_sparse_bitset_get_or_add_node(struct u_sparse_bitset *s, unsigned offset)
+{
+   assert(!_u_sparse_bitset_is_small(s));
+   assert((offset & U_SPARSE_BITSET_BIT_INDEX_MASK) == 0);
+
+   struct u_sparse_bitset_node *node = _u_sparse_bitset_get_node(s, offset);
+   if (!node) {
+      node = rzalloc(s->mem_ctx, struct u_sparse_bitset_node);
+      node->offset = offset;
+      rb_tree_insert(&s->tree, &node->node, _u_sparse_bitset_node_compare);
+   }
+
+   return node;
+}
+
+static inline void
+u_sparse_bitset_set(struct u_sparse_bitset *s, unsigned bit)
+{
+   if (_u_sparse_bitset_is_small(s)) {
+      assert(bit < s->capacity);
+      BITSET_SET(s->vals, bit);
+   } else {
+      struct u_sparse_bitset_node *node =
+         _u_sparse_bitset_get_or_add_node(s, bit & U_SPARSE_BITSET_OFFSET_MASK);
+
+      BITSET_SET(node->vals, bit & U_SPARSE_BITSET_BIT_INDEX_MASK);
+   }
+}
+
+static inline void
+u_sparse_bitset_clear(struct u_sparse_bitset *s, unsigned bit)
+{
+   if (_u_sparse_bitset_is_small(s)) {
+      assert(bit < s->capacity);
+      BITSET_CLEAR(s->vals, bit);
+   } else {
+      struct u_sparse_bitset_node *node =
+         _u_sparse_bitset_get_node(s, bit & U_SPARSE_BITSET_OFFSET_MASK);
+
+      if (node)
+         BITSET_CLEAR(node->vals, bit & U_SPARSE_BITSET_BIT_INDEX_MASK);
+   }
+}
+
+static inline bool
+u_sparse_bitset_test(struct u_sparse_bitset *s, unsigned bit)
+{
+   if (_u_sparse_bitset_is_small(s)) {
+      assert(bit < s->capacity);
+      return BITSET_TEST(s->vals, bit);
+   } else {
+      struct u_sparse_bitset_node *node =
+         _u_sparse_bitset_get_node(s, bit & U_SPARSE_BITSET_OFFSET_MASK);
+
+      return node != NULL &&
+             BITSET_TEST(node->vals, bit & U_SPARSE_BITSET_BIT_INDEX_MASK);
+   }
+}
+
+static inline int
+u_sparse_bitset_cmp(struct u_sparse_bitset *a, struct u_sparse_bitset *b)
+{
+   assert(a->capacity == b->capacity);
+
+   if (_u_sparse_bitset_is_small(a)) {
+      return memcmp(a->vals, b->vals, BITSET_BYTES(a->capacity));
+   }
+
+   struct rb_node *a_iter = rb_tree_first(&a->tree);
+   struct rb_node *b_iter = rb_tree_first(&b->tree);
+
+   while (a_iter && b_iter) {
+      struct u_sparse_bitset_node *node_a = to_u_sparse_bitset_node(a_iter);
+      struct u_sparse_bitset_node *node_b = to_u_sparse_bitset_node(b_iter);
+
+      int node_cmp = _u_sparse_bitset_node_compare(a_iter, b_iter);
+      if (node_cmp)
+         return node_cmp;
+
+      int cmp_res = memcmp(node_a->vals, node_b->vals, sizeof(node_a->vals));
+      if (cmp_res)
+         return cmp_res;
+
+      a_iter = rb_node_next(a_iter);
+      b_iter = rb_node_next(b_iter);
+   }
+
+   return (a_iter != NULL) - (b_iter != NULL);
+}
+
+static inline void
+u_sparse_bitset_dup_with_ctx(struct u_sparse_bitset *dst,
+                             struct u_sparse_bitset *src, void *mem_ctx)
+{
+   u_sparse_bitset_init(dst, src->capacity, mem_ctx);
+
+   if (_u_sparse_bitset_is_small(src)) {
+      memcpy(dst->vals, src->vals, BITSET_BYTES(src->capacity));
+      return;
+   }
+
+   rb_tree_foreach(struct u_sparse_bitset_node, node, &src->tree, node) {
+      if (BITSET_IS_EMPTY(node->vals))
+         continue;
+
+      struct u_sparse_bitset_node *dst_node =
+         (struct u_sparse_bitset_node *)ralloc_memdup(
+            mem_ctx, node, sizeof(struct u_sparse_bitset_node));
+
+      rb_tree_insert(&dst->tree, &dst_node->node,
+                     _u_sparse_bitset_node_compare);
+   }
+}
+
+static inline void
+u_sparse_bitset_dup(struct u_sparse_bitset *dst, struct u_sparse_bitset *src)
+{
+   u_sparse_bitset_dup_with_ctx(dst, src, src->mem_ctx);
+}
+
+static inline bool
+_u_bitset_merge(BITSET_WORD *dst, const BITSET_WORD *src, unsigned words)
+{
+   bool changed = false;
+
+   for (unsigned i = 0; i < words; i++) {
+      changed |= (bool)(src[i] & ~dst[i]);
+      dst[i] |= src[i];
+   }
+
+   return changed;
+}
+
+static inline bool
+u_sparse_bitset_merge(struct u_sparse_bitset *dst, struct u_sparse_bitset *src)
+{
+   bool changed = false;
+   assert(dst->capacity == src->capacity);
+
+   if (_u_sparse_bitset_is_small(src)) {
+      return _u_bitset_merge(dst->vals, src->vals, BITSET_WORDS(src->capacity));
+   }
+
+   rb_tree_foreach(struct u_sparse_bitset_node, node, &src->tree, node) {
+      if (!BITSET_IS_EMPTY(node->vals)) {
+         struct u_sparse_bitset_node *dst_node =
+            _u_sparse_bitset_get_or_add_node(dst, node->offset);
+
+         changed |=
+            _u_bitset_merge(dst_node->vals, node->vals, ARRAY_SIZE(node->vals));
+      }
+   }
+
+   return changed;
+}
+
+static inline unsigned
+u_sparse_bitset_count(struct u_sparse_bitset *s)
+{
+   if (_u_sparse_bitset_is_small(s)) {
+      return __bitset_count(s->vals, s->capacity);
+   } else {
+      unsigned sum = 0;
+
+      rb_tree_foreach_safe(struct u_sparse_bitset_node, node, &s->tree, node) {
+         sum += BITSET_COUNT(node->vals);
+      }
+
+      return sum;
+   }
+}
+
+static inline void
+u_sparse_bitset_free(struct u_sparse_bitset *s)
+{
+   if (_u_sparse_bitset_is_small(s)) {
+      ralloc_free(s->vals);
+   } else {
+      rb_tree_foreach_safe(struct u_sparse_bitset_node, node, &s->tree, node) {
+         rb_tree_remove(&s->tree, &node->node);
+         ralloc_free(node);
+      }
+   }
+}
+
+static inline unsigned
+_u_sparse_bitset_next_set_dense(BITSET_WORD *set, unsigned size, unsigned from)
+{
+   /* Check if there even is a first node */
+   if (from >= size) {
+      return UINT_MAX;
+   }
+
+   unsigned i = BITSET_BITWORD(from);
+   unsigned offset = from % BITSET_WORDBITS;
+
+   /* Check for a next bit in the first node */
+   if (set[i] >> offset) {
+      return from + ffs(set[i] >> offset) - 1;
+   }
+
+   /* Else look for the next node */
+   for (i++; i < BITSET_WORDS(size); ++i) {
+      if (set[i]) {
+         return (i * BITSET_WORDBITS) + ffs(set[i]) - 1;
+      }
+   }
+
+   return UINT_MAX;
+}
+
+static inline unsigned
+_u_sparse_bitset_next_set(const struct u_sparse_bitset *s, uintptr_t *node_,
+                          const unsigned from)
+{
+   unsigned ret = UINT_MAX;
+
+   if (_u_sparse_bitset_is_small(s)) {
+      unsigned i = _u_sparse_bitset_next_set_dense(s->vals, s->capacity, from);
+      if (i < s->capacity) {
+         ret = i;
+      }
+   } else {
+      unsigned node_from = from;
+
+      /* Look for the next set bit in the current node */
+      for (struct rb_node *rb = (struct rb_node *)*node_; rb != NULL;
+           rb = rb_node_next(rb), node_from = 0) {
+
+         struct u_sparse_bitset_node *it = to_u_sparse_bitset_node(rb);
+         unsigned num = U_SPARSE_BITSET_BITS_PER_NODE;
+
+         /* Deal with from at the end of the node */
+         if (node_from != 0 &&
+             (node_from - it->offset) >= U_SPARSE_BITSET_BITS_PER_NODE) {
+            continue;
+         }
+
+         node_from &= U_SPARSE_BITSET_BIT_INDEX_MASK;
+
+         unsigned i = _u_sparse_bitset_next_set_dense(it->vals, num, node_from);
+
+         if (i < num) {
+            *node_ = (uintptr_t)rb;
+            ret = it->offset + i;
+            break;
+         }
+      }
+   }
+
+   assert(from <= ret);
+   return ret;
+}
+
+static inline unsigned
+_u_sparse_bitset_first_set(struct u_sparse_bitset *s, uintptr_t *node_)
+{
+   uintptr_t node = 0;
+   if (!_u_sparse_bitset_is_small(s)) {
+      node = (uintptr_t)rb_tree_first(&s->tree);
+   }
+
+   unsigned ret = _u_sparse_bitset_next_set(s, &node, 0);
+   *node_ = node;
+   return ret;
+}
+
+#define U_SPARSE_BITSET_FOREACH_SET(s, it)                                     \
+   for (uintptr_t _node, it = _u_sparse_bitset_first_set((s), &_node);         \
+        it != UINT_MAX; it = _u_sparse_bitset_next_set((s), &_node, it + 1))
+
+#endif //_UTIL_SPARSE_BITSET_H