Skip to content

Commit

Permalink
Add bit_set library
Browse files Browse the repository at this point in the history
Necessary for upcoming `epoch_scheduler` improvements.

Signed-off-by: Klimenty Tsoutsman <klim@tsoutsman.com>
  • Loading branch information
tsoutsman committed Dec 27, 2023
1 parent 813015b commit 50c045f
Show file tree
Hide file tree
Showing 3 changed files with 238 additions and 0 deletions.
6 changes: 6 additions & 0 deletions libs/bit_set/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
[package]
name = "bit_set"
version = "0.1.0"
authors = ["Klim Tsoutsman <klim@tsoutsman.com>"]
description = "A bit set storing integers less than 64"
edition = "2021"
67 changes: 67 additions & 0 deletions libs/bit_set/src/iter.rs
Original file line number Diff line number Diff line change
@@ -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]
);
}
}
165 changes: 165 additions & 0 deletions libs/bit_set/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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));
}
}

0 comments on commit 50c045f

Please sign in to comment.