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 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
8 pub struct $bit_index_name {
9 bits: $bit_index_type,
11 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 #[inline]
162 pub fn set_bit(&mut self, bit_nb: u8) {
163 self.bits |= self.single_bit(bit_nb);
164 }
165
166 #[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 #[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}