Files
mesa/src/util/sparse_bitset.h
Natalie Vock a8b75dd0f4 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>
2025-11-06 21:34:33 +00:00

386 lines
11 KiB
C

/*
* 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