Add bit_set library

Necessary for upcoming `epoch_scheduler` improvements.

Signed-off-by: Klimenty Tsoutsman <[email protected]>

libs/bit_set/Cargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "bit_set"
+version = "0.1.0"
+authors = ["Klim Tsoutsman <[email protected]>"]
+description = "A bit set storing integers less than 64"
+edition = "2021"

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<Self::Item> {
+        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::<Vec<_>>(), [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::<Vec<_>>(),
+            [0, 8, 16, 24, 32, 40, 48, 56]
+        );
+    }
+}

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<u8>` 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<u8> {
+        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<u8> {
+        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));
+    }
+}