Rollup merge of #139847 - thaliaarchi:args/delegate-iter, r=workingjubilee

Delegate to inner `vec::IntoIter` from `env::ArgsOs`

Delegate from `std::env::ArgsOs` to the methods of the inner platform-specific iterators, when it would be more efficient than just using the default methods of its own impl. Most platforms use `vec::IntoIter` as the inner type, so prioritize delegating to the methods it provides.

`std::env::Args` is implemented atop `std::env::ArgsOs` and performs UTF-8 validation with a panic for invalid data. This is a visible effect which users certainly rely on, so we can't skip any arguments. Any further iterator methods would skip some elements, so no change is needed for that type.

Add `#[inline]` for any methods which simply wrap the inner iterator.

Author: Zalathar

library/std/src/env.rs

@@ -12,9 +12,11 @@
 
 use crate::error::Error;
 use crate::ffi::{OsStr, OsString};
+use crate::num::NonZero;
+use crate::ops::Try;
 use crate::path::{Path, PathBuf};
 use crate::sys::{env as env_imp, os as os_imp};
-use crate::{fmt, io, sys};
+use crate::{array, fmt, io, sys};
 
 /// Returns the current working directory as a [`PathBuf`].
 ///
@@ -872,19 +874,36 @@ impl !Sync for Args {}
 #[stable(feature = "env", since = "1.0.0")]
 impl Iterator for Args {
     type Item = String;
+
     fn next(&mut self) -> Option<String> {
         self.inner.next().map(|s| s.into_string().unwrap())
     }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.inner.size_hint()
     }
+
+    // Methods which skip args cannot simply delegate to the inner iterator,
+    // because `env::args` states that we will "panic during iteration if any
+    // argument to the process is not valid Unicode".
+    //
+    // This offers two possible interpretations:
+    // - a skipped argument is never encountered "during iteration"
+    // - even a skipped argument is encountered "during iteration"
+    //
+    // As a panic can be observed, we err towards validating even skipped
+    // arguments for now, though this is not explicitly promised by the API.
 }
 
 #[stable(feature = "env", since = "1.0.0")]
 impl ExactSizeIterator for Args {
+    #[inline]
     fn len(&self) -> usize {
         self.inner.len()
     }
+
+    #[inline]
     fn is_empty(&self) -> bool {
         self.inner.is_empty()
     }
@@ -914,29 +933,81 @@ impl !Sync for ArgsOs {}
 #[stable(feature = "env", since = "1.0.0")]
 impl Iterator for ArgsOs {
     type Item = OsString;
+
+    #[inline]
     fn next(&mut self) -> Option<OsString> {
         self.inner.next()
     }
+
+    #[inline]
+    fn next_chunk<const N: usize>(
+        &mut self,
+    ) -> Result<[OsString; N], array::IntoIter<OsString, N>> {
+        self.inner.next_chunk()
+    }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.inner.size_hint()
     }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.inner.len()
+    }
+
+    #[inline]
+    fn last(self) -> Option<OsString> {
+        self.inner.last()
+    }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.inner.advance_by(n)
+    }
+
+    #[inline]
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        self.inner.try_fold(init, f)
+    }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.inner.fold(init, f)
+    }
 }
 
 #[stable(feature = "env", since = "1.0.0")]
 impl ExactSizeIterator for ArgsOs {
+    #[inline]
     fn len(&self) -> usize {
         self.inner.len()
     }
+
+    #[inline]
     fn is_empty(&self) -> bool {
         self.inner.is_empty()
     }
 }
 
 #[stable(feature = "env_iterators", since = "1.12.0")]
 impl DoubleEndedIterator for ArgsOs {
+    #[inline]
     fn next_back(&mut self) -> Option<OsString> {
         self.inner.next_back()
     }
+
+    #[inline]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.inner.advance_back_by(n)
+    }
 }
 
 #[stable(feature = "std_debug", since = "1.16.0")]

library/std/src/lib.rs

@@ -301,6 +301,8 @@
 #![feature(formatting_options)]
 #![feature(if_let_guard)]
 #![feature(intra_doc_pointers)]
+#![feature(iter_advance_by)]
+#![feature(iter_next_chunk)]
 #![feature(lang_items)]
 #![feature(let_chains)]
 #![feature(link_cfg)]
@@ -321,6 +323,7 @@
 #![feature(strict_provenance_lints)]
 #![feature(thread_local)]
 #![feature(try_blocks)]
+#![feature(try_trait_v2)]
 #![feature(type_alias_impl_trait)]
 // tidy-alphabetical-end
 //

library/std/src/sys/args/common.rs

@@ -1,5 +1,7 @@
 use crate::ffi::OsString;
-use crate::{fmt, vec};
+use crate::num::NonZero;
+use crate::ops::Try;
+use crate::{array, fmt, vec};
 
 pub struct Args {
     iter: vec::IntoIter<OsString>,
@@ -9,6 +11,7 @@ impl !Send for Args {}
 impl !Sync for Args {}
 
 impl Args {
+    #[inline]
     pub(super) fn new(args: Vec<OsString>) -> Self {
         Args { iter: args.into_iter() }
     }
@@ -22,22 +25,77 @@ impl fmt::Debug for Args {
 
 impl Iterator for Args {
     type Item = OsString;
+
+    #[inline]
     fn next(&mut self) -> Option<OsString> {
         self.iter.next()
     }
+
+    #[inline]
+    fn next_chunk<const N: usize>(
+        &mut self,
+    ) -> Result<[OsString; N], array::IntoIter<OsString, N>> {
+        self.iter.next_chunk()
+    }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.iter.size_hint()
     }
-}
 
-impl ExactSizeIterator for Args {
-    fn len(&self) -> usize {
+    #[inline]
+    fn count(self) -> usize {
         self.iter.len()
     }
+
+    #[inline]
+    fn last(mut self) -> Option<OsString> {
+        self.iter.next_back()
+    }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.iter.advance_by(n)
+    }
+
+    #[inline]
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        self.iter.try_fold(init, f)
+    }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.iter.fold(init, f)
+    }
 }
 
 impl DoubleEndedIterator for Args {
+    #[inline]
     fn next_back(&mut self) -> Option<OsString> {
         self.iter.next_back()
     }
+
+    #[inline]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.iter.advance_back_by(n)
+    }
+}
+
+impl ExactSizeIterator for Args {
+    #[inline]
+    fn len(&self) -> usize {
+        self.iter.len()
+    }
+
+    #[inline]
+    fn is_empty(&self) -> bool {
+        self.iter.is_empty()
+    }
 }

library/std/src/sys/args/sgx.rs

@@ -1,6 +1,8 @@
 #![allow(fuzzy_provenance_casts)] // FIXME: this module systematically confuses pointers and integers
 
 use crate::ffi::OsString;
+use crate::num::NonZero;
+use crate::ops::Try;
 use crate::sync::atomic::{Atomic, AtomicUsize, Ordering};
 use crate::sys::os_str::Buf;
 use crate::sys::pal::abi::usercalls::alloc;
@@ -28,35 +30,81 @@ pub unsafe fn init(argc: isize, argv: *const *const u8) {
 
 pub fn args() -> Args {
     let args = unsafe { (ARGS.load(Ordering::Relaxed) as *const ArgsStore).as_ref() };
-    if let Some(args) = args { Args(args.iter()) } else { Args([].iter()) }
+    let slice = args.map(|args| args.as_slice()).unwrap_or(&[]);
+    Args { iter: slice.iter() }
 }
 
-pub struct Args(slice::Iter<'static, OsString>);
+pub struct Args {
+    iter: slice::Iter<'static, OsString>,
+}
 
 impl fmt::Debug for Args {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.0.as_slice().fmt(f)
+        self.iter.as_slice().fmt(f)
     }
 }
 
 impl Iterator for Args {
     type Item = OsString;
+
     fn next(&mut self) -> Option<OsString> {
-        self.0.next().cloned()
+        self.iter.next().cloned()
     }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        self.0.size_hint()
+        self.iter.size_hint()
     }
-}
 
-impl ExactSizeIterator for Args {
-    fn len(&self) -> usize {
-        self.0.len()
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.len()
+    }
+
+    fn last(self) -> Option<OsString> {
+        self.iter.last().cloned()
+    }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.iter.advance_by(n)
+    }
+
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        self.iter.by_ref().cloned().try_fold(init, f)
+    }
+
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.iter.cloned().fold(init, f)
     }
 }
 
 impl DoubleEndedIterator for Args {
     fn next_back(&mut self) -> Option<OsString> {
-        self.0.next_back().cloned()
+        self.iter.next_back().cloned()
+    }
+
+    #[inline]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.iter.advance_back_by(n)
+    }
+}
+
+impl ExactSizeIterator for Args {
+    #[inline]
+    fn len(&self) -> usize {
+        self.iter.len()
+    }
+
+    #[inline]
+    fn is_empty(&self) -> bool {
+        self.iter.is_empty()
     }
 }

library/std/src/sys/args/unsupported.rs

@@ -15,22 +15,28 @@ impl fmt::Debug for Args {
 
 impl Iterator for Args {
     type Item = OsString;
+
+    #[inline]
     fn next(&mut self) -> Option<OsString> {
         None
     }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         (0, Some(0))
     }
 }
 
-impl ExactSizeIterator for Args {
-    fn len(&self) -> usize {
-        0
-    }
-}
-
 impl DoubleEndedIterator for Args {
+    #[inline]
     fn next_back(&mut self) -> Option<OsString> {
         None
     }
 }
+
+impl ExactSizeIterator for Args {
+    #[inline]
+    fn len(&self) -> usize {
+        0
+    }
+}