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) -> u32 {
60        self.low.count_ones() + self.high.count_ones()
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    fn test_bitmap() {
103        let mut bitmap = Bitmap::new();
104
105        let indexes = vec![0, 10, 17, 120, 127, 128, 129, 200, 255];
106        for i in &indexes {
107            bitmap.set(*i);
108            assert!(bitmap.get(*i));
109        }
110
111        for i in 0..=255 {
112            assert_eq!(bitmap.get(i), indexes.contains(&i));
113        }
114        #[allow(clippy::cast_possible_truncation)]
115        {
116            assert_eq!(bitmap.count_ones(), indexes.len() as u32);
117        }
118
119        let value = bitmap.iter().collect::<Vec<_>>();
120        assert_eq!(value, indexes);
121
122        bitmap.invert();
123        for i in 0..=255 {
124            assert_eq!(bitmap.get(i), !indexes.contains(&i));
125        }
126    }
127
128    #[test]
129    fn test_bitmap_all() {
130        let mut bitmap = Bitmap::new();
131        assert_eq!(bitmap.iter().count(), 0);
132        bitmap.invert();
133        assert_eq!(bitmap.iter().count(), 256);
134    }
135
136    #[test]
137    fn test_bitmap_or_assign() {
138        let mut bitmap = Bitmap::new();
139        bitmap.set(10);
140        bitmap.set(30);
141
142        let mut bitmap2 = Bitmap::new();
143        bitmap2.set(20);
144        bitmap2 |= bitmap;
145
146        assert_eq!(vec![10, 20, 30], bitmap2.iter().collect::<Vec<_>>());
147    }
148}
boreal/bitmaps.rs

boreal/
bitmaps.rs
