yui_core/conc/misc/
bitmap.rs

1//! Compact bitmap over a small element type. `BitMap<E, S>` is a single
2//! word `S` (the storage) whose bit `e.into()` is set iff `e` is present.
3//! `S = u128` covers element indices `0..128`, and [`U256`](super::u256::U256)
4//! doubles that; extend by impl-ing [`BitStorage`] for a wider type.
5
6use std::marker::PhantomData;
7use std::ops::{BitAnd, BitOr, BitOrAssign, Shl, Sub};
8
9/// Single-word storage backing a [`BitMap`]. Implemented for `u8`..`u128` and
10/// [`U256`](super::u256::U256).
11pub trait BitStorage:
12    Copy + Eq + Default
13    + BitAnd<Output = Self> + BitOr<Output = Self> + BitOrAssign
14    + Shl<u32, Output = Self> + Sub<Output = Self>
15{
16    /// Number of bits this storage can hold (capacity of the bitmap).
17    const WIDTH: u32;
18
19    fn one() -> Self;
20    fn is_zero(self) -> bool;
21    fn count_ones(self) -> u32;
22    fn trailing_zeros(self) -> u32;
23}
24
25macro_rules! impl_bit_storage {
26    ($($t:ty),+ $(,)?) => {
27        $(
28            impl BitStorage for $t {
29                const WIDTH: u32 = <$t>::BITS;
30                fn one() -> Self { 1 }
31                fn is_zero(self) -> bool { self == 0 }
32                fn count_ones(self) -> u32 { <$t>::count_ones(self) }
33                fn trailing_zeros(self) -> u32 { <$t>::trailing_zeros(self) }
34            }
35        )+
36    };
37}
38
39impl_bit_storage!(u8, u16, u32, u64, u128);
40
41#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Default)]
42pub struct BitMap<E, S> {
43    bits: S,
44    _phantom: PhantomData<E>,
45}
46
47impl<E, S> BitMap<E, S>
48where
49    E: Copy + Into<u32> + TryFrom<u32>,
50    S: BitStorage,
51{
52    pub fn new() -> Self {
53        Self { bits: S::default(), _phantom: PhantomData }
54    }
55
56    pub fn insert(&mut self, e: E) {
57        let bit = e.into();
58        // note: not a `debug_assert` — in release the shift is masked and the write wraps silently.
59        assert!(bit < S::WIDTH, "BitMap index {bit} ≥ storage width {}", S::WIDTH);
60        self.bits |= S::one() << bit;
61    }
62
63    pub fn contains(&self, e: E) -> bool {
64        let bit = e.into();
65        // hot read path; every present bit already went through the checked `insert`.
66        debug_assert!(bit < S::WIDTH, "BitMap index {bit} ≥ storage width {}", S::WIDTH);
67        !(self.bits & (S::one() << bit)).is_zero()
68    }
69
70    pub fn len(&self) -> usize {
71        self.bits.count_ones() as usize
72    }
73
74    pub fn is_empty(&self) -> bool {
75        self.bits.is_zero()
76    }
77
78    /// Iterate elements in ascending order. The smallest element (if any) is
79    /// `self.iter().next()`.
80    pub fn iter(&self) -> Iter<E, S> {
81        Iter { bits: self.bits, _phantom: PhantomData }
82    }
83}
84
85impl<E, S> FromIterator<E> for BitMap<E, S>
86where
87    E: Copy + Into<u32> + TryFrom<u32>,
88    S: BitStorage,
89{
90    fn from_iter<I: IntoIterator<Item = E>>(iter: I) -> Self {
91        let mut m = Self::new();
92        for e in iter { m.insert(e); }
93        m
94    }
95}
96
97impl<E, S: BitStorage> BitOrAssign for BitMap<E, S> {
98    fn bitor_assign(&mut self, rhs: Self) {
99        self.bits |= rhs.bits;
100    }
101}
102
103impl<E, S: BitStorage> BitOr for BitMap<E, S> {
104    type Output = Self;
105    fn bitor(self, rhs: Self) -> Self {
106        Self { bits: self.bits | rhs.bits, _phantom: PhantomData }
107    }
108}
109
110pub struct Iter<E, S> {
111    bits: S,
112    _phantom: PhantomData<E>,
113}
114
115impl<E, S> Iterator for Iter<E, S>
116where
117    E: TryFrom<u32>,
118    S: BitStorage,
119{
120    type Item = E;
121    fn next(&mut self) -> Option<E> {
122        if self.bits.is_zero() { return None; }
123        let bit = self.bits.trailing_zeros();
124        self.bits = self.bits & (self.bits - S::one());
125        E::try_from(bit).ok()
126    }
127}
128
129impl<E, S> ExactSizeIterator for Iter<E, S>
130where
131    E: TryFrom<u32>,
132    S: BitStorage,
133{
134    fn len(&self) -> usize {
135        self.bits.count_ones() as usize
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    type B = BitMap<u8, u128>;
143
144    #[test]
145    fn empty() {
146        let m = B::new();
147        assert!(m.is_empty());
148        assert_eq!(m.len(), 0);
149        assert_eq!(m.iter().next(), None);
150        assert!(!m.contains(0));
151        assert_eq!(m.iter().collect::<Vec<u8>>(), Vec::<u8>::new());
152    }
153
154    #[test]
155    fn insert_contains() {
156        let mut m = B::new();
157        m.insert(3);
158        m.insert(7);
159        m.insert(3);
160        assert_eq!(m.len(), 2);
161        assert!(m.contains(3));
162        assert!(m.contains(7));
163        assert!(!m.contains(5));
164    }
165
166    #[test]
167    fn from_iter_iter_ascending() {
168        let m: B = [10u8, 3, 7, 0].into_iter().collect();
169        assert_eq!(m.len(), 4);
170        assert_eq!(m.iter().next(), Some(0));
171        assert_eq!(m.iter().collect::<Vec<u8>>(), vec![0, 3, 7, 10]);
172    }
173
174    #[test]
175    fn bitor_union() {
176        let a: B = [1, 5].into_iter().collect();
177        let b: B = [5, 9].into_iter().collect();
178        let u = a | b;
179        assert_eq!(u.iter().collect::<Vec<u8>>(), vec![1, 5, 9]);
180    }
181
182    #[test]
183    fn small_storage_u8() {
184        // Max-8 BitMap: indices 0..8.
185        type Small = BitMap<u8, u8>;
186        let mut m = Small::new();
187        m.insert(0);
188        m.insert(7);
189        assert_eq!(m.len(), 2);
190        assert_eq!(m.iter().collect::<Vec<u8>>(), vec![0, 7]);
191    }
192
193    #[test]
194    #[should_panic(expected = "BitMap index")]
195    fn small_storage_out_of_range_panics() {
196        let mut m: BitMap<u8, u8> = BitMap::new();
197        m.insert(8);  // u8 only has bits 0..8
198    }
199}
yui_core/conc/misc/bitmap.rs

yui_core/conc/misc/
bitmap.rs
