util: Add sparse bitset data structure

Useful for potentially huge bitsets that are expected to be mostly
filled with zeroes, reducing memory consumption by assuming bits being
zero by default (without wasting memory to store zeroes).

Co-authored-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37908>

4 files changed, 515 insertions, 0 deletions

diff --git a/src/.clang-format b/src/.clang-format
index 099120f2cca..22ba63a5b63 100644
--- a/src/.clang-format
+++ b/src/.clang-format
@@ -103,6 +103,8 @@ ForEachMacros:
   - foreach_list_typed_safe
   - foreach_two_lists
 
+  - U_SPARSE_BITSET_FOREACH_SET
+
 # nir
   - nir_foreach_function_temp_variable
   - nir_foreach_function_temp_variable_safe
diff --git a/src/util/meson.build b/src/util/meson.build
index 0152a53126e..14b3af1176d 100644
--- a/src/util/meson.build
+++ b/src/util/meson.build
@@ -447,6 +447,7 @@ if with_tests
     'tests/register_allocate_test.cpp',
     'tests/roundeven_test.cpp',
     'tests/set_test.cpp',
+    'tests/sparse_bitset_test.cpp',
     'tests/string_buffer_test.cpp',
     'tests/timespec_test.cpp',
     'tests/u_atomic_test.cpp',
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
diff --git a/src/util/tests/sparse_bitset_test.cpp b/src/util/tests/sparse_bitset_test.cpp
new file mode 100644
index 00000000000..00de3d0bc65
--- /dev/null
+++ b/src/util/tests/sparse_bitset_test.cpp
@@ -0,0 +1,127 @@
+/*
+ * Copyright © 2025 Valve Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include <gtest/gtest.h>
+#include "util/sparse_bitset.h"
+
+
+TEST(sparse_bitset, tree)
+{
+   struct u_sparse_bitset set;
+   u_sparse_bitset_init(&set, 1048577, NULL);
+
+   u_sparse_bitset_set(&set, 65535);
+   u_sparse_bitset_set(&set, 1048576);
+
+   unsigned num_nodes = 0;
+   rb_tree_foreach(struct u_sparse_bitset_node, _, &set.tree, node) {
+      ++num_nodes;
+   }
+   EXPECT_EQ(num_nodes, 2);
+
+   u_sparse_bitset_free(&set);
+}
+
+TEST(sparse_bitset, set_clear)
+{
+   struct u_sparse_bitset set;
+   u_sparse_bitset_init(&set, 1048577, NULL);
+
+   u_sparse_bitset_set(&set, 65535);
+   u_sparse_bitset_set(&set, 1048576);
+   u_sparse_bitset_set(&set, 16383);
+
+   EXPECT_EQ(u_sparse_bitset_test(&set, 128), false);
+   EXPECT_EQ(u_sparse_bitset_test(&set, 65535), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set, 16383), true);
+
+   u_sparse_bitset_clear(&set, 1236749);
+   u_sparse_bitset_clear(&set, 65535);
+
+   EXPECT_EQ(u_sparse_bitset_test(&set, 65535), false);
+
+   u_sparse_bitset_free(&set);
+}
+
+TEST(sparse_bitset, set_dup)
+{
+   struct u_sparse_bitset set;
+   u_sparse_bitset_init(&set, 1048577, NULL);
+
+   u_sparse_bitset_set(&set, 65535);
+   u_sparse_bitset_set(&set, 1048576);
+
+   struct u_sparse_bitset set2;
+   u_sparse_bitset_dup(&set2, &set);
+
+   u_sparse_bitset_clear(&set, 65535);
+
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 128), false);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 65535), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 1048576), true);
+
+   u_sparse_bitset_free(&set);
+   u_sparse_bitset_free(&set2);
+}
+
+TEST(sparse_bitset, set_merge)
+{
+   struct u_sparse_bitset set;
+   u_sparse_bitset_init(&set, 1048577, NULL);
+
+   u_sparse_bitset_set(&set, 65535);
+   u_sparse_bitset_set(&set, 1048576);
+
+   struct u_sparse_bitset set2;
+   u_sparse_bitset_init(&set2, 1048577, NULL);
+   u_sparse_bitset_set(&set2, 128);
+   u_sparse_bitset_set(&set2, 16383);
+
+   EXPECT_EQ(u_sparse_bitset_merge(&set2, &set), true);
+
+   u_sparse_bitset_clear(&set, 65535);
+
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 128), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 16383), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 65535), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 1048576), true);
+
+   EXPECT_EQ(u_sparse_bitset_merge(&set2, &set), false);
+
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 128), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 16383), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 65535), true);
+   EXPECT_EQ(u_sparse_bitset_test(&set2, 1048576), true);
+
+   u_sparse_bitset_free(&set);
+   u_sparse_bitset_free(&set2);
+}
+
+TEST(sparse_bitset, set_foreach)
+{
+   struct u_sparse_bitset set;
+   u_sparse_bitset_init(&set, 1048577, NULL);
+
+   u_sparse_bitset_set(&set, 65535);
+   u_sparse_bitset_set(&set, 1048576);
+   u_sparse_bitset_set(&set, 16383);
+   u_sparse_bitset_set(&set, 19);
+   u_sparse_bitset_set(&set, 422);
+   u_sparse_bitset_set(&set, 65539);
+
+   uint32_t arr[] = {19, 422, 16383, 65535, 65539, 1048576};
+   unsigned i = 0;
+
+   U_SPARSE_BITSET_FOREACH_SET(&set, it) {
+      EXPECT_LE(i, ARRAY_SIZE(arr));
+      EXPECT_EQ(it, arr[i]);
+      ++i;
+   }
+
+   EXPECT_EQ(i, ARRAY_SIZE(arr));
+
+   u_sparse_bitset_free(&set);
+}