compiler/rust/bitset: Add a BitIndex helper struct

Reviewed-by: Mel Henning <mhenning@darkrefraction.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41915>
This commit is contained in:
Faith Ekstrand 2026-05-21 13:35:37 -04:00 committed by Marge Bot
parent fa867cf189
commit 81c9eddb69

View file

@ -23,8 +23,8 @@
use std::cmp::{max, min};
use std::marker::PhantomData;
use std::ops::{
BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, RangeFull,
Sub, SubAssign,
Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor,
BitXorAssign, RangeFull, Sub, SubAssign,
};
/// Converts a value into a bit index
@ -64,6 +64,54 @@ impl FromBitIndex for usize {
}
}
#[derive(Clone, Copy)]
struct BitIndex {
word: usize,
bit: u8,
}
impl BitIndex {
const fn flatten(self) -> usize {
self.word * 32 + (self.bit as usize)
}
const fn from_flat_index(idx: usize) -> BitIndex {
BitIndex {
word: idx / 32,
bit: (idx % 32) as u8,
}
}
}
impl<K: IntoBitIndex> From<K> for BitIndex {
fn from(key: K) -> BitIndex {
BitIndex::from_flat_index(key.into_bit_index())
}
}
impl From<BitIndex> for usize {
fn from(idx: BitIndex) -> usize {
idx.flatten()
}
}
impl AddAssign<usize> for BitIndex {
fn add_assign(&mut self, rhs: usize) {
let bit = usize::from(self.bit) + rhs;
self.bit = (bit % 32) as u8;
self.word += bit / 32;
}
}
impl Add<usize> for BitIndex {
type Output = BitIndex;
fn add(mut self, rhs: usize) -> BitIndex {
self += rhs;
self
}
}
/// A set implemented as an array of bits
///
/// Unlike `HashSet` and similar containers which actually store the provided
@ -115,33 +163,27 @@ impl<K> BitSet<K> {
impl<K: IntoBitIndex> BitSet<K> {
pub fn contains(&self, key: K) -> bool {
let idx = key.into_bit_index();
let w = idx / 32;
let b = idx % 32;
if w < self.words.len() {
self.words[w] & (1_u32 << b) != 0
let idx = BitIndex::from(key);
if idx.word < self.words.len() {
self.words[idx.word] & (1_u32 << idx.bit) != 0
} else {
false
}
}
pub fn insert(&mut self, key: K) -> bool {
let idx = key.into_bit_index();
let w = idx / 32;
let b = idx % 32;
self.reserve_words(w + 1);
let exists = self.words[w] & (1_u32 << b) != 0;
self.words[w] |= 1_u32 << b;
let idx = BitIndex::from(key);
self.reserve_words(idx.word + 1);
let exists = self.words[idx.word] & (1_u32 << idx.bit) != 0;
self.words[idx.word] |= 1_u32 << idx.bit;
!exists
}
pub fn remove(&mut self, key: K) -> bool {
let idx = key.into_bit_index();
let w = idx / 32;
let b = idx % 32;
self.reserve_words(w + 1);
let exists = self.words[w] & (1_u32 << b) != 0;
self.words[w] &= !(1_u32 << b);
let idx = BitIndex::from(key);
self.reserve_words(idx.word + 1);
let exists = self.words[idx.word] & (1_u32 << idx.bit) != 0;
self.words[idx.word] &= !(1_u32 << idx.bit);
exists
}
}