Tweak iterator filtering

Prep for lint truncation improvement.

- Kill the handling of the "max 0" case, it's pathological; this
  requires passing NonZeroUsize
- Drop the requirement for the iterator to be clonable as it
  improves ergonomics, and return a vector of the collected elements
  instead

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

Author: cgwalters

lib/src/fsck.rs

@@ -8,9 +8,10 @@
 
 use std::fmt::Write as _;
 use std::future::Future;
+use std::num::NonZeroUsize;
 use std::pin::Pin;
 
-use bootc_utils::iterator_split_nonempty_rest_count;
+use bootc_utils::collect_until;
 use camino::Utf8PathBuf;
 use cap_std::fs::{Dir, MetadataExt as _};
 use cap_std_ext::cap_std;
@@ -250,9 +251,10 @@ async fn check_fsverity_inner(storage: &Storage) -> FsckResult {
     let verity_found_state =
         verity_state_of_all_objects(&storage.repo(), verity_state.desired == Tristate::Enabled)
             .await?;
-    let Some((missing, rest)) =
-        iterator_split_nonempty_rest_count(verity_found_state.missing.iter(), 5)
-    else {
+    let Some((missing, rest)) = collect_until(
+        verity_found_state.missing.iter(),
+        const { NonZeroUsize::new(5).unwrap() },
+    ) else {
         return fsck_ok();
     };
     let mut err = String::from("fsverity enabled, but objects without fsverity:\n");

lib/src/lints.rs

@@ -8,6 +8,7 @@
 use std::collections::{BTreeMap, BTreeSet};
 use std::env::consts::ARCH;
 use std::fmt::Write as WriteFmt;
+use std::num::NonZeroUsize;
 use std::ops::ControlFlow;
 use std::os::unix::ffi::OsStrExt;
 use std::path::Path;
@@ -31,6 +32,9 @@ const BASEIMAGE_REF: &str = "usr/share/doc/bootc/baseimage/base";
 // https://systemd.io/API_FILE_SYSTEMS/ with /var added for us
 const API_DIRS: &[&str] = &["dev", "proc", "sys", "run", "tmp", "var"];
 
+/// Only output this many items by default
+const DEFAULT_TRUNCATED_OUTPUT: NonZeroUsize = const { NonZeroUsize::new(5).unwrap() };
+
 /// A lint check has failed.
 #[derive(thiserror::Error, Debug)]
 struct LintError(String);
@@ -640,7 +644,7 @@ fn check_var_tmpfiles(_root: &Dir) -> LintResult {
     }
     let mut msg = String::new();
     if let Some((samples, rest)) =
-        bootc_utils::iterator_split_nonempty_rest_count(r.tmpfiles.iter(), 5)
+        bootc_utils::collect_until(r.tmpfiles.iter(), DEFAULT_TRUNCATED_OUTPUT)
     {
         msg.push_str("Found content in /var missing systemd tmpfiles.d entries:\n");
         for elt in samples {
@@ -651,10 +655,10 @@ fn check_var_tmpfiles(_root: &Dir) -> LintResult {
         }
     }
     if let Some((samples, rest)) =
-        bootc_utils::iterator_split_nonempty_rest_count(r.unsupported.iter(), 5)
+        bootc_utils::collect_until(r.unsupported.iter(), DEFAULT_TRUNCATED_OUTPUT)
     {
         msg.push_str("Found non-directory/non-symlink files in /var:\n");
-        for elt in samples.map(PathQuotedDisplay::new) {
+        for elt in samples.iter().map(PathQuotedDisplay::new) {
             writeln!(msg, "  {elt}")?;
         }
         if rest > 0 {
@@ -688,7 +692,7 @@ fn check_sysusers(rootfs: &Dir) -> LintResult {
     }
     let mut msg = String::new();
     if let Some((samples, rest)) =
-        bootc_utils::iterator_split_nonempty_rest_count(r.missing_users.iter(), 5)
+        bootc_utils::collect_until(r.missing_users.iter(), DEFAULT_TRUNCATED_OUTPUT)
     {
         msg.push_str("Found /etc/passwd entry without corresponding systemd sysusers.d:\n");
         for elt in samples {
@@ -699,7 +703,7 @@ fn check_sysusers(rootfs: &Dir) -> LintResult {
         }
     }
     if let Some((samples, rest)) =
-        bootc_utils::iterator_split_nonempty_rest_count(r.missing_groups.iter(), 5)
+        bootc_utils::collect_until(r.missing_groups.iter(), DEFAULT_TRUNCATED_OUTPUT)
     {
         msg.push_str("Found /etc/group entry without corresponding systemd sysusers.d:\n");
         for elt in samples {

ostree-ext/src/container/store.rs

@@ -31,6 +31,7 @@ use ostree::{gio, glib};
 use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::fmt::Write as _;
 use std::iter::FromIterator;
+use std::num::NonZeroUsize;
 use tokio::sync::mpsc::{Receiver, Sender};
 
 /// Configuration for the proxy.
@@ -963,9 +964,10 @@ impl ImageImporter {
                 tracing::debug!("Imported layer: {}", r.commit.as_str());
                 layer_commits.push(r.commit);
                 let filtered_owned = HashMap::from_iter(r.filtered.clone());
-                if let Some((filtered, n_rest)) =
-                    bootc_utils::iterator_split_nonempty_rest_count(r.filtered.iter(), 5)
-                {
+                if let Some((filtered, n_rest)) = bootc_utils::collect_until(
+                    r.filtered.iter(),
+                    const { NonZeroUsize::new(5).unwrap() },
+                ) {
                     let mut msg = String::new();
                     for (path, n) in filtered {
                         write!(msg, "{path}: {n} ").unwrap();

utils/src/iterators.rs

@@ -1,3 +1,5 @@
+use std::num::NonZeroUsize;
+
 /// Given an iterator that's clonable, split it into two iterators
 /// at a given maximum number of elements.
 pub fn iterator_split<I>(
@@ -11,23 +13,31 @@ where
     (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)>
+/// Gather the first N items, and provide the count of the remaining items.
+/// The max count cannot be zero as that's a pathological case.
+pub fn collect_until<I>(it: I, max: NonZeroUsize) -> Option<(Vec<I::Item>, usize)>
 where
-    I: Iterator + Clone,
+    I: Iterator,
 {
-    let rest = it.clone();
+    let mut items = Vec::with_capacity(max.get());
+
     let mut it = it.peekable();
-    if it.peek().is_some() {
-        Some((it.take(max), rest.skip(max).count()))
-    } else {
-        None
+    if it.peek().is_none() {
+        return None;
+    }
+
+    while let Some(next) = it.next() {
+        items.push(next);
+
+        // If we've reached max items, stop collecting
+        if items.len() == max.get() {
+            break;
+        }
     }
+    // Count remaining items
+    let remaining = it.count();
+    items.shrink_to_fit();
+    Some((items, remaining))
 }
 
 #[cfg(test)]
@@ -60,34 +70,31 @@ mod tests {
     #[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());
+        for v in [1, 5].into_iter().map(|v| NonZeroUsize::new(v).unwrap()) {
+            assert!(collect_until(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 {
+        let Some((elts, 0)) = collect_until(a.iter(), NonZeroUsize::new(1).unwrap()) else {
             panic!()
         };
-        assert_eq!(elts.count(), 1);
+        assert_eq!(elts.len(), 1);
 
-        let Some((elts, 0)) = iterator_split_nonempty_rest_count(a.iter(), 5) else {
+        let Some((elts, 0)) = collect_until(a.iter(), const { NonZeroUsize::new(5).unwrap() })
+        else {
             panic!()
         };
-        assert_eq!(elts.count(), 1);
+        assert_eq!(elts.len(), 1);
 
         let a = &["foo", "bar", "baz", "blah", "other"];
-        let Some((elts, 3)) = iterator_split_nonempty_rest_count(a.iter(), 2) else {
+        let Some((elts, 3)) = collect_until(a.iter(), const { NonZeroUsize::new(2).unwrap() })
+        else {
             panic!()
         };
-        assert_eq!(elts.count(), 2);
+        assert_eq!(elts.len(), 2);
     }
 }