eta_algorithms/data_structs/bitmap/
mod.rs

1use std::alloc::Layout;
2use std::ptr;
3
4use crate::data_structs::bitmap::consts::{BIT_END_OFFSET, BIT_MASK, DIV_SHIFT};
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; // Divide by 64 => 2^6
13    pub(crate) const BIT_END_OFFSET: usize = 63;
14    pub(crate) const BIT_MASK: usize = 0xFFFFFFFFFFFFFFFF;
15}
16
17#[cfg(target_pointer_width = "32")]
18pub(self) mod consts {
19    pub(crate) const DIV_SHIFT: usize = 5;
20    pub(crate) const BIT_END_OFFSET: usize = 31;
21    pub(crate) const BIT_MASK: usize = 0xFFFFFFFF;
22}
23
24#[cfg(target_pointer_width = "16")]
25pub(self) mod consts {
26    pub(crate) const DIV_SHIFT: usize = 4;
27    pub(crate) const BIT_END_OFFSET: usize = 15;
28    pub(crate) const BIT_MASK: usize = 0xFFFF;
29}
30
31pub struct Bitmap {
32    data: *mut usize,
33    bit_capacity: usize,
34    capacity: usize,
35    layout: Layout,
36}
37
38impl Bitmap {
39    pub fn new(bit_count: usize) -> Self {
40        let size = (bit_count >> DIV_SHIFT) + 1;
41        let layout = Layout::array::<usize>(size).expect("Failed to create layout");
42        let data = unsafe { std::alloc::alloc(layout) as *mut usize };
43        unsafe { ptr::write_bytes(data, 0, size) };
44        Bitmap {
45            data,
46            capacity: size,
47            bit_capacity: bit_count,
48            layout,
49        }
50    }
51
52    #[inline(always)]
53    pub fn count_zeros(&self, lower_bound: usize, upper_bound: usize) -> usize {
54        if upper_bound == 0 {
55            panic!("Upper bound cannot be zero");
56        }
57
58        if lower_bound > upper_bound {
59            panic!("Lower bound cannot be greater than upper bound");
60        }
61
62        if lower_bound >= self.bit_capacity {
63            panic!("Bit index out of bounds");
64        }
65        if upper_bound > self.bit_capacity {
66            panic!("Upper bound out of bounds");
67        }
68
69        if lower_bound == upper_bound {
70            return 0;
71        }
72
73        unsafe { self.count_zeros_unchecked(lower_bound, upper_bound - 1) }
74    }
75    #[inline(always)]
76    pub fn count_ones(&self, lower_bound: usize, upper_bound: usize) -> usize {
77        if upper_bound == 0 {
78            panic!("Upper bound cannot be zero");
79        }
80
81        if lower_bound > upper_bound {
82            panic!("Lower bound cannot be greater than upper bound");
83        }
84
85        if lower_bound >= self.bit_capacity {
86            panic!("Bit index out of bounds");
87        }
88
89        if upper_bound > self.bit_capacity {
90            panic!("Upper bound out of bounds");
91        }
92
93        if lower_bound == upper_bound {
94            return 0;
95        }
96
97        unsafe { self.count_ones_unchecked(lower_bound, upper_bound - 1) }
98    }
99
100    pub unsafe fn count_zeros_unchecked(&self, lower_bound: usize, upper_bound: usize) -> usize {
101        let lower_offset = lower_bound >> DIV_SHIFT;
102        let lower_bit_offset = lower_bound & BIT_END_OFFSET;
103        let upper_offset = upper_bound >> DIV_SHIFT;
104        let upper_bit_offset = upper_bound & BIT_END_OFFSET;
105        let mut data_ptr = self.data.add(lower_offset);
106        let mut data = if lower_offset != upper_offset {
107            *data_ptr | ((1 << lower_bit_offset) - 1)
108        } else {
109            (*data_ptr | ((1 << lower_bit_offset) - 1)) | !(((1 << upper_bit_offset) - 1) | (1 << upper_bit_offset))
110            // + 1 because inclusive and we need to do upper and lower bound masking
111        };
112        let mut counter = data.count_zeros() as usize;
113        let end = self.data.add(upper_offset);
114        loop {
115            data_ptr = data_ptr.add(1);
116            data = *data_ptr;
117            if data_ptr >= end {
118                break;
119            }
120            counter += data.count_zeros() as usize;
121        }
122        // The remainder
123        if upper_offset != lower_offset {
124            let data_end_ptr = data_ptr;
125            // The other OR is there to add 1 because upper_bound is inclusive
126            let data_end = *data_end_ptr | !(((1 << upper_bit_offset) - 1) | (1 << upper_bit_offset)); // + 1 because inclusive
127            counter += data_end.count_zeros() as usize;
128        }
129
130        counter
131    }
132
133    pub unsafe fn count_ones_unchecked(&self, lower_bound: usize, upper_bound: usize) -> usize {
134        let lower_offset = lower_bound >> DIV_SHIFT;
135        let lower_bit_offset = lower_bound & BIT_END_OFFSET;
136        let upper_offset = upper_bound >> DIV_SHIFT;
137        let upper_bit_offset = upper_bound & BIT_END_OFFSET;
138        let mut data_ptr = self.data.add(lower_offset);
139        let mut data = if lower_offset != upper_offset {
140            *data_ptr & (!((1 << lower_bit_offset) - 1))
141        } else {
142            (*data_ptr & (!((1 << lower_bit_offset) - 1))) & (((1 << upper_bit_offset) - 1) | (1 << upper_bit_offset))
143            // + 1 because inclusive and we need to do upper and lower bound masking
144        };
145        let mut counter = data.count_ones() as usize;
146        let end = self.data.add(upper_offset);
147        loop {
148            data_ptr = data_ptr.add(1);
149            data = *data_ptr;
150            if data_ptr >= end {
151                break;
152            }
153            counter += data.count_ones() as usize;
154        }
155        // The remainder
156        if upper_offset != lower_offset {
157            let data_end_ptr = data_ptr;
158            // The other OR is there to add 1 because upper_bound is inclusive
159            let data_end = *data_end_ptr & (((1 << upper_bit_offset) - 1) | (1 << upper_bit_offset));
160            counter += data_end.count_ones() as usize;
161        }
162        counter
163    }
164    #[inline(always)]
165    pub fn first_one_bounds(&self, lower_bound: usize, upper_bound: usize) -> Option<usize> {
166        if upper_bound == 0 {
167            panic!("Upper bound cannot be zero");
168        }
169
170        if lower_bound > upper_bound {
171            panic!("Lower bound cannot be greater than upper bound");
172        }
173
174        if lower_bound >= self.bit_capacity {
175            panic!("Bit index out of bounds");
176        }
177        if upper_bound > self.bit_capacity {
178            panic!("Upper bound out of bounds");
179        }
180        let index = unsafe { self.first_one_unchecked(lower_bound) };
181        if index >= upper_bound {
182            None
183        } else {
184            Some(index)
185        }
186    }
187    #[inline(always)]
188    pub fn first_zero_bounds(&self, lower_bound: usize, upper_bound: usize) -> Option<usize> {
189        if upper_bound == 0 {
190            panic!("Upper bound cannot be zero");
191        }
192
193        if lower_bound > upper_bound {
194            panic!("Lower bound cannot be greater than upper bound");
195        }
196
197        if lower_bound >= self.bit_capacity {
198            panic!("Bit index out of bounds");
199        }
200        if upper_bound > self.bit_capacity {
201            panic!("Upper bound out of bounds");
202        }
203        let index = unsafe { self.first_zero_unchecked(lower_bound) };
204        if index >= upper_bound {
205            None
206        } else {
207            Some(index)
208        }
209    }
210    #[inline(always)]
211    pub fn first_zero(&self, bit_index: usize) -> Option<usize> {
212        if bit_index >= self.bit_capacity {
213            panic!("Bit index out of bounds");
214        }
215
216        let index = unsafe { self.first_zero_unchecked(bit_index) };
217        if index >= self.bit_capacity {
218            None
219        } else {
220            Some(index)
221        }
222    }
223    #[inline(always)]
224    pub fn first_one(&self, bit_index: usize) -> Option<usize> {
225        if bit_index >= self.bit_capacity {
226            panic!("Bit index out of bounds");
227        }
228
229        let index = unsafe { self.first_one_unchecked(bit_index) };
230        if index >= self.bit_capacity {
231            None
232        } else {
233            Some(index)
234        }
235    }
236    /// Returns higher > bit_capacity in case of not found
237    pub unsafe fn first_zero_unchecked(&self, bit_index: usize) -> usize {
238        let offset = bit_index >> DIV_SHIFT;
239        let bit_offset = bit_index & BIT_END_OFFSET;
240
241        //We scan the last chunk with trailing, if it fails it shall return the max
242        let last_chunk = self.data.add(self.capacity - 1);
243        let mut data_ptr = self.data.add(offset);
244        let mut counter = offset * (BIT_END_OFFSET + 1);
245        let mut data = *data_ptr | ((1 << bit_offset) - 1);
246
247        while data & BIT_MASK == BIT_MASK {
248            if data_ptr != last_chunk {
249                data_ptr = data_ptr.add(1);
250                data = *data_ptr;
251                counter += BIT_END_OFFSET + 1;
252                continue;
253            }
254            break;
255        }
256        counter += data.trailing_ones() as usize;
257        counter
258    }
259
260    /// Returns higher > bit_capacity in case of not found
261    pub unsafe fn first_one_unchecked(&self, bit_index: usize) -> usize {
262        let offset = bit_index >> DIV_SHIFT;
263        let bit_offset = bit_index & BIT_END_OFFSET;
264
265        //We scan the last chunk with trailing, if it fails it shall return the max
266        let last_chunk = self.data.add(self.capacity - 1);
267        let mut data_ptr = self.data.add(offset);
268        let mut counter = offset * (BIT_END_OFFSET + 1);
269        let mut data = *data_ptr & !((1 << bit_offset) - 1);
270
271        while !data & BIT_MASK == BIT_MASK {
272            if data_ptr != last_chunk {
273                data_ptr = data_ptr.add(1);
274                data = *data_ptr;
275                counter += BIT_END_OFFSET + 1;
276                continue;
277            }
278            break;
279        }
280        counter += data.trailing_zeros() as usize;
281        counter
282    }
283    pub fn check_batch(&self, handles: &[Handle]) -> bool {
284        for handle in handles {
285            let val = unsafe { *self.data.add(handle.chunk as usize) };
286            if (val & handle.bit_mask) != handle.bit_mask {
287                return false;
288            }
289        }
290        true
291    }
292
293    pub fn to_indices_true(&self) -> Vec<usize> {
294        let mut indices = Vec::new();
295        for i in 0..self.bit_capacity {
296            if unsafe { self.get_unchecked(i) } {
297                indices.push(i);
298            }
299        }
300        indices
301    }
302
303    pub fn to_indices_true_bounded(&self, start: usize, end: usize) -> Vec<usize> {
304        if start >= end {
305            panic!("Start must be less than end");
306        }
307        if end > self.bit_capacity {
308            panic!("End must be less than or equal to bit capacity");
309        }
310
311        let mut indices = Vec::new();
312        for i in start..end {
313            if unsafe { self.get_unchecked(i) } {
314                indices.push(i);
315            }
316        }
317        indices
318    }
319
320    pub fn to_indices_false(&self) -> Vec<usize> {
321        let mut indices = Vec::new();
322        for i in 0..self.bit_capacity {
323            if unsafe { self.get_unchecked(i) == false } {
324                indices.push(i);
325            }
326        }
327        indices
328    }
329
330    pub fn to_indices_false_bounded(&self, start: usize, end: usize) -> Vec<usize> {
331        if start >= end {
332            panic!("Start must be less than end");
333        }
334        if end > self.bit_capacity {
335            panic!("End must be less than or equal to bit capacity");
336        }
337
338        let mut indices = Vec::new();
339        for i in start..end {
340            if unsafe { self.get_unchecked(i) == false } {
341                indices.push(i);
342            }
343        }
344        indices
345    }
346
347    #[inline(always)]
348    pub fn bit_capacity(&self) -> usize {
349        self.bit_capacity
350    }
351    #[inline(always)]
352    pub fn capacity(&self) -> usize {
353        self.capacity
354    }
355    #[inline(always)]
356    pub fn set(&mut self, bit_index: usize, value: bool) {
357        if bit_index >= self.bit_capacity {
358            panic!("Bit index out of bounds");
359        }
360
361        let offset = bit_index >> DIV_SHIFT;
362        let bit_offset = bit_index & BIT_END_OFFSET;
363        unsafe {
364            let ptr = self.data.add(offset);
365            *ptr = (*ptr & !(1 << bit_offset)) | ((value as usize) << bit_offset);
366        }
367    }
368    #[inline(always)]
369    pub fn get(&self, bit_index: usize) -> Option<bool> {
370        if bit_index >= self.bit_capacity {
371            return None;
372        }
373
374        let offset = bit_index >> DIV_SHIFT;
375        let bit_offset = bit_index & (BIT_END_OFFSET);
376        unsafe {
377            let ptr = self.data.add(offset);
378            Some((*ptr & (1 << bit_offset)) != 0)
379        }
380    }
381    #[inline(always)]
382    pub unsafe fn set_unchecked(&mut self, bit_index: usize, value: bool) {
383        let offset = bit_index >> DIV_SHIFT;
384        let bit_offset = bit_index & (BIT_END_OFFSET);
385        unsafe {
386            let ptr = self.data.add(offset);
387            *ptr = (*ptr & !(1 << bit_offset)) | ((value as usize) << bit_offset);
388        }
389    }
390    #[inline(always)]
391    pub unsafe fn get_unchecked(&self, bit_index: usize) -> bool {
392        let offset = bit_index >> DIV_SHIFT;
393        let bit_offset = bit_index & (BIT_END_OFFSET);
394        unsafe {
395            let ptr = self.data.add(offset);
396            (*ptr & (1 << bit_offset)) != 0
397        }
398    }
399}
400
401impl Drop for Bitmap {
402    fn drop(&mut self) {
403        unsafe {
404            std::alloc::dealloc(self.data as *mut u8, self.layout);
405        }
406    }
407}
eta_algorithms/data_structs/bitmap/mod.rs

eta_algorithms/data_structs/bitmap/
mod.rs
