index_set/atomic_bitset.rs
1use crate::*;
2
3/// A thread-safe bitset that uses atomic operations to manage bits.
4pub struct AtomicBitSet<const N: usize> {
5 bitset: [AtomicUsize; N],
6 // used for optimizing the search to find the next free bit
7 rotation: AtomicUsize,
8}
9
10impl<const N: usize> AtomicBitSet<N> {
11 /// Creates a new `AtomicBitSet` with the specified number of slots.
12 /// Each slot can hold 32/64 bits depending on the architecture.
13 ///
14 /// ## Examples
15 ///
16 /// ```rust
17 /// use index_set::{AtomicBitSet, slot_count};
18 ///
19 /// let bitset: AtomicBitSet<{ slot_count::from_bits(128) }> = AtomicBitSet::new();
20 /// ```
21 #[inline]
22 #[allow(clippy::new_without_default)]
23 pub const fn new() -> Self {
24 Self {
25 bitset: [const { AtomicUsize::new(0) }; N],
26 rotation: AtomicUsize::new(0),
27 }
28 }
29
30 /// Atomically finds the next free bit (unset bit with value `0`) in the bitset, sets it to `1`,
31 /// and returns its index.
32 ///
33 /// This method is thread-safe and can be used in concurrent environments.
34 ///
35 /// ## Examples
36 ///
37 /// ```rust
38 /// # use index_set::{AtomicBitSet, slot_count};
39 /// let ids: AtomicBitSet<{ slot_count::from_bits(128) }> = AtomicBitSet::new();
40 /// assert_eq!(ids.set_next_free_bit(), Some(0));
41 /// assert_eq!(ids.set_next_free_bit(), Some(1));
42 /// ```
43 pub fn set_next_free_bit(&self) -> Option<usize> {
44 // rotate the slots to find the next free id
45 let skip = self.rotation.load(Ordering::Relaxed);
46 let mut slot_idx = skip;
47
48 let slots = utils::rotate_left(&self.bitset, skip);
49
50 for slot in slots {
51 let available_slot = slot.fetch_update(Ordering::AcqRel, Ordering::Acquire, |curr| {
52 // slot is full
53 if curr == usize::MAX {
54 return None;
55 }
56 let next_available_bit = (!curr).trailing_zeros() as usize;
57 Some(curr | (1 << next_available_bit))
58 });
59
60 if let Ok(curr) = available_slot {
61 if skip != slot_idx {
62 self.rotation.store(slot_idx, Ordering::Relaxed);
63 }
64 let next_available_bit = (!curr).trailing_zeros() as usize;
65 return Some(slot_idx * usize::BITS as usize + next_available_bit);
66 }
67
68 slot_idx += 1;
69 if slot_idx >= N {
70 slot_idx = 0;
71 }
72 }
73 None
74 }
75}
76
77impl<const N: usize> std::ops::Deref for AtomicBitSet<N> {
78 type Target = [AtomicUsize];
79
80 #[inline]
81 fn deref(&self) -> &Self::Target {
82 &self.bitset
83 }
84}