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 there's nothing, just return
26    let _ = it.peek()?;
27
28    for next in it.by_ref() {
29        items.push(next);
30
31        // If we've reached max items, stop collecting
32        if items.len() == max.get() {
33            break;
34        }
35    }
36    // Count remaining items
37    let remaining = it.count();
38    items.shrink_to_fit();
39    Some((items, remaining))
40}
41
42#[cfg(test)]
43mod tests {
44    use super::*;
45
46    #[test]
47    fn test_it_split() {
48        let a: &[&str] = &[];
49        for v in [0, 1, 5] {
50            let (first, rest) = iterator_split(a.iter(), v);
51            assert_eq!(first.count(), 0);
52            assert_eq!(rest.count(), 0);
53        }
54        let a = &["foo"];
55        for v in [1, 5] {
56            let (first, rest) = iterator_split(a.iter(), v);
57            assert_eq!(first.count(), 1);
58            assert_eq!(rest.count(), 0);
59        }
60        let (first, rest) = iterator_split(a.iter(), 1);
61        assert_eq!(first.count(), 1);
62        assert_eq!(rest.count(), 0);
63        let a = &["foo", "bar", "baz", "blah", "other"];
64        let (first, rest) = iterator_split(a.iter(), 2);
65        assert_eq!(first.count(), 2);
66        assert_eq!(rest.count(), 3);
67    }
68
69    #[test]
70    fn test_split_empty_iterator() {
71        let a: &[&str] = &[];
72        for v in [1, 5].into_iter().map(|v| NonZeroUsize::new(v).unwrap()) {
73            assert!(collect_until(a.iter(), v).is_none());
74        }
75    }
76
77    #[test]
78    fn test_split_nonempty_iterator() {
79        let a = &["foo"];
80
81        let Some((elts, 0)) = collect_until(a.iter(), NonZeroUsize::new(1).unwrap()) else {
82            panic!()
83        };
84        assert_eq!(elts.len(), 1);
85
86        let Some((elts, 0)) = collect_until(a.iter(), const { NonZeroUsize::new(5).unwrap() })
87        else {
88            panic!()
89        };
90        assert_eq!(elts.len(), 1);
91
92        let a = &["foo", "bar", "baz", "blah", "other"];
93        let Some((elts, 3)) = collect_until(a.iter(), const { NonZeroUsize::new(2).unwrap() })
94        else {
95            panic!()
96        };
97        assert_eq!(elts.len(), 2);
98    }
99}
bootc_internal_utils/iterators.rs

bootc_internal_utils/
iterators.rs
