boreal/
bitmaps.rs

1use std::ops::BitOrAssign;
2
3#[derive(Debug, Clone, Copy, Default)]
4/// A bitmap with 256 bits
5pub struct Bitmap {
6    low: u128,
7    high: u128,
8}
9
10impl Bitmap {
11    const HALF: u8 = 128;
12
13    pub fn new() -> Self {
14        Self::default()
15    }
16
17    fn mask(bit: u8) -> u128 {
18        1u128 << (bit & 127)
19    }
20
21    fn get_half(&self, bit: u8) -> u128 {
22        if bit < Self::HALF {
23            self.low
24        } else {
25            self.high
26        }
27    }
28
29    fn get_half_mut(&mut self, bit: u8) -> &mut u128 {
30        if bit < Self::HALF {
31            &mut self.low
32        } else {
33            &mut self.high
34        }
35    }
36
37    #[must_use]
38    #[inline(always)]
39    pub fn get(&self, bit: u8) -> bool {
40        let mask = Self::mask(bit);
41        let half = self.get_half(bit);
42        half & mask != 0
43    }
44
45    #[inline(always)]
46    pub fn set(&mut self, bit: u8) {
47        let mask = Self::mask(bit);
48        let half = self.get_half_mut(bit);
49        *half |= mask;
50    }
51
52    #[inline(always)]
53    pub fn invert(&mut self) {
54        self.low = !self.low;
55        self.high = !self.high;
56    }
57
58    #[inline(always)]
59    pub fn count_ones(&self) -> usize {
60        (self.low.count_ones() + self.high.count_ones()) as usize
61    }
62
63    pub fn iter(&self) -> Iter {
64        Iter(*self)
65    }
66}
67
68/// implement `|=`
69impl BitOrAssign for Bitmap {
70    fn bitor_assign(&mut self, rhs: Self) {
71        self.low |= rhs.low;
72        self.high |= rhs.high;
73    }
74}
75
76pub struct Iter(Bitmap);
77
78impl Iterator for Iter {
79    type Item = u8;
80
81    fn next(&mut self) -> Option<Self::Item> {
82        let (v, offset) = if self.0.low != 0 {
83            (&mut self.0.low, 0)
84        } else if self.0.high != 0 {
85            (&mut self.0.high, 128)
86        } else {
87            return None;
88        };
89
90        // Safety: this value is contained in [0; 127] so it always fits in a u8.
91        let t: u8 = v.trailing_zeros().try_into().unwrap();
92        *v &= !(1 << t);
93        Some(offset + t)
94    }
95}
96
97#[cfg(test)]
98mod test {
99    use super::Bitmap;
100
101    #[test]
102
103    fn test_bitmap() {
104        let mut bitmap = Bitmap::new();
105
106        let indexes = vec![0, 10, 17, 120, 127, 128, 129, 200, 255];
107        for i in &indexes {
108            bitmap.set(*i);
109            assert!(bitmap.get(*i));
110        }
111
112        for i in 0..=255 {
113            assert_eq!(bitmap.get(i), indexes.contains(&i));
114        }
115        assert_eq!(bitmap.count_ones(), indexes.len());
116
117        let value = bitmap.iter().collect::<Vec<_>>();
118        assert_eq!(value, indexes);
119
120        bitmap.invert();
121        for i in 0..=255 {
122            assert_eq!(bitmap.get(i), !indexes.contains(&i));
123        }
124    }
125
126    #[test]
127    fn test_bitmap_all() {
128        let mut bitmap = Bitmap::new();
129        assert_eq!(bitmap.iter().count(), 0);
130        bitmap.invert();
131        assert_eq!(bitmap.iter().count(), 256);
132    }
133
134    #[test]
135    fn test_bitmap_or_assign() {
136        let mut bitmap = Bitmap::new();
137        bitmap.set(10);
138        bitmap.set(30);
139
140        let mut bitmap2 = Bitmap::new();
141        bitmap2.set(20);
142        bitmap2 |= bitmap;
143
144        assert_eq!(vec![10, 20, 30], bitmap2.iter().collect::<Vec<_>>());
145    }
146}
boreal/bitmaps.rs

boreal/
bitmaps.rs
