diff --git a/libs/bit_set/Cargo.toml b/libs/bit_set/Cargo.toml new file mode 100644 index 0000000000..af52627a32 --- /dev/null +++ b/libs/bit_set/Cargo.toml @@ -0,0 +1,6 @@ +[package] +name = "bit_set" +version = "0.1.0" +authors = ["Klim Tsoutsman "] +description = "A bit set storing integers less than 64" +edition = "2021" diff --git a/libs/bit_set/src/iter.rs b/libs/bit_set/src/iter.rs new file mode 100644 index 0000000000..135f36c3c2 --- /dev/null +++ b/libs/bit_set/src/iter.rs @@ -0,0 +1,67 @@ +use core::intrinsics::unlikely; + +/// An iterator over a [`BitSet`]. +/// +/// [`BitSet`]: crate::BitSet +pub struct Iter { + set: u64, + current_mask: u64, +} + +impl Iter { + pub(crate) const fn new(set: u64) -> Self { + Self { + set, + current_mask: u64::MAX, + } + } +} + +impl Iterator for Iter { + type Item = usize; + + fn next(&mut self) -> Option { + let next_index = (self.set & self.current_mask).trailing_zeros(); + + if unlikely(next_index == 64) { + None + } else { + // https://users.rust-lang.org/t/how-to-make-an-integer-with-n-bits-set-without-overflow/63078 + self.current_mask = u64::MAX.checked_shl(next_index + 1).unwrap_or(0); + Some(next_index as usize) + } + } +} + +#[cfg(test)] +mod tests { + extern crate alloc; + + use alloc::vec::Vec; + + use crate::BitSet; + + #[test] + fn test_iter() { + let mut set = BitSet::new(); + set.insert(57); + set.insert(58); + set.insert(61); + set.insert(63); + assert_eq!(set.iter().collect::>(), [57, 58, 61, 63]); + + let mut set = BitSet::new(); + set.insert(0); + set.insert(8); + set.insert(16); + set.insert(24); + set.insert(32); + set.insert(40); + set.insert(48); + set.insert(56); + assert_eq!( + set.iter().collect::>(), + [0, 8, 16, 24, 32, 40, 48, 56] + ); + } +} diff --git a/libs/bit_set/src/lib.rs b/libs/bit_set/src/lib.rs new file mode 100644 index 0000000000..1724f46fc7 --- /dev/null +++ b/libs/bit_set/src/lib.rs @@ -0,0 +1,165 @@ +//! A bit set backed by a [`u64`]. +//! +//! See [`BitSet`] for more details. + +#![no_std] +#![feature(const_likely, core_intrinsics)] + +mod iter; + +use core::intrinsics::likely; + +pub use iter::Iter; + +/// A bit set backed by a [`u64`]. +/// +/// This is equivalent to a `HashSet` storing integers in the range `[0, +/// 64)`. +#[derive(Debug, Clone)] +pub struct BitSet { + inner: u64, +} + +impl BitSet { + /// Constructs a new, empty `BitSet`. + pub const fn new() -> Self { + Self { inner: 0 } + } + + /// Returns an iterator over the elements of the set. + pub const fn iter(&self) -> Iter { + Iter::new(self.inner) + } + + /// Returns `true` if the set contains the given element. + /// + /// # Panics + /// + /// Panics if `element` is greater than 63. + #[must_use] + pub const fn contains(&self, element: u8) -> bool { + assert!(element < 64); + self.inner & (1 << element) != 0 + } + + /// Adds an element to the set. + /// + /// # Panics + /// + /// Panics if `element` is greater than 63. + pub fn insert(&mut self, element: u8) { + assert!(element < 64); + self.inner |= 1 << element; + } + + /// Removes an element from the set. + /// + /// # Panics + /// + /// Panics if `element` is greater than 63. + pub fn remove(&mut self, element: u8) { + assert!(element < 64); + self.inner &= !(1 << element); + } + + /// Returns the smallest element in the set. + /// + /// Returns `None` if the set is empty. + #[must_use] + pub const fn min(&self) -> Option { + if likely(self.inner != 0) { + Some(self.inner.trailing_zeros() as u8) + } else { + None + } + } + + /// Returns the largest element in the set. + /// + /// Returns `None` if the set is empty. + #[must_use] + pub const fn max(&self) -> Option { + if likely(self.inner != 0) { + // self.inner.leading_zeros() <= 63 + Some(63 - self.inner.leading_zeros() as u8) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_contains() { + let mut set = BitSet::new(); + + for i in 0..64 { + assert!(!set.contains(i)); + } + + set.insert(3); + + for i in 0..64 { + if i != 3 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + + set.insert(0); + + for i in 0..64 { + if i != 0 && i != 3 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + + set.insert(63); + + for i in 0..64 { + if i != 0 && i != 3 && i != 63 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + } + + #[test] + fn test_remove() { + let mut set = BitSet::new(); + + set.insert(3); + set.insert(63); + set.remove(3); + + for i in 0..64 { + if i != 63 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + } + + #[test] + fn test_min_max() { + let mut set = BitSet::new(); + assert_eq!(set.min(), None); + assert_eq!(set.max(), None); + + set.insert(5); + assert_eq!(set.min(), Some(5)); + assert_eq!(set.max(), Some(5)); + + set.insert(3); + assert_eq!(set.min(), Some(3)); + assert_eq!(set.max(), Some(5)); + } +}