bootc_internal_utils/
iterators.rs

1use std::num::NonZeroUsize;
2
3/// Given an iterator that's cloneable, split it into two iterators
4/// at a given maximum number of elements.
5pub fn iterator_split<I>(
6    it: I,
7    max: usize,
8) -> (impl Iterator<Item = I::Item>, impl Iterator<Item = I::Item>)
9where
10    I: Iterator + Clone,
11{
12    let rest = it.clone();
13    (it.take(max), rest.skip(max))
14}
15
16/// Gather the first N items, and provide the count of the remaining items.
17/// The max count cannot be zero as that's a pathological case.
18pub fn collect_until<I>(it: I, max: NonZeroUsize) -> Option<(Vec<I::Item>, usize)>
19where
20    I: Iterator,
21{
22    let mut items = Vec::with_capacity(max.get());
23
24    let mut it = it.peekable();
25    if it.peek().is_none() {
26        return None;
27    }
28
29    while let Some(next) = it.next() {
30        items.push(next);
31
32        // If we've reached max items, stop collecting
33        if items.len() == max.get() {
34            break;
35        }
36    }
37    // Count remaining items
38    let remaining = it.count();
39    items.shrink_to_fit();
40    Some((items, remaining))
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    #[test]
48    fn test_it_split() {
49        let a: &[&str] = &[];
50        for v in [0, 1, 5] {
51            let (first, rest) = iterator_split(a.iter(), v);
52            assert_eq!(first.count(), 0);
53            assert_eq!(rest.count(), 0);
54        }
55        let a = &["foo"];
56        for v in [1, 5] {
57            let (first, rest) = iterator_split(a.iter(), v);
58            assert_eq!(first.count(), 1);
59            assert_eq!(rest.count(), 0);
60        }
61        let (first, rest) = iterator_split(a.iter(), 1);
62        assert_eq!(first.count(), 1);
63        assert_eq!(rest.count(), 0);
64        let a = &["foo", "bar", "baz", "blah", "other"];
65        let (first, rest) = iterator_split(a.iter(), 2);
66        assert_eq!(first.count(), 2);
67        assert_eq!(rest.count(), 3);
68    }
69
70    #[test]
71    fn test_split_empty_iterator() {
72        let a: &[&str] = &[];
73        for v in [1, 5].into_iter().map(|v| NonZeroUsize::new(v).unwrap()) {
74            assert!(collect_until(a.iter(), v).is_none());
75        }
76    }
77
78    #[test]
79    fn test_split_nonempty_iterator() {
80        let a = &["foo"];
81
82        let Some((elts, 0)) = collect_until(a.iter(), NonZeroUsize::new(1).unwrap()) else {
83            panic!()
84        };
85        assert_eq!(elts.len(), 1);
86
87        let Some((elts, 0)) = collect_until(a.iter(), const { NonZeroUsize::new(5).unwrap() })
88        else {
89            panic!()
90        };
91        assert_eq!(elts.len(), 1);
92
93        let a = &["foo", "bar", "baz", "blah", "other"];
94        let Some((elts, 3)) = collect_until(a.iter(), const { NonZeroUsize::new(2).unwrap() })
95        else {
96            panic!()
97        };
98        assert_eq!(elts.len(), 2);
99    }
100}
bootc_internal_utils/iterators.rs

bootc_internal_utils/
iterators.rs
