utils: Add a method to split an iterator

Signed-off-by: Colin Walters <[email protected]>

Author: cgwalters

utils/src/lib.rs

@@ -6,3 +6,97 @@ mod command;
 mod tracing_util;
 pub use command::*;
 pub use tracing_util::*;
+
+/// Given an iterator that's clonable, split it into two iterators
+/// at a given maximum number of elements.
+pub fn iterator_split<I>(
+    it: I,
+    max: usize,
+) -> (impl Iterator<Item = I::Item>, impl Iterator<Item = I::Item>)
+where
+    I: Iterator + Clone,
+{
+    let rest = it.clone();
+    (it.take(max), rest.skip(max))
+}
+
+/// Given an iterator that's clonable, split off the first N elements
+/// at the pivot point. If the iterator would be empty, return None.
+/// Return the count of the remainder.
+pub fn iterator_split_nonempty_rest_count<I>(
+    it: I,
+    max: usize,
+) -> Option<(impl Iterator<Item = I::Item>, usize)>
+where
+    I: Iterator + Clone,
+{
+    let rest = it.clone();
+    let mut it = it.peekable();
+    if it.peek().is_some() {
+        Some((it.take(max), rest.skip(max).count()))
+    } else {
+        None
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_it_split() {
+        let a: &[&str] = &[];
+        for v in [0, 1, 5] {
+            let (first, rest) = iterator_split(a.iter(), v);
+            assert_eq!(first.count(), 0);
+            assert_eq!(rest.count(), 0);
+        }
+        let a = &["foo"];
+        for v in [1, 5] {
+            let (first, rest) = iterator_split(a.iter(), v);
+            assert_eq!(first.count(), 1);
+            assert_eq!(rest.count(), 0);
+        }
+        let (first, rest) = iterator_split(a.iter(), 1);
+        assert_eq!(first.count(), 1);
+        assert_eq!(rest.count(), 0);
+        let a = &["foo", "bar", "baz", "blah", "other"];
+        let (first, rest) = iterator_split(a.iter(), 2);
+        assert_eq!(first.count(), 2);
+        assert_eq!(rest.count(), 3);
+    }
+
+    #[test]
+    fn test_split_empty_iterator() {
+        let a: &[&str] = &[];
+        for v in [0, 1, 5] {
+            assert!(iterator_split_nonempty_rest_count(a.iter(), v).is_none());
+        }
+    }
+
+    #[test]
+    fn test_split_nonempty_iterator() {
+        let a = &["foo"];
+
+        let Some((elts, 1)) = iterator_split_nonempty_rest_count(a.iter(), 0) else {
+            panic!()
+        };
+        assert_eq!(elts.count(), 0);
+
+        let Some((elts, 0)) = iterator_split_nonempty_rest_count(a.iter(), 1) else {
+            panic!()
+        };
+        assert_eq!(elts.count(), 1);
+
+        let Some((elts, 0)) = iterator_split_nonempty_rest_count(a.iter(), 5) else {
+            panic!()
+        };
+        assert_eq!(elts.count(), 1);
+
+        let a = &["foo", "bar", "baz", "blah", "other"];
+        let Some((elts, 3)) = iterator_split_nonempty_rest_count(a.iter(), 2) else {
+            panic!()
+        };
+        assert_eq!(elts.count(), 2);
+    }
+}