vortex_mask/
iter_bools.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::iter;
5
6use crate::AllOr;
7use crate::Mask;
8
9impl Mask {
10    /// Provides a closure with an iterator over the boolean values of the mask.
11    ///
12    /// This allows us to provide different implementations of the iterator based on the underlying
13    /// representation of the mask, while avoiding a heap allocation to return a boxed iterator.
14    ///
15    /// Note that bool iteration might not be the fastest way to achieve whatever is it you're
16    /// trying to do!
17    #[inline]
18    pub fn iter_bools<F, T>(&self, mut f: F) -> T
19    where
20        F: FnMut(&mut dyn Iterator<Item = bool>) -> T,
21    {
22        match self.bit_buffer() {
23            AllOr::All => f(&mut iter::repeat_n(true, self.len())),
24            AllOr::None => f(&mut iter::repeat_n(false, self.len())),
25            AllOr::Some(buffer) => f(&mut buffer.iter()),
26        }
27    }
28}
29
30#[cfg(test)]
31mod test {
32    use itertools::Itertools;
33
34    use crate::Mask;
35
36    #[test]
37    fn iter_bools_all_true() {
38        let mask = Mask::new_true(10);
39        assert_eq!(mask.iter_bools(|iter| iter.collect_vec()), vec![true; 10]);
40    }
41
42    #[test]
43    fn iter_bools_all_false() {
44        let mask = Mask::new_false(10);
45        assert_eq!(mask.iter_bools(|iter| iter.collect_vec()), vec![false; 10]);
46    }
47
48    #[test]
49    fn iter_bools_indices() {
50        assert_eq!(
51            Mask::from_indices(5, vec![]).iter_bools(|iter| iter.collect_vec()),
52            vec![false; 5],
53        );
54        assert_eq!(
55            Mask::from_indices(5, vec![0, 1, 2, 3, 4]).iter_bools(|iter| iter.collect_vec()),
56            vec![true; 5],
57        );
58        assert_eq!(
59            Mask::from_indices(5, vec![0, 4]).iter_bools(|iter| iter.collect_vec()),
60            vec![true, false, false, false, true],
61        );
62        assert_eq!(
63            Mask::from_indices(5, vec![1, 2, 3]).iter_bools(|iter| iter.collect_vec()),
64            vec![false, true, true, true, false],
65        );
66    }
67
68    #[test]
69    fn iter_bools_slices() {
70        assert_eq!(
71            Mask::from_slices(5, vec![]).iter_bools(|iter| iter.collect_vec()),
72            vec![false; 5],
73        );
74        assert_eq!(
75            Mask::from_slices(5, vec![(0, 5)]).iter_bools(|iter| iter.collect_vec()),
76            vec![true; 5],
77        );
78        assert_eq!(
79            Mask::from_slices(5, vec![(0, 1), (4, 5)]).iter_bools(|iter| iter.collect_vec()),
80            vec![true, false, false, false, true],
81        );
82        assert_eq!(
83            Mask::from_slices(5, vec![(1, 4)]).iter_bools(|iter| iter.collect_vec()),
84            vec![false, true, true, true, false],
85        );
86    }
87}
vortex_mask/iter_bools.rs

vortex_mask/
iter_bools.rs
