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

vortex_mask/
iter_bools.rs
