bitmaps/
bitmap.rs

Help
1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5use core::ops::*;
6
7use typenum::*;
8
9use crate::types::{BitOps, Bits};
10
11#[cfg(feature = "std")]
12use std::fmt::{Debug, Error, Formatter};
13
14/// A compact array of bits.
15///
16/// The type used to store the bitmap will be the minimum unsigned integer type
17/// required to fit the number of bits, from `u8` to `u128`. If the size is 1,
18/// `bool` is used. If the size exceeds 128, an array of `u128` will be used,
19/// sized as appropriately. The maximum supported size is currently 1024,
20/// represented by an array `[u128; 8]`.
21pub struct Bitmap<Size: Bits> {
22    pub(crate) data: Size::Store,
23}
24
25impl<Size: Bits> Clone for Bitmap<Size> {
26    fn clone(&self) -> Self {
27        Bitmap { data: self.data }
28    }
29}
30
31impl<Size: Bits> Copy for Bitmap<Size> {}
32
33impl<Size: Bits> Default for Bitmap<Size> {
34    fn default() -> Self {
35        Bitmap {
36            data: Size::Store::default(),
37        }
38    }
39}
40
41impl<Size: Bits> PartialEq for Bitmap<Size> {
42    fn eq(&self, other: &Self) -> bool {
43        self.data == other.data
44    }
45}
46
47#[cfg(feature = "std")]
48impl<Size: Bits> Debug for Bitmap<Size> {
49    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
50        write!(f, "{}", Size::Store::to_hex(&self.data))
51    }
52}
53
54impl<Size: Bits> Bitmap<Size> {
55    /// Construct a bitmap with every bit set to `false`.
56    #[inline]
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    /// Construct a bitmap where every bit with index less than `bits` is
62    /// `true`, and every other bit is `false`.
63    #[inline]
64    pub fn mask(bits: usize) -> Self {
65        debug_assert!(bits < Size::USIZE);
66        Self {
67            data: Size::Store::make_mask(bits),
68        }
69    }
70
71    /// Construct a bitmap from a value of the same type as its backing store.
72    #[inline]
73    pub fn from_value(data: Size::Store) -> Self {
74        Self { data }
75    }
76
77    /// Convert this bitmap into a value of the type of its backing store.
78    #[inline]
79    pub fn into_value(self) -> Size::Store {
80        self.data
81    }
82
83    /// Count the number of `true` bits in the bitmap.
84    #[inline]
85    pub fn len(self) -> usize {
86        Size::Store::len(&self.data)
87    }
88
89    /// Test if the bitmap contains only `false` bits.
90    #[inline]
91    pub fn is_empty(self) -> bool {
92        self.first_index().is_none()
93    }
94
95    /// Get the value of the bit at a given index.
96    #[inline]
97    pub fn get(self, index: usize) -> bool {
98        debug_assert!(index < Size::USIZE);
99        Size::Store::get(&self.data, index)
100    }
101
102    /// Set the value of the bit at a given index.
103    ///
104    /// Returns the previous value of the bit.
105    #[inline]
106    pub fn set(&mut self, index: usize, value: bool) -> bool {
107        debug_assert!(index < Size::USIZE);
108        Size::Store::set(&mut self.data, index, value)
109    }
110
111    /// Find the index of the first `true` bit in the bitmap.
112    #[inline]
113    pub fn first_index(self) -> Option<usize> {
114        Size::Store::first_index(&self.data)
115    }
116
117    /// Invert all the bits in the bitmap.
118    #[inline]
119    pub fn invert(&mut self) {
120        Size::Store::invert(&mut self.data);
121    }
122}
123
124impl<'a, Size: Bits> IntoIterator for &'a Bitmap<Size> {
125    type Item = usize;
126    type IntoIter = Iter<'a, Size>;
127
128    fn into_iter(self) -> Self::IntoIter {
129        Iter {
130            index: 0,
131            data: self,
132        }
133    }
134}
135
136impl<Size: Bits> BitAnd for Bitmap<Size> {
137    type Output = Self;
138    fn bitand(mut self, rhs: Self) -> Self::Output {
139        Size::Store::bit_and(&mut self.data, &rhs.data);
140        self
141    }
142}
143
144impl<Size: Bits> BitOr for Bitmap<Size> {
145    type Output = Self;
146    fn bitor(mut self, rhs: Self) -> Self::Output {
147        Size::Store::bit_or(&mut self.data, &rhs.data);
148        self
149    }
150}
151
152impl<Size: Bits> BitXor for Bitmap<Size> {
153    type Output = Self;
154    fn bitxor(mut self, rhs: Self) -> Self::Output {
155        Size::Store::bit_xor(&mut self.data, &rhs.data);
156        self
157    }
158}
159
160impl<Size: Bits> Not for Bitmap<Size> {
161    type Output = Self;
162    fn not(mut self) -> Self::Output {
163        Size::Store::invert(&mut self.data);
164        self
165    }
166}
167
168impl<Size: Bits> BitAndAssign for Bitmap<Size> {
169    fn bitand_assign(&mut self, rhs: Self) {
170        Size::Store::bit_and(&mut self.data, &rhs.data);
171    }
172}
173
174impl<Size: Bits> BitOrAssign for Bitmap<Size> {
175    fn bitor_assign(&mut self, rhs: Self) {
176        Size::Store::bit_or(&mut self.data, &rhs.data);
177    }
178}
179
180impl<Size: Bits> BitXorAssign for Bitmap<Size> {
181    fn bitxor_assign(&mut self, rhs: Self) {
182        Size::Store::bit_xor(&mut self.data, &rhs.data);
183    }
184}
185
186impl From<[u128; 2]> for Bitmap<U256> {
187    fn from(data: [u128; 2]) -> Self {
188        Bitmap { data }
189    }
190}
191
192impl From<[u128; 3]> for Bitmap<U384> {
193    fn from(data: [u128; 3]) -> Self {
194        Bitmap { data }
195    }
196}
197
198impl From<[u128; 4]> for Bitmap<U512> {
199    fn from(data: [u128; 4]) -> Self {
200        Bitmap { data }
201    }
202}
203
204impl From<[u128; 5]> for Bitmap<U640> {
205    fn from(data: [u128; 5]) -> Self {
206        Bitmap { data }
207    }
208}
209
210impl From<[u128; 6]> for Bitmap<U768> {
211    fn from(data: [u128; 6]) -> Self {
212        Bitmap { data }
213    }
214}
215
216impl From<[u128; 7]> for Bitmap<U896> {
217    fn from(data: [u128; 7]) -> Self {
218        Bitmap { data }
219    }
220}
221
222impl From<[u128; 8]> for Bitmap<U1024> {
223    fn from(data: [u128; 8]) -> Self {
224        Bitmap { data }
225    }
226}
227
228impl Into<[u128; 2]> for Bitmap<U256> {
229    fn into(self) -> [u128; 2] {
230        self.data
231    }
232}
233
234impl Into<[u128; 3]> for Bitmap<U384> {
235    fn into(self) -> [u128; 3] {
236        self.data
237    }
238}
239
240impl Into<[u128; 4]> for Bitmap<U512> {
241    fn into(self) -> [u128; 4] {
242        self.data
243    }
244}
245
246impl Into<[u128; 5]> for Bitmap<U640> {
247    fn into(self) -> [u128; 5] {
248        self.data
249    }
250}
251
252impl Into<[u128; 6]> for Bitmap<U768> {
253    fn into(self) -> [u128; 6] {
254        self.data
255    }
256}
257
258impl Into<[u128; 7]> for Bitmap<U896> {
259    fn into(self) -> [u128; 7] {
260        self.data
261    }
262}
263
264impl Into<[u128; 8]> for Bitmap<U1024> {
265    fn into(self) -> [u128; 8] {
266        self.data
267    }
268}
269
270/// An iterator over the indices in a bitmap which are `true`.
271///
272/// This yields a sequence of `usize` indices, not their contents (which are
273/// always `true` anyway, by definition).
274///
275/// # Examples
276///
277/// ```rust
278/// # use bitmaps::Bitmap;
279/// # use typenum::U10;
280/// let mut bitmap: Bitmap<U10> = Bitmap::new();
281/// bitmap.set(3, true);
282/// bitmap.set(5, true);
283/// bitmap.set(8, true);
284/// let true_indices: Vec<usize> = bitmap.into_iter().collect();
285/// assert_eq!(vec![3, 5, 8], true_indices);
286/// ```
287pub struct Iter<'a, Size: Bits> {
288    index: usize,
289    data: &'a Bitmap<Size>,
290}
291
292impl<'a, Size: Bits> Iterator for Iter<'a, Size> {
293    type Item = usize;
294
295    fn next(&mut self) -> Option<Self::Item> {
296        if self.index >= Size::USIZE {
297            return None;
298        }
299        if self.data.get(self.index) {
300            self.index += 1;
301            Some(self.index - 1)
302        } else {
303            self.index += 1;
304            self.next()
305        }
306    }
307}
308
309#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
310#[allow(clippy::cast_ptr_alignment)]
311mod x86_arch {
312    use super::*;
313    #[cfg(target_arch = "x86")]
314    use core::arch::x86::*;
315    #[cfg(target_arch = "x86_64")]
316    use core::arch::x86_64::*;
317
318    impl Bitmap<U128> {
319        #[target_feature(enable = "sse2")]
320        pub unsafe fn load_m128i(&self) -> __m128i {
321            _mm_loadu_si128(&self.data as *const _ as *const __m128i)
322        }
323    }
324
325    impl Bitmap<U256> {
326        #[target_feature(enable = "sse2")]
327        pub unsafe fn load_m128i(&self) -> [__m128i; 2] {
328            let ptr = &self.data as *const _ as *const __m128i;
329            [_mm_loadu_si128(ptr), _mm_loadu_si128(ptr.add(1))]
330        }
331
332        #[target_feature(enable = "avx")]
333        pub unsafe fn load_m256i(&self) -> __m256i {
334            _mm256_loadu_si256(&self.data as *const _ as *const __m256i)
335        }
336    }
337
338    impl Bitmap<U512> {
339        #[target_feature(enable = "sse2")]
340        pub unsafe fn load_m128i(&self) -> [__m128i; 4] {
341            let ptr = &self.data as *const _ as *const __m128i;
342            [
343                _mm_loadu_si128(ptr),
344                _mm_loadu_si128(ptr.add(1)),
345                _mm_loadu_si128(ptr.add(2)),
346                _mm_loadu_si128(ptr.add(3)),
347            ]
348        }
349
350        #[target_feature(enable = "avx")]
351        pub unsafe fn load_m256i(&self) -> [__m256i; 2] {
352            let ptr = &self.data as *const _ as *const __m256i;
353            [_mm256_loadu_si256(ptr), _mm256_loadu_si256(ptr.add(1))]
354        }
355    }
356
357    impl Bitmap<U768> {
358        #[target_feature(enable = "sse2")]
359        pub unsafe fn load_m128i(&self) -> [__m128i; 6] {
360            let ptr = &self.data as *const _ as *const __m128i;
361            [
362                _mm_loadu_si128(ptr),
363                _mm_loadu_si128(ptr.add(1)),
364                _mm_loadu_si128(ptr.add(2)),
365                _mm_loadu_si128(ptr.add(3)),
366                _mm_loadu_si128(ptr.add(4)),
367                _mm_loadu_si128(ptr.add(5)),
368            ]
369        }
370
371        #[target_feature(enable = "avx")]
372        pub unsafe fn load_m256i(&self) -> [__m256i; 3] {
373            let ptr = &self.data as *const _ as *const __m256i;
374            [
375                _mm256_loadu_si256(ptr),
376                _mm256_loadu_si256(ptr.add(1)),
377                _mm256_loadu_si256(ptr.add(2)),
378            ]
379        }
380    }
381
382    impl Bitmap<U1024> {
383        #[target_feature(enable = "sse2")]
384        pub unsafe fn load_m128i(&self) -> [__m128i; 8] {
385            let ptr = &self.data as *const _ as *const __m128i;
386            [
387                _mm_loadu_si128(ptr),
388                _mm_loadu_si128(ptr.add(1)),
389                _mm_loadu_si128(ptr.add(2)),
390                _mm_loadu_si128(ptr.add(3)),
391                _mm_loadu_si128(ptr.add(4)),
392                _mm_loadu_si128(ptr.add(5)),
393                _mm_loadu_si128(ptr.add(6)),
394                _mm_loadu_si128(ptr.add(7)),
395            ]
396        }
397
398        #[target_feature(enable = "avx")]
399        pub unsafe fn load_m256i(&self) -> [__m256i; 4] {
400            let ptr = &self.data as *const _ as *const __m256i;
401            [
402                _mm256_loadu_si256(ptr),
403                _mm256_loadu_si256(ptr.add(1)),
404                _mm256_loadu_si256(ptr.add(2)),
405                _mm256_loadu_si256(ptr.add(3)),
406            ]
407        }
408    }
409
410    impl From<__m128i> for Bitmap<U128> {
411        fn from(data: __m128i) -> Self {
412            Self {
413                data: unsafe { core::mem::transmute(data) },
414            }
415        }
416    }
417
418    impl From<__m256i> for Bitmap<U256> {
419        fn from(data: __m256i) -> Self {
420            Self {
421                data: unsafe { core::mem::transmute(data) },
422            }
423        }
424    }
425
426    impl Into<__m128i> for Bitmap<U128> {
427        fn into(self) -> __m128i {
428            unsafe { self.load_m128i() }
429        }
430    }
431
432    impl Into<__m256i> for Bitmap<U256> {
433        fn into(self) -> __m256i {
434            unsafe { self.load_m256i() }
435        }
436    }
437
438    #[cfg(test)]
439    mod test {
440        use super::*;
441
442        struct AlignmentTester<B>
443        where
444            B: Bits,
445        {
446            _byte: u8,
447            bits: Bitmap<B>,
448        }
449
450        #[test]
451        fn load_128() {
452            let mut t: AlignmentTester<U128> = AlignmentTester {
453                _byte: 0,
454                bits: Bitmap::new(),
455            };
456            t.bits.set(5, true);
457            let m = unsafe { t.bits.load_m128i() };
458            let mut bits: Bitmap<U128> = m.into();
459            assert!(bits.set(5, false));
460            assert!(bits.is_empty());
461        }
462
463        #[test]
464        fn load_256() {
465            let mut t: AlignmentTester<U256> = AlignmentTester {
466                _byte: 0,
467                bits: Bitmap::new(),
468            };
469            t.bits.set(5, true);
470            let m = unsafe { t.bits.load_m256i() };
471            let mut bits: Bitmap<U256> = m.into();
472            assert!(bits.set(5, false));
473            assert!(bits.is_empty());
474        }
475    }
476}
477
478#[cfg(test)]
479mod test {
480    use super::*;
481    use proptest::collection::btree_set;
482    use proptest::proptest;
483    use typenum::{U1024, U64};
484
485    proptest! {
486        #[test]
487        fn get_set_and_iter_64(bits in btree_set(0..64usize, 0..64)) {
488            let mut bitmap = Bitmap::<U64>::new();
489            for i in &bits {
490                bitmap.set(*i, true);
491            }
492            for i in 0..64 {
493                assert_eq!(bitmap.get(i), bits.contains(&i));
494            }
495            assert!(bitmap.into_iter().eq(bits.into_iter()));
496        }
497
498        #[test]
499        fn get_set_and_iter_1024(bits in btree_set(0..1024usize, 0..1024)) {
500            let mut bitmap = Bitmap::<U1024>::new();
501            for i in &bits {
502                bitmap.set(*i, true);
503            }
504            for i in 0..1024 {
505                assert_eq!(bitmap.get(i), bits.contains(&i));
506            }
507            assert!(bitmap.into_iter().eq(bits.into_iter()));
508        }
509    }
510}
bitmaps/bitmap.rs

bitmaps/
bitmap.rs
