vers_vecs/bit_vec/fast_rs_vec/
bitset.rs

1//! Module that contains the bitset iterator over a `RsVec`.
2//! The iterator does the same as the `iter1`/`iter0` methods of `RsVec`, but it is faster for dense vectors.
3//! It only exists with the `simd` feature enabled, and since it is slower for sparse vectors,
4//! it is not used as a replacement for the `iter1`/`iter0` methods.
5
6use crate::RsVec;
7use std::mem::size_of;
8
9/// The number of bits in a `RsVec` that can be processed by AVX instructions at once.
10const VECTOR_SIZE: usize = 16;
11
12// add iterator functions to RsVec
13impl RsVec {
14    /// Get an iterator over the 0-bits in the vector.
15    /// The iterator returns the indices of the 0-bits in the vector, just as [`iter0`]
16    /// and [`select0`] do.
17    ///
18    /// This method is faster than [`iter0`] for dense vectors, but slower for sparse vectors.
19    ///
20    /// See [`BitSetIter`] for more information.
21    ///
22    /// [`iter0`]: RsVec::iter0
23    /// [`select0`]: RsVec::select0
24    /// [`BitSetIter`]: BitSetIter
25    #[must_use]
26    pub fn bit_set_iter0(&self) -> BitSetIter<'_, true> {
27        BitSetIter::new(self)
28    }
29
30    /// Get an iterator over the 1-bits in the vector.
31    /// The iterator returns the indices of the 1-bits in the vector, just as [`iter1`]
32    /// and [`select1`] do.
33    ///
34    /// This method is faster than [`iter1`] for dense vectors, but slower for sparse vectors.
35    ///
36    /// See [`BitSetIter`] for more information.
37    ///
38    /// [`iter1`]: RsVec::iter1
39    /// [`select1`]: RsVec::select1
40    /// [`BitSetIter`]: BitSetIter
41    #[must_use]
42    pub fn bit_set_iter1(&self) -> BitSetIter<'_, false> {
43        BitSetIter::new(self)
44    }
45}
46
47/// An iterator that iterates over 1-bits or 0-bits and returns their indices.
48/// It uses AVX vector instructions to process 16 bits at once.
49/// It is faster than [`SelectIter`] for dense vectors.
50///
51/// This is also faster than manually calling `select` on each rank,
52/// because the select data structures are not parsed by this iterator.
53///
54/// The iterator can be constructed by calling [`bit_set_iter0`] or [`bit_set_iter1`].
55///
56/// # Example
57/// ```rust
58/// use vers_vecs::{BitVec, RsVec};
59///
60/// let mut bit_vec = BitVec::new();
61/// bit_vec.append_word(u64::MAX);
62/// bit_vec.append_word(u64::MAX);
63/// bit_vec.flip_bit(4);
64///
65/// let rs_vec = RsVec::from_bit_vec(bit_vec);
66///
67/// let mut iter = rs_vec.bit_set_iter0();
68///
69/// assert_eq!(iter.next(), Some(4));
70/// assert_eq!(iter.next(), None);
71/// ```
72///
73/// [`bit_set_iter0`]: RsVec::bit_set_iter0
74/// [`bit_set_iter1`]: RsVec::bit_set_iter1
75/// [`SelectIter`]: super::SelectIter
76pub struct BitSetIter<'a, const ZERO: bool> {
77    vec: &'a RsVec,
78    base: usize,
79    offsets: [u32; VECTOR_SIZE],
80    content_len: u8,
81    cursor: u8,
82}
83
84impl<'a, const ZERO: bool> BitSetIter<'a, ZERO> {
85    pub(super) fn new(vec: &'a RsVec) -> Self {
86        let mut iter = Self {
87            vec,
88            base: 0,
89            offsets: [0; VECTOR_SIZE],
90            content_len: 0,
91            cursor: 0,
92        };
93
94        if vec.len() > VECTOR_SIZE {
95            iter.load_chunk(vec.get_bits_unchecked(0, VECTOR_SIZE) as u16);
96        }
97
98        iter
99    }
100
101    fn load_chunk(&mut self, data: u16) {
102        use std::arch::x86_64::{__mmask16, _mm512_mask_compressstoreu_epi32, _mm512_setr_epi32};
103
104        unsafe {
105            let offsets = _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
106            assert!(VECTOR_SIZE <= size_of::<u16>() * 8, "change data types");
107            let mut mask = __mmask16::from(data);
108            if ZERO {
109                mask = !mask;
110            }
111            _mm512_mask_compressstoreu_epi32(self.offsets.as_mut_ptr() as *mut _, mask, offsets);
112            self.content_len = mask.count_ones() as u8;
113            self.cursor = 0;
114        }
115    }
116
117    fn load_next_chunk(&mut self) -> Option<()> {
118        while self.cursor == self.content_len {
119            if self.base + VECTOR_SIZE >= self.vec.len() {
120                return None;
121            }
122
123            self.base += VECTOR_SIZE;
124            let data = self.vec.get_bits_unchecked(self.base, VECTOR_SIZE) as u16;
125            self.load_chunk(data);
126        }
127        Some(())
128    }
129}
130
131impl<const ZERO: bool> Iterator for BitSetIter<'_, ZERO> {
132    type Item = usize;
133
134    fn next(&mut self) -> Option<Self::Item> {
135        if self.base >= self.vec.len() {
136            return None;
137        }
138
139        if self.cursor == self.content_len {
140            if self.load_next_chunk().is_none() {
141                if ZERO {
142                    while self.base < self.vec.len() && self.vec.get_unchecked(self.base) != 0 {
143                        self.base += 1;
144                    }
145                } else {
146                    while self.base < self.vec.len() && self.vec.get_unchecked(self.base) != 1 {
147                        self.base += 1;
148                    }
149                }
150
151                return if self.base < self.vec.len() {
152                    self.base += 1;
153                    Some(self.base - 1)
154                } else {
155                    None
156                };
157            }
158        }
159
160        let offset = self.offsets[self.cursor as usize];
161        self.cursor += 1;
162        Some(self.base + offset as usize)
163    }
164}
vers_vecs/bit_vec/fast_rs_vec/bitset.rs

vers_vecs/bit_vec/fast_rs_vec/
bitset.rs
