eta_algorithms/data_structs/bitmap/
mod.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
use std::alloc::Layout;
use std::ptr;

use crate::data_structs::bitmap::consts::{DIV_SHIFT, MASK};
use crate::data_structs::bitmap::handle::Handle;

pub mod atomic_bitmap;
pub mod handle;

#[cfg(target_pointer_width = "64")]
pub(crate) mod consts {
    pub(crate) const DIV_SHIFT: usize = 6;
    pub(crate) const MASK: usize = 63;
}

#[cfg(target_pointer_width = "32")]
pub(self) mod consts {
    pub(crate) const DIV_SHIFT: usize = 5;
    pub(crate) const MASK: usize = 31;
}

#[cfg(target_pointer_width = "16")]
pub(self) mod consts {
    pub(crate) const DIV_SHIFT: usize = 4;
    pub(crate) const MASK: usize = 15;
}

pub struct Bitmap {
    data: *mut usize,
    bit_capacity: usize,
    capacity: usize,
    layout: Layout,
}

impl Bitmap {
    pub fn new(bit_count: usize) -> Self {
        let size = (bit_count >> DIV_SHIFT) + 1;
        let layout = Layout::array::<usize>(size).expect("Failed to create layout");
        let data = unsafe { std::alloc::alloc(layout) as *mut usize };
        unsafe { ptr::write_bytes(data, 0, size) };
        Bitmap {
            data,
            capacity: size,
            bit_capacity: bit_count,
            layout,
        }
    }

    pub fn check_batch(&self, handles: &[Handle]) -> bool {
        for handle in handles {
            let val = unsafe { *self.data.add(handle.chunk as usize) };
            if (val & handle.bit_mask) != handle.bit_mask {
                return false;
            }
        }
        true
    }

    pub fn to_indices_true(&self) -> Vec<usize> {
        let mut indices = Vec::new();
        for i in 0..self.bit_capacity {
            if unsafe { self.get_unchecked(i) } {
                indices.push(i);
            }
        }
        indices
    }
    pub fn to_indices_false(&self) -> Vec<usize> {
        let mut indices = Vec::new();
        for i in 0..self.bit_capacity {
            if unsafe { self.get_unchecked(i) == false } {
                indices.push(i);
            }
        }
        indices
    }

    #[inline(always)]
    pub fn bit_capacity(&self) -> usize {
        self.bit_capacity
    }
    #[inline(always)]
    pub fn capacity(&self) -> usize {
        self.capacity
    }
    #[inline(always)]
    pub fn set(&mut self, bit_index: usize, value: bool) {
        if bit_index >= self.bit_capacity {
            panic!("Bit index out of bounds");
        }

        let offset = bit_index >> DIV_SHIFT;
        let bit_offset = bit_index & MASK;
        unsafe {
            let ptr = self.data.add(offset);
            *ptr = (*ptr & !(1 << bit_offset)) | ((value as usize) << bit_offset);
        }
    }
    #[inline(always)]
    pub fn get(&self, bit_index: usize) -> Option<bool> {
        if bit_index >= self.bit_capacity {
            return None;
        }

        let offset = bit_index >> DIV_SHIFT;
        let bit_offset = bit_index & (MASK);
        unsafe {
            let ptr = self.data.add(offset);
            Some((*ptr & (1 << bit_offset)) != 0)
        }
    }
    #[inline(always)]
    pub unsafe fn set_unchecked(&mut self, bit_index: usize, value: bool) {
        let offset = bit_index >> DIV_SHIFT;
        let bit_offset = bit_index & (MASK);
        unsafe {
            let ptr = self.data.add(offset);
            *ptr = (*ptr & !(1 << bit_offset)) | ((value as usize) << bit_offset);
        }
    }
    #[inline(always)]
    pub unsafe fn get_unchecked(&self, bit_index: usize) -> bool {
        let offset = bit_index >> DIV_SHIFT;
        let bit_offset = bit_index & (MASK);
        unsafe {
            let ptr = self.data.add(offset);
            (*ptr & (1 << bit_offset)) != 0
        }
    }
}

impl Drop for Bitmap {
    fn drop(&mut self) {
        unsafe {
            std::alloc::dealloc(self.data as *mut u8, self.layout);
        }
    }
}