egglog_concurrency/
bitset.rs

1//! A simple dense bitset that leverages [`crate::ReadOptimizedLock`] for safe
2//! resizing and provides wait-free reads and writes when no resizing is
3//! required.
4
5use std::{
6    mem,
7    sync::atomic::{AtomicU64, Ordering},
8};
9
10use crate::ReadOptimizedLock;
11
12pub struct BitSet {
13    data: ReadOptimizedLock<Vec<AtomicU64>>,
14}
15
16impl BitSet {
17    pub fn with_capacity(n: usize) -> Self {
18        let n = n.next_multiple_of(64).next_power_of_two();
19        let cells = n / 64;
20        let mut data = Vec::with_capacity(cells);
21        data.resize_with(cells, AtomicU64::default);
22        BitSet {
23            data: ReadOptimizedLock::new(data),
24        }
25    }
26
27    /// Get the value of the bit at index `i`. If the bitset has not been
28    /// initialized up to `i`, then this method returns `false`.
29    pub fn get(&self, i: usize) -> bool {
30        let cell = i / 64;
31        let bit = i % 64;
32        let reader = self.data.read();
33        reader
34            .get(cell)
35            .map(|x| x.load(Ordering::Acquire) & (1 << bit) != 0)
36            .unwrap_or(false)
37    }
38
39    /// Set the bit at index `i` to `val`. If the bitset has not been
40    /// initialized up to `i`, the bitset will be resized. The resizing
41    /// operation will block until all current readers have finished.
42    pub fn set(&self, i: usize, val: bool) {
43        let cell = i / 64;
44        let bit = i % 64;
45        let handle = self.data.read();
46        if let Some(cell) = handle.get(cell) {
47            if val {
48                cell.fetch_or(1 << bit, Ordering::Release);
49            } else {
50                cell.fetch_and(!(1 << bit), Ordering::Release);
51            }
52            return;
53        }
54        mem::drop(handle);
55        let mut writer = self.data.lock();
56        if cell >= writer.len() {
57            writer.resize_with((cell + 1).next_power_of_two(), AtomicU64::default);
58        }
59        mem::drop(writer);
60        self.set(i, val);
61    }
62}
egglog_concurrency/bitset.rs

egglog_concurrency/
bitset.rs
