bitset_fixed/
lib.rs

1#![allow(clippy::suspicious_op_assign_impl)]
2
3const TRUE: &bool = &true;
4const FALSE: &bool = &false;
5
6#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
7/// Fixed sized bitset
8pub struct BitSet {
9    buf: Vec<u64>,
10    size: usize,
11}
12
13impl BitSet {
14    /// Construct a new, zero bitset with specified capacity.
15    /// This method allocates O(size) bits
16    pub fn new(size: usize) -> BitSet {
17        BitSet {
18            buf: vec![0; (size + 63) / 64],
19            size,
20        }
21    }
22
23    /// Set i-th bit to `b`
24    #[inline]
25    pub fn set(&mut self, i: usize, b: bool) {
26        assert!(i < self.size);
27        if b {
28            self.buf[i >> 6] |= 1 << (i & 63);
29        } else {
30            self.buf[i >> 6] &= !(1 << (i & 63));
31        }
32    }
33
34    /// Get the size of bits
35    #[inline]
36    pub fn size(&self) -> usize {
37        self.size
38    }
39
40    /// Get the number of ones
41    #[inline]
42    pub fn count_ones(&self) -> u32 {
43        self.buf.iter().map(|x| x.count_ones()).sum()
44    }
45
46    /// Get the number of zeros
47    #[inline]
48    pub fn count_zeros(&self) -> u32 {
49        self.size as u32 - self.count_ones()
50    }
51
52    /// Faster left shift and or
53    ///
54    /// `bitset | (bitset << x)`
55    #[inline]
56    pub fn shl_or(&mut self, x: usize) {
57        let q = x >> 6;
58        let r = x & 63;
59
60        if q >= self.buf.len() {
61            return;
62        }
63
64        if r == 0 {
65            for i in (q..self.buf.len()).rev() {
66                *unsafe { self.buf.get_unchecked_mut(i) } |=
67                    *unsafe { self.buf.get_unchecked(i - q) };
68            }
69        } else {
70            for i in (q + 1..self.buf.len()).rev() {
71                *unsafe { self.buf.get_unchecked_mut(i) } |=
72                    (unsafe { self.buf.get_unchecked(i - q) } << r)
73                        | (unsafe { self.buf.get_unchecked(i - q - 1) } >> (64 - r));
74            }
75            *unsafe { self.buf.get_unchecked_mut(q) } |= unsafe { self.buf.get_unchecked(0) } << r;
76        }
77
78        self.chomp();
79    }
80
81    /// Get inner buffer
82    #[inline]
83    pub fn buffer(&self) -> &[u64] {
84        &self.buf
85    }
86
87    /// Get inner buffer with mutable reference
88    #[inline]
89    pub fn buffer_mut(&mut self) -> &mut [u64] {
90        &mut self.buf
91    }
92
93    /// Set tailing bits in inner buffer whose index are greater than size to `0`
94    #[inline]
95    pub fn chomp(&mut self) {
96        let r = self.size & 63;
97        if r != 0 {
98            if let Some(x) = self.buf.last_mut() {
99                let d = 64 - r;
100                *x = (*x << d) >> d;
101            }
102        }
103    }
104}
105
106impl std::ops::Index<usize> for BitSet {
107    type Output = bool;
108    #[inline]
109    fn index(&self, index: usize) -> &bool {
110        assert!(index < self.size);
111        [FALSE, TRUE][(self.buf[index >> 6] >> (index & 63)) as usize & 1]
112    }
113}
114
115impl std::ops::ShlAssign<usize> for BitSet {
116    #[inline]
117    fn shl_assign(&mut self, x: usize) {
118        let q = x >> 6;
119        let r = x & 63;
120
121        if q >= self.buf.len() {
122            for x in &mut self.buf {
123                *x = 0;
124            }
125            return;
126        }
127
128        if r == 0 {
129            for i in (q..self.buf.len()).rev() {
130                *unsafe { self.buf.get_unchecked_mut(i) } =
131                    *unsafe { self.buf.get_unchecked(i - q) };
132            }
133        } else {
134            for i in (q + 1..self.buf.len()).rev() {
135                *unsafe { self.buf.get_unchecked_mut(i) } =
136                    (unsafe { self.buf.get_unchecked(i - q) } << r)
137                        | (unsafe { self.buf.get_unchecked(i - q - 1) } >> (64 - r));
138            }
139            *unsafe { self.buf.get_unchecked_mut(q) } = unsafe { self.buf.get_unchecked(0) } << r;
140        }
141
142        for x in &mut self.buf[..q] {
143            *x = 0;
144        }
145
146        self.chomp();
147    }
148}
149
150impl std::ops::Shl<usize> for BitSet {
151    type Output = Self;
152
153    #[inline]
154    fn shl(mut self, x: usize) -> Self {
155        self <<= x;
156        self
157    }
158}
159
160impl<'a> std::ops::Shl<usize> for &'a BitSet {
161    type Output = BitSet;
162    #[inline]
163    fn shl(self, x: usize) -> Self::Output {
164        let mut result = self.clone();
165        result <<= x;
166        result
167    }
168}
169
170impl std::ops::ShrAssign<usize> for BitSet {
171    #[inline]
172    fn shr_assign(&mut self, x: usize) {
173        let q = x >> 6;
174        let r = x & 63;
175
176        if q >= self.buf.len() {
177            for x in &mut self.buf {
178                *x = 0;
179            }
180            return;
181        }
182
183        if r == 0 {
184            for i in 0..self.buf.len() - q {
185                *unsafe { self.buf.get_unchecked_mut(i) } =
186                    *unsafe { self.buf.get_unchecked(i + q) };
187            }
188        } else {
189            for i in 0..self.buf.len() - q - 1 {
190                *unsafe { self.buf.get_unchecked_mut(i) } =
191                    (unsafe { self.buf.get_unchecked(i + q) } >> r)
192                        | (unsafe { self.buf.get_unchecked(i + q + 1) } << (64 - r));
193            }
194            let len = self.buf.len();
195            *unsafe { self.buf.get_unchecked_mut(len - q - 1) } =
196                unsafe { self.buf.get_unchecked(len - 1) } >> r;
197        }
198
199        let len = self.buf.len();
200        for x in &mut self.buf[len - q..] {
201            *x = 0;
202        }
203    }
204}
205
206impl std::ops::Shr<usize> for BitSet {
207    type Output = Self;
208
209    #[inline]
210    fn shr(mut self, x: usize) -> Self {
211        self >>= x;
212        self
213    }
214}
215
216impl<'a> std::ops::Shr<usize> for &'a BitSet {
217    type Output = BitSet;
218
219    #[inline]
220    fn shr(self, x: usize) -> Self::Output {
221        let mut result = self.clone();
222        result >>= x;
223        result
224    }
225}
226
227impl<'a> std::ops::BitAndAssign<&'a BitSet> for BitSet {
228    #[inline]
229    fn bitand_assign(&mut self, rhs: &'a Self) {
230        for (a, b) in self.buf.iter_mut().zip(rhs.buf.iter()) {
231            *a &= *b;
232        }
233    }
234}
235
236impl<'a> std::ops::BitAnd<&'a BitSet> for BitSet {
237    type Output = Self;
238    #[inline]
239    fn bitand(mut self, rhs: &'a Self) -> Self::Output {
240        self &= rhs;
241        self
242    }
243}
244
245impl<'a, 'b> std::ops::BitAnd<&'b BitSet> for &'a BitSet {
246    type Output = BitSet;
247    #[inline]
248    fn bitand(self, rhs: &'b BitSet) -> Self::Output {
249        let mut result = self.clone();
250        result &= rhs;
251        result
252    }
253}
254
255impl<'a> std::ops::BitOrAssign<&'a BitSet> for BitSet {
256    #[inline]
257    fn bitor_assign(&mut self, rhs: &'a Self) {
258        for (a, b) in self.buf.iter_mut().zip(rhs.buf.iter()) {
259            *a |= *b;
260        }
261        self.chomp();
262    }
263}
264
265impl<'a> std::ops::BitOr<&'a BitSet> for BitSet {
266    type Output = Self;
267    #[inline]
268    fn bitor(mut self, rhs: &'a Self) -> Self::Output {
269        self |= rhs;
270        self
271    }
272}
273
274impl<'a, 'b> std::ops::BitOr<&'b BitSet> for &'a BitSet {
275    type Output = BitSet;
276    #[inline]
277    fn bitor(self, rhs: &'b BitSet) -> Self::Output {
278        let mut result = self.clone();
279        result |= rhs;
280        result
281    }
282}
283
284impl<'a> std::ops::BitXorAssign<&'a BitSet> for BitSet {
285    #[inline]
286    fn bitxor_assign(&mut self, rhs: &'a Self) {
287        for (a, b) in self.buf.iter_mut().zip(rhs.buf.iter()) {
288            *a ^= *b;
289        }
290        self.chomp();
291    }
292}
293
294impl<'a> std::ops::BitXor<&'a BitSet> for BitSet {
295    type Output = Self;
296    #[inline]
297    fn bitxor(mut self, rhs: &'a Self) -> Self {
298        self ^= rhs;
299        self
300    }
301}
302
303impl<'a, 'b> std::ops::BitXor<&'b BitSet> for &'a BitSet {
304    type Output = BitSet;
305    #[inline]
306    fn bitxor(self, rhs: &'b BitSet) -> Self::Output {
307        let mut result = self.clone();
308        result ^= rhs;
309        result
310    }
311}
312
313impl std::ops::Not for BitSet {
314    type Output = Self;
315    #[inline]
316    fn not(mut self) -> Self::Output {
317        for x in &mut self.buf {
318            *x = !*x;
319        }
320        self.chomp();
321        self
322    }
323}
324
325impl<'a> std::ops::Not for &'a BitSet {
326    type Output = BitSet;
327    #[inline]
328    fn not(self) -> Self::Output {
329        !self.clone()
330    }
331}
332
333#[test]
334fn test_bitset_set_read() {
335    use rand::prelude::*;
336    let size = 6400;
337    let mut set = BitSet::new(size);
338    let mut v = vec![false; size];
339    let mut rng = StdRng::seed_from_u64(114514);
340
341    for i in 0..size {
342        let b = rng.next_u32() % 2 == 0;
343        set.set(i, b);
344        v[i] = b;
345    }
346
347    for i in 0..size {
348        assert_eq!(set[i], v[i]);
349    }
350}
351
352#[test]
353fn test_bitset_shl() {
354    let do_test = |size, shift| {
355        use rand::prelude::*;
356        let mut set = BitSet::new(size);
357        let mut v = vec![false; size];
358        let mut rng = StdRng::seed_from_u64(114514);
359
360        for i in 0..size {
361            let b = rng.next_u32() % 2 == 0;
362            set.set(i, b);
363            v[i] = b;
364        }
365        for i in (shift..v.len()).rev() {
366            v[i] = v[i - shift];
367        }
368        for i in 0..std::cmp::min(size, shift) {
369            v[i] = false;
370        }
371
372        set <<= shift;
373        for i in 0..size {
374            assert_eq!(set[i], v[i]);
375        }
376    };
377
378    do_test(6400, 640);
379    do_test(6400, 114);
380    do_test(6400, 514);
381    do_test(6400, 6400);
382    do_test(6400, 16400);
383
384    let mut t = BitSet::new(310);
385
386    for i in 0..31000 {
387        t <<= i;
388    }
389}
390
391#[test]
392fn test_bitset_shr() {
393    let do_test = |size, shift| {
394        use rand::prelude::*;
395        let mut set = BitSet::new(size);
396        let mut v = vec![false; size];
397        let mut rng = StdRng::seed_from_u64(114514);
398
399        for i in 0..size {
400            let b = rng.next_u32() % 2 == 0;
401            set.set(i, b);
402            v[i] = b;
403        }
404
405        let s = if size >= shift { size - shift } else { 0 };
406
407        for i in 0..s {
408            v[i] = v[i + shift];
409        }
410
411        for i in s..size {
412            v[i] = false;
413        }
414
415        set >>= shift;
416        for i in 0..size {
417            assert_eq!(set[i], v[i]);
418        }
419    };
420
421    do_test(6400, 640);
422    do_test(6400, 114);
423    do_test(6400, 514);
424    do_test(63, 65);
425    do_test(6400, 6400);
426    do_test(6400, 16400);
427
428    let mut t = BitSet::new(310);
429
430    for i in 0..31000 {
431        t >>= i;
432    }
433}
434
435#[test]
436fn test_bitset_chomp() {
437    let mut set1 = BitSet::new(4);
438    let mut set2 = BitSet::new(8);
439
440    for i in 0..4 {
441        set1.set(i, true);
442        set2.set(i, true);
443    }
444
445    for i in 4..8 {
446        set2.set(i, true);
447    }
448
449    set1 <<= 2;
450    assert_eq!(set1.count_ones(), 2);
451    assert_eq!(set1.count_zeros(), 2);
452    assert_eq!((&set1 | &set2).count_ones(), 4);
453    assert_eq!((&set1 & &set2).count_ones(), 2);
454    assert_eq!((&set1 ^ &set2).count_ones(), 2);
455}
bitset_fixed/lib.rs

bitset_fixed/
lib.rs
