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    pub fn iter_bools<F, T>(&self, mut f: F) -> T
17    where
18        F: FnMut(&mut dyn Iterator<Item = bool>) -> T,
19    {
20        match self.boolean_buffer() {
21            AllOr::All => f(&mut iter::repeat_n(true, self.len())),
22            AllOr::None => f(&mut iter::repeat_n(false, self.len())),
23            AllOr::Some(buffer) => f(&mut buffer.iter()),
24        }
25    }
26}
27
28#[cfg(test)]
29mod test {
30    use itertools::Itertools;
31
32    use crate::Mask;
33
34    #[test]
35    fn iter_bools_all_true() {
36        let mask = Mask::new_true(10);
37        assert_eq!(mask.iter_bools(|iter| iter.collect_vec()), vec![true; 10]);
38    }
39
40    #[test]
41    fn iter_bools_all_false() {
42        let mask = Mask::new_false(10);
43        assert_eq!(mask.iter_bools(|iter| iter.collect_vec()), vec![false; 10]);
44    }
45
46    #[test]
47    fn iter_bools_indices() {
48        assert_eq!(
49            Mask::from_indices(5, vec![]).iter_bools(|iter| iter.collect_vec()),
50            vec![false; 5],
51        );
52        assert_eq!(
53            Mask::from_indices(5, vec![0, 1, 2, 3, 4]).iter_bools(|iter| iter.collect_vec()),
54            vec![true; 5],
55        );
56        assert_eq!(
57            Mask::from_indices(5, vec![0, 4]).iter_bools(|iter| iter.collect_vec()),
58            vec![true, false, false, false, true],
59        );
60        assert_eq!(
61            Mask::from_indices(5, vec![1, 2, 3]).iter_bools(|iter| iter.collect_vec()),
62            vec![false, true, true, true, false],
63        );
64    }
65
66    #[test]
67    fn iter_bools_slices() {
68        assert_eq!(
69            Mask::from_slices(5, vec![]).iter_bools(|iter| iter.collect_vec()),
70            vec![false; 5],
71        );
72        assert_eq!(
73            Mask::from_slices(5, vec![(0, 5)]).iter_bools(|iter| iter.collect_vec()),
74            vec![true; 5],
75        );
76        assert_eq!(
77            Mask::from_slices(5, vec![(0, 1), (4, 5)]).iter_bools(|iter| iter.collect_vec()),
78            vec![true, false, false, false, true],
79        );
80        assert_eq!(
81            Mask::from_slices(5, vec![(1, 4)]).iter_bools(|iter| iter.collect_vec()),
82            vec![false, true, true, true, false],
83        );
84    }
85}
vortex_mask/iter_bools.rs

vortex_mask/
iter_bools.rs
