diff --git a/src/compiler/rust/bitset.rs b/src/compiler/rust/bitset.rs index eb706fdc6c4..d27383bd0a2 100644 --- a/src/compiler/rust/bitset.rs +++ b/src/compiler/rust/bitset.rs @@ -24,7 +24,7 @@ use std::cmp::{max, min}; use std::marker::PhantomData; use std::ops::{ Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, - BitXorAssign, RangeFull, Sub, SubAssign, + BitXorAssign, Range, RangeFull, Sub, SubAssign, }; /// Converts a value into a bit index @@ -179,6 +179,228 @@ fn find_next_unset(words: &[u32], start: BitIndex) -> BitIndex { BitIndex::from_word(words.len()) } +/// A set implemented as an array of bits, able to be used as constant data +/// +/// The fixed size W is in units of 32-bit words. This is due to a Rust +/// restriction which prevents us from doing math on constants which size +/// arrays. +#[derive(Clone, Copy)] +pub struct ConstBitSet { + words: [u32; W], + phantom: PhantomData, +} + +impl ConstBitSet { + pub const fn new() -> Self { + ConstBitSet { + words: [0_u32; W], + phantom: PhantomData, + } + } + + pub const fn clear(&mut self) { + let mut w = 0_usize; + while w < W { + self.words[w] = 0; + w += 1; + } + } + + pub const fn is_empty(&self) -> bool { + let mut w = 0_usize; + while w < W { + if self.words[w] != 0 { + return false; + } + w += 1; + } + true + } +} + +macro_rules! impl_const_bit_set_binop { + ( + $K:ident, + $BinOp:ident, + $bin_op:ident, + $AssignBinOp:ident, + $assign_bin_op:ident, + |$a:ident, $b:ident| $impl:expr, + ) => { + impl $AssignBinOp<&ConstBitSet> + for ConstBitSet + { + fn $assign_bin_op(&mut self, rhs: &ConstBitSet) { + for w in 0..W { + let $a = self.words[w]; + let $b = rhs.words[w]; + self.words[w] = $impl; + } + } + } + + impl $AssignBinOp> + for ConstBitSet + { + fn $assign_bin_op(&mut self, rhs: ConstBitSet) { + self.$assign_bin_op(&rhs); + } + } + + impl $BinOp<&ConstBitSet> + for ConstBitSet + { + type Output = ConstBitSet; + + fn $bin_op( + mut self, + rhs: &ConstBitSet, + ) -> ConstBitSet { + self.$assign_bin_op(rhs); + self + } + } + + impl $BinOp> for ConstBitSet { + type Output = ConstBitSet; + + fn $bin_op( + mut self, + rhs: ConstBitSet, + ) -> ConstBitSet { + self.$assign_bin_op(rhs); + self + } + } + }; +} + +macro_rules! impl_const_bit_set { + ($K:ident) => { + impl ConstBitSet { + pub const fn contains(&self, key: $K) -> bool { + let idx = BitIndex::from_flat_index(key as usize); + if idx.word < self.words.len() { + self.words[idx.word] & (1_u32 << idx.bit) != 0 + } else { + false + } + } + + pub const fn insert(&mut self, key: $K) -> bool { + let idx = BitIndex::from_flat_index(key as usize); + assert!(idx.word < W, "ConstBitSet index out of bounds"); + let exists = self.contains(key); + self.words[idx.word] |= 1_u32 << idx.bit; + !exists + } + + pub const fn insert_range(&mut self, range: Range<$K>) { + assert!( + range.end as usize <= W * 32, + "ConstBitSet index out of bounds", + ); + + if range.start >= range.end { + return; + } + + let start = BitIndex::from_flat_index(range.start as usize); + let end = BitIndex::from_flat_index(range.end as usize); + + let mut word = start.word; + let mut mask = u32::MAX << start.bit; + while word < end.word { + self.words[word] |= mask; + mask = u32::MAX; + word += 1; + } + + debug_assert!(word == end.word); + mask &= !(u32::MAX << end.bit); + self.words[word] |= mask; + } + + pub const fn remove(&mut self, key: $K) -> bool { + let idx = BitIndex::from_flat_index(key as usize); + if idx.word < self.words.len() { + let exists = self.contains(key); + self.words[idx.word] &= !(1_u32 << idx.bit); + exists + } else { + false + } + } + + pub const fn from_array(arr: [$K; N]) -> Self { + let mut set = ConstBitSet::::new(); + let mut i = 0_usize; + while i < N { + set.insert(arr[i]); + i += 1; + } + set + } + + pub const fn from_range(range: Range<$K>) -> Self { + let mut set = ConstBitSet::::new(); + set.insert_range(range); + set + } + + pub fn iter(&self) -> impl '_ + Iterator { + BitSetIter::new(&self.words) + } + } + + impl From> for ConstBitSet { + fn from(range: Range<$K>) -> ConstBitSet { + ConstBitSet::::from_range(range) + } + } + + impl_const_bit_set_binop!( + $K, + BitAnd, + bitand, + BitAndAssign, + bitand_assign, + |a, b| a & b, + ); + + impl_const_bit_set_binop!( + $K, + BitOr, + bitor, + BitOrAssign, + bitor_assign, + |a, b| a | b, + ); + + impl_const_bit_set_binop!( + $K, + BitXor, + bitxor, + BitXorAssign, + bitxor_assign, + |a, b| a ^ b, + ); + + impl_const_bit_set_binop!( + $K, + Sub, + sub, + SubAssign, + sub_assign, + |a, b| a & !b, + ); + }; +} + +impl_const_bit_set!(u8); +impl_const_bit_set!(u16); +impl_const_bit_set!(usize); + /// A set implemented as an array of bits /// /// Unlike `HashSet` and similar containers which actually store the provided