eta_algorithms/data_structs/bitmap/
mod.rs

1use std::alloc::Layout;
2use std::ptr;
3
4use crate::data_structs::bitmap::consts::{DIV_SHIFT, MASK};
5use crate::data_structs::bitmap::handle::Handle;
6
7pub mod atomic_bitmap;
8pub mod handle;
9
10#[cfg(target_pointer_width = "64")]
11pub(crate) mod consts {
12    pub(crate) const DIV_SHIFT: usize = 6;
13    pub(crate) const MASK: usize = 63;
14}
15
16#[cfg(target_pointer_width = "32")]
17pub(self) mod consts {
18    pub(crate) const DIV_SHIFT: usize = 5;
19    pub(crate) const MASK: usize = 31;
20}
21
22#[cfg(target_pointer_width = "16")]
23pub(self) mod consts {
24    pub(crate) const DIV_SHIFT: usize = 4;
25    pub(crate) const MASK: usize = 15;
26}
27
28pub struct Bitmap {
29    data: *mut usize,
30    bit_capacity: usize,
31    capacity: usize,
32    layout: Layout,
33}
34
35impl Bitmap {
36    pub fn new(bit_count: usize) -> Self {
37        let size = (bit_count >> DIV_SHIFT) + 1;
38        let layout = Layout::array::<usize>(size).expect("Failed to create layout");
39        let data = unsafe { std::alloc::alloc(layout) as *mut usize };
40        unsafe { ptr::write_bytes(data, 0, size) };
41        Bitmap {
42            data,
43            capacity: size,
44            bit_capacity: bit_count,
45            layout,
46        }
47    }
48
49    pub fn check_batch(&self, handles: &[Handle]) -> bool {
50        for handle in handles {
51            let val = unsafe { *self.data.add(handle.chunk as usize) };
52            if (val & handle.bit_mask) != handle.bit_mask {
53                return false;
54            }
55        }
56        true
57    }
58
59    pub fn to_indices_true(&self) -> Vec<usize> {
60        let mut indices = Vec::new();
61        for i in 0..self.bit_capacity {
62            if unsafe { self.get_unchecked(i) } {
63                indices.push(i);
64            }
65        }
66        indices
67    }
68    pub fn to_indices_false(&self) -> Vec<usize> {
69        let mut indices = Vec::new();
70        for i in 0..self.bit_capacity {
71            if unsafe { self.get_unchecked(i) == false } {
72                indices.push(i);
73            }
74        }
75        indices
76    }
77
78    #[inline(always)]
79    pub fn bit_capacity(&self) -> usize {
80        self.bit_capacity
81    }
82    #[inline(always)]
83    pub fn capacity(&self) -> usize {
84        self.capacity
85    }
86    #[inline(always)]
87    pub fn set(&mut self, bit_index: usize, value: bool) {
88        if bit_index >= self.bit_capacity {
89            panic!("Bit index out of bounds");
90        }
91
92        let offset = bit_index >> DIV_SHIFT;
93        let bit_offset = bit_index & MASK;
94        unsafe {
95            let ptr = self.data.add(offset);
96            *ptr = (*ptr & !(1 << bit_offset)) | ((value as usize) << bit_offset);
97        }
98    }
99    #[inline(always)]
100    pub fn get(&self, bit_index: usize) -> Option<bool> {
101        if bit_index >= self.bit_capacity {
102            return None;
103        }
104
105        let offset = bit_index >> DIV_SHIFT;
106        let bit_offset = bit_index & (MASK);
107        unsafe {
108            let ptr = self.data.add(offset);
109            Some((*ptr & (1 << bit_offset)) != 0)
110        }
111    }
112    #[inline(always)]
113    pub unsafe fn set_unchecked(&mut self, bit_index: usize, value: bool) {
114        let offset = bit_index >> DIV_SHIFT;
115        let bit_offset = bit_index & (MASK);
116        unsafe {
117            let ptr = self.data.add(offset);
118            *ptr = (*ptr & !(1 << bit_offset)) | ((value as usize) << bit_offset);
119        }
120    }
121    #[inline(always)]
122    pub unsafe fn get_unchecked(&self, bit_index: usize) -> bool {
123        let offset = bit_index >> DIV_SHIFT;
124        let bit_offset = bit_index & (MASK);
125        unsafe {
126            let ptr = self.data.add(offset);
127            (*ptr & (1 << bit_offset)) != 0
128        }
129    }
130}
131
132impl Drop for Bitmap {
133    fn drop(&mut self) {
134        unsafe {
135            std::alloc::dealloc(self.data as *mut u8, self.layout);
136        }
137    }
138}