Skip to main content

commonware_utils/bitmap/
atomic.rs

1//! A fixed-length bitmap whose bits can be set concurrently through shared references.
2
3use super::BitMap;
4use std::{
5    collections::VecDeque,
6    sync::atomic::{AtomicU64, Ordering},
7};
8
9/// A fixed-length bitmap whose bits can be set concurrently through shared references.
10///
11/// [Atomic::set] uses relaxed ordering, which is published by ownership transfer rather than
12/// by the operations themselves: the bits can only be read back by consuming the map
13/// ([Atomic::into_bitmap]), so whatever reclaims exclusive ownership from the setters
14/// (joining their tasks, `Arc::into_inner`) is what makes their writes visible.
15pub struct Atomic {
16    /// The bits, packed into words (bit `i` is bit `i % 64` of word `i / 64`).
17    ///
18    /// Invariant: `words.len() == len.div_ceil(64)`
19    /// Invariant: All bits at index `i` where `i >= len` are 0.
20    words: Vec<AtomicU64>,
21
22    /// The number of bits in the bitmap.
23    len: u64,
24}
25
26impl Atomic {
27    /// The size of a word in bits.
28    const WORD_BITS: u64 = u64::BITS as u64;
29
30    /// Create a bitmap of `len` zero bits.
31    pub fn zeroes(len: u64) -> Self {
32        let words = (0..len.div_ceil(Self::WORD_BITS))
33            .map(|_| AtomicU64::new(0))
34            .collect();
35        Self { words, len }
36    }
37
38    /// Set `bit` to 1.
39    ///
40    /// # Panics
41    ///
42    /// Panics if the bit doesn't exist.
43    pub fn set(&self, bit: u64) {
44        assert!(
45            bit < self.len,
46            "bit {} out of bounds (len: {})",
47            bit,
48            self.len
49        );
50        self.words[(bit / Self::WORD_BITS) as usize]
51            .fetch_or(1 << (bit % Self::WORD_BITS), Ordering::Relaxed);
52    }
53
54    /// Convert into a [BitMap] holding the same bits.
55    ///
56    /// Taking `self` by value means the caller already reclaimed exclusive ownership from
57    /// every setter, which is the synchronization that makes their relaxed writes visible.
58    pub fn into_bitmap(self) -> BitMap {
59        let chunks: VecDeque<_> = self
60            .words
61            .into_iter()
62            .map(|word| word.into_inner().to_le_bytes())
63            .collect();
64        BitMap::from_chunks(chunks, self.len)
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    /// The converted bitmap must match a [BitMap] built with plain sets, across lengths
73    /// covering the empty, sub-word, word-aligned, and multi-word-with-tail shapes.
74    #[test]
75    fn test_set_and_into_bitmap_matches_plain_sets() {
76        for len in [0u64, 1, 63, 64, 65, 128, 200] {
77            let atomic = Atomic::zeroes(len);
78            let mut expected = BitMap::zeroes(len);
79            for bit in (0..len).step_by(3).chain(len.checked_sub(1)) {
80                atomic.set(bit);
81                expected.set(bit, true);
82            }
83            assert_eq!(atomic.into_bitmap(), expected, "len {len}");
84        }
85    }
86
87    #[test]
88    fn test_shared_setters() {
89        let atomic = Atomic::zeroes(1000);
90        std::thread::scope(|s| {
91            for stripe in 0..4u64 {
92                let atomic = &atomic;
93                s.spawn(move || {
94                    for bit in (stripe..1000).step_by(4) {
95                        atomic.set(bit);
96                    }
97                });
98            }
99        });
100        assert_eq!(atomic.into_bitmap(), BitMap::ones(1000));
101    }
102
103    #[test]
104    #[should_panic(expected = "out of bounds")]
105    fn test_set_past_len_panics() {
106        Atomic::zeroes(64).set(64);
107    }
108
109    /// A bit past `len` must be rejected even when it lands inside the trailing word's
110    /// allocation, where the word indexing alone would accept it.
111    #[test]
112    #[should_panic(expected = "out of bounds")]
113    fn test_set_past_len_in_tail_word_panics() {
114        Atomic::zeroes(65).set(70);
115    }
116
117    #[test]
118    #[should_panic(expected = "out of bounds")]
119    fn test_set_on_empty_panics() {
120        Atomic::zeroes(0).set(0);
121    }
122}