bit_index/
lib.rs

1use std::cmp::max;
2use std::fmt::{self, Debug};
3
4macro_rules! impl_bit_index {
5    ($bit_index_name:ident, $bit_index_type:ty) => {
6        /// A list of bits to track elements. Little-endian and zero-indexed.`
7        #[derive(Copy, Clone, PartialEq, Eq, Hash)]
8        pub struct $bit_index_name {
9            /// The bits to track elements
10            bits: $bit_index_type,
11            /// The number of elements to track. Leading zeros do not represent anything, the zero's in the least `nb_bits` positions represent the absence of the corresponding element.
12            nb_bits: u8,
13        }
14
15        impl $bit_index_name {
16            const SIZE: u8 = std::mem::size_of::<$bit_index_type>() as u8 * 8;
17
18            pub fn new(nb_bits: u8) -> Result<Self, String> {
19                if nb_bits > Self::SIZE {
20                    Err(format!(
21                        "This BitIndex can only keep {} bits, not {}",
22                        Self::SIZE,
23                        nb_bits
24                    ))
25                } else {
26                    Ok(Self {
27                        bits: Self::init(nb_bits),
28                        nb_bits,
29                    })
30                }
31            }
32
33            pub fn empty(nb_bits: u8) -> Result<Self, String> {
34                Self::new(nb_bits).map(|mut bi| {
35                    bi.clear();
36                    bi
37                })
38            }
39
40            pub fn unwrap(&self) -> $bit_index_type {
41                self.bits
42            }
43
44            #[inline]
45            pub fn is_empty(&self) -> bool {
46                self.bits == 0
47            }
48
49            #[inline]
50            pub fn clear(&mut self) {
51                self.bits = 0;
52            }
53
54            pub fn restore(&mut self) {
55                self.bits = Self::init(self.nb_bits);
56            }
57
58            pub fn nb_elements(&self) -> u8 {
59                self.bits.count_ones() as u8
60            }
61
62            pub fn get(&mut self, idx: u8) -> Option<u8> {
63                self.get_from_low_end(idx)
64            }
65
66            pub fn get_from_low_end(&self, idx: u8) -> Option<u8> {
67                let nb_elements = self.nb_elements();
68                self.get_check(idx).and_then(|_| match idx {
69                    0 => self.smallest(),
70                    i if i == nb_elements - 1 => self.largest(),
71                    i if i > (nb_elements / 2) => self.get_from_high_end(nb_elements - idx - 1),
72                    i => {
73                        let mut my_clone = self.clone();
74                        for _ in (0..i) {
75                            my_clone.pop_smallest();
76                        }
77                        let res = my_clone.smallest();
78                        res
79                    }
80                })
81            }
82
83            pub fn get_from_high_end(&self, idx: u8) -> Option<u8> {
84                let nb_elements = self.nb_elements();
85                self.get_check(idx).and_then(|_| match idx {
86                    0 => self.largest(),
87                    i if i == nb_elements - 1 => self.smallest(),
88                    i if i > (nb_elements / 2) => self.get_from_low_end(nb_elements - idx - 1),
89                    i => {
90                        let mut my_clone = self.clone();
91                        for _ in (0..i) {
92                            my_clone.pop_largest();
93                        }
94                        let res = my_clone.largest();
95                        res
96                    }
97                })
98            }
99
100            fn get_check(&self, idx: u8) -> Option<u8> {
101                if idx >= self.nb_bits {
102                    panic!(
103                        "This {} can only handle inputs upto {}",
104                        stringify!($bit_index_name),
105                        self.nb_bits
106                    );
107                }
108                if self.is_empty() || idx >= self.nb_elements() {
109                    return None;
110                }
111                Some(0)
112            }
113
114            pub fn pop(&mut self, idx: u8) -> Option<u8> {
115                let res = self.get(idx);
116                res.map(|bit_nb| self.unset_bit(bit_nb));
117                res
118            }
119
120            pub fn pop_from_low_end(&mut self, idx: u8) -> Option<u8> {
121                let res = self.get_from_low_end(idx);
122                res.map(|bit_nb| self.unset_bit(bit_nb));
123                res
124            }
125
126            pub fn pop_from_high_end(&mut self, idx: u8) -> Option<u8> {
127                let res = self.get_from_high_end(idx);
128                res.map(|bit_nb| self.unset_bit(bit_nb));
129                res
130            }
131
132            pub fn smallest(&self) -> Option<u8> {
133                if self.is_empty() {
134                    None
135                } else {
136                    Some(self.bits.trailing_zeros() as u8)
137                }
138            }
139
140            pub fn pop_smallest(&mut self) -> Option<u8> {
141                let res = self.smallest();
142                res.map(|bit_nb| self.unset_bit(bit_nb));
143                res
144            }
145
146            pub fn largest(&self) -> Option<u8> {
147                if self.is_empty() {
148                    None
149                } else {
150                    Some((Self::SIZE as u8) - self.bits.leading_zeros() as u8 - 1)
151                }
152            }
153
154            pub fn pop_largest(&mut self) -> Option<u8> {
155                let res = self.largest();
156                res.map(|bit_nb| self.unset_bit(bit_nb));
157                res
158            }
159
160            // explicit check not necessary: handled by `single_bit`
161            #[inline]
162            pub fn set_bit(&mut self, bit_nb: u8) {
163                self.bits |= self.single_bit(bit_nb);
164            }
165
166            // explicit check not necessary: handled by `all_but_single_bit`
167            #[inline]
168            pub fn unset_bit(&mut self, bit_nb: u8) {
169                self.bits &= self.all_but_single_bit(bit_nb);
170            }
171
172            pub fn add(&mut self, bits: $bit_index_type) {
173                self.bits |= bits
174            }
175
176            pub fn absorb(&mut self, other: $bit_index_name) {
177                self.add(other.bits);
178                self.nb_bits = max(self.nb_bits, other.nb_bits);
179            }
180
181            #[inline]
182            fn single_bit(&self, bit_nb: u8) -> $bit_index_type {
183                self.check_input(bit_nb);
184                1 << bit_nb
185            }
186
187            // explicit check not necessary: handled by `single_bit`
188            #[inline]
189            fn all_but_single_bit(&self, bit_nb: u8) -> $bit_index_type {
190                <$bit_index_type>::MAX ^ self.single_bit(bit_nb)
191            }
192
193            #[inline]
194            fn check_input(&self, i: u8) {
195                if i >= self.nb_bits {
196                    panic!(
197                        "This {} can only handle inputs upto {}",
198                        stringify!($bit_index_name),
199                        self.nb_bits
200                    )
201                }
202            }
203
204            #[inline]
205            fn init(nb_bits: u8) -> $bit_index_type {
206                if nb_bits == Self::SIZE {
207                    <$bit_index_type>::MAX
208                } else {
209                    (1 << nb_bits) - 1
210                }
211            }
212        }
213
214        impl Debug for $bit_index_name {
215            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
216                writeln!(f, "{} {{", stringify!($bit_index_name))?;
217                writeln!(f, "    nb_bits: {}", self.nb_bits)?;
218                writeln!(f, "    bits: {:b}", self.bits)?;
219                writeln!(f, "}}")
220            }
221        }
222    };
223}
224
225impl_bit_index!(BitIndex8, u8);
226impl_bit_index!(BitIndex16, u16);
227impl_bit_index!(BitIndex32, u32);
228impl_bit_index!(BitIndex64, u64);
229impl_bit_index!(BitIndex128, u128);
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn new() {
237        let bi = BitIndex8::new(4).unwrap();
238        assert_eq!(0b1111, bi.unwrap());
239        assert!(BitIndex8::new(9).is_err());
240
241        let bi = BitIndex64::new(44).unwrap();
242        assert_eq!(0b11111111111111111111111111111111111111111111, bi.unwrap());
243        assert!(BitIndex64::new(69).is_err());
244    }
245
246    #[test]
247    fn empty() {
248        let mut bi = BitIndex8::empty(5).unwrap();
249        assert_eq!(None, bi.largest());
250        assert_eq!(0, bi.nb_elements());
251        bi.restore();
252        assert_eq!(5, bi.nb_elements());
253        assert_eq!(Some(4), bi.largest());
254    }
255
256    #[test]
257    fn nb_elements() {
258        let mut bi = BitIndex8::new(5).unwrap();
259        assert_eq!(5, bi.nb_elements());
260        bi.pop_largest();
261        assert_eq!(4, bi.nb_elements());
262    }
263
264    #[test]
265    fn set_unset_bits() {
266        let mut bi = BitIndex8::new(4).unwrap();
267        assert_eq!(0b1111, bi.unwrap());
268        bi.unset_bit(2);
269        assert_eq!(0b1011, bi.unwrap());
270        bi.unset_bit(0);
271        assert_eq!(0b1010, bi.unwrap());
272        bi.unset_bit(0);
273        assert_eq!(0b1010, bi.unwrap());
274        bi.set_bit(2);
275        assert_eq!(0b1110, bi.unwrap());
276        bi.set_bit(2);
277        assert_eq!(0b1110, bi.unwrap());
278    }
279
280    #[test]
281    #[should_panic]
282    fn set_panic() {
283        BitIndex8::new(4).unwrap().set_bit(4);
284    }
285
286    #[test]
287    #[should_panic]
288    fn unset_panic() {
289        BitIndex8::new(4).unwrap().unset_bit(4);
290    }
291
292    #[test]
293    fn smallest() {
294        let mut bi = BitIndex8::new(4).unwrap();
295        assert_eq!(Some(0), bi.smallest());
296        bi.unset_bit(0);
297        assert_eq!(Some(1), bi.smallest());
298        bi.unset_bit(3);
299        assert_eq!(Some(1), bi.smallest());
300        bi.unset_bit(1);
301        bi.unset_bit(2);
302        assert_eq!(None, bi.smallest());
303    }
304
305    #[test]
306    fn largest() {
307        let mut bi = BitIndex8::new(4).unwrap();
308        assert_eq!(Some(3), bi.largest());
309        bi.unset_bit(3);
310        assert_eq!(Some(2), bi.largest());
311        bi.unset_bit(2);
312        bi.unset_bit(1);
313        assert_eq!(Some(0), bi.largest());
314        bi.unset_bit(0);
315        assert_eq!(None, bi.largest());
316    }
317
318    #[test]
319    fn pop_smallest() {
320        let mut bi = BitIndex8::new(4).unwrap();
321        assert_eq!(Some(0), bi.pop_smallest());
322        assert_eq!(Some(1), bi.pop_smallest());
323        assert_eq!(Some(2), bi.pop_smallest());
324        assert_eq!(Some(3), bi.pop_smallest());
325        assert_eq!(None, bi.pop_smallest());
326    }
327
328    #[test]
329    fn pop_largest() {
330        let mut bi = BitIndex8::new(4).unwrap();
331        assert_eq!(Some(3), bi.pop_largest());
332        assert_eq!(Some(2), bi.pop_largest());
333        assert_eq!(Some(1), bi.pop_largest());
334        assert_eq!(Some(0), bi.pop_largest());
335        assert_eq!(None, bi.pop_largest());
336
337        let mut bi = BitIndex8::new(4).unwrap();
338        bi.unset_bit(1);
339        assert_eq!(Some(3), bi.pop_largest());
340        assert_eq!(Some(2), bi.pop_largest());
341        assert_eq!(Some(0), bi.pop_largest());
342        assert_eq!(None, bi.pop_largest());
343    }
344
345    #[test]
346    fn get() {
347        let mut bi = BitIndex8::new(4).unwrap();
348        bi.unset_bit(1);
349        assert_eq!(3, bi.nb_elements());
350        assert_eq!(Some(0), bi.get(0));
351        assert_eq!(Some(2), bi.get(1));
352        assert_eq!(Some(3), bi.get(2));
353        assert_eq!(None, bi.get(3));
354
355        let mut bi = BitIndex64::new(64).unwrap();
356        bi.unset_bit(1);
357        assert_eq!(Some(62), bi.get_from_low_end(61));
358        assert_eq!(Some(3), bi.get_from_high_end(60));
359
360        let mut bi = BitIndex8::new(8).unwrap();
361        bi.unset_bit(1);
362        assert_eq!(Some(2), bi.get_from_high_end(5));
363        assert_eq!(Some(6), bi.get_from_low_end(5));
364    }
365
366    #[test]
367    fn get_largest() {
368        let mut bi = BitIndex8::new(4).unwrap();
369        bi.unset_bit(1);
370        assert_eq!(3, bi.nb_elements());
371        assert_eq!(Some(3), bi.get_from_high_end(0));
372        assert_eq!(Some(2), bi.get_from_high_end(1));
373        assert_eq!(Some(0), bi.get_from_high_end(2));
374        assert_eq!(None, bi.get(3));
375    }
376
377    #[test]
378    #[should_panic]
379    fn get_panic() {
380        let mut bi = BitIndex8::new(4).unwrap();
381
382        assert_eq!(None, bi.get(4));
383        assert_eq!(None, bi.get(10));
384    }
385}