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}