Auto merge of #139849 - thaliaarchi:args/zkvm, r=ibraheemdev

Fix `env::ArgsOs` for zkVM

The zkVM implementation of `env::ArgsOs` incorrectly reports the full length even after having iterated. Instead, use a range approach which works out to be simpler. Also, implement more iterator methods like the other platforms in #139847.

cc `@flaub` `@jbruestle` `@SchmErik`

1 files changed, 58 insertions, 45 deletions

diff --git a/library/std/src/sys/args/zkvm.rs b/library/std/src/sys/args/zkvm.rs
index 194ba7159d4..d26bf1eaff9 100644
--- a/library/std/src/sys/args/zkvm.rs
+++ b/library/std/src/sys/args/zkvm.rs
@@ -1,25 +1,20 @@
-use crate::ffi::OsString;
-use crate::fmt;
-use crate::sys::os_str;
+use crate::ffi::{OsStr, OsString};
+use crate::num::NonZero;
+use crate::sync::OnceLock;
 use crate::sys::pal::{WORD_SIZE, abi};
-use crate::sys_common::FromInner;
-
-pub struct Args {
-    i_forward: usize,
-    i_back: usize,
-    count: usize,
-}
+use crate::{fmt, ptr, slice};
 
 pub fn args() -> Args {
-    let count = unsafe { abi::sys_argc() };
-    Args { i_forward: 0, i_back: 0, count }
+    Args { iter: ARGS.get_or_init(|| get_args()).iter() }
 }
 
-impl Args {
-    /// Use sys_argv to get the arg at the requested index. Does not check that i is less than argc
-    /// and will not return if the index is out of bounds.
-    fn argv(i: usize) -> OsString {
-        let arg_len = unsafe { abi::sys_argv(crate::ptr::null_mut(), 0, i) };
+fn get_args() -> Vec<&'static OsStr> {
+    let argc = unsafe { abi::sys_argc() };
+    let mut args = Vec::with_capacity(argc);
+
+    for i in 0..argc {
+        // Get the size of the argument then the data.
+        let arg_len = unsafe { abi::sys_argv(ptr::null_mut(), 0, i) };
 
         let arg_len_words = (arg_len + WORD_SIZE - 1) / WORD_SIZE;
         let words = unsafe { abi::sys_alloc_words(arg_len_words) };
@@ -27,20 +22,24 @@ impl Args {
         let arg_len2 = unsafe { abi::sys_argv(words, arg_len_words, i) };
         debug_assert_eq!(arg_len, arg_len2);
 
-        // Convert to OsString.
-        //
-        // FIXME: We can probably get rid of the extra copy here if we
-        // reimplement "os_str" instead of just using the generic unix
-        // "os_str".
-        let arg_bytes: &[u8] =
-            unsafe { crate::slice::from_raw_parts(words.cast() as *const u8, arg_len) };
-        OsString::from_inner(os_str::Buf { inner: arg_bytes.to_vec() })
+        let arg_bytes = unsafe { slice::from_raw_parts(words.cast(), arg_len) };
+        args.push(unsafe { OsStr::from_encoded_bytes_unchecked(arg_bytes) });
     }
+    args
 }
 
+static ARGS: OnceLock<Vec<&'static OsStr>> = OnceLock::new();
+
+pub struct Args {
+    iter: slice::Iter<'static, &'static OsStr>,
+}
+
+impl !Send for Args {}
+impl !Sync for Args {}
+
 impl fmt::Debug for Args {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_list().finish()
+        self.iter.as_slice().fmt(f)
     }
 }
 
@@ -48,34 +47,48 @@ impl Iterator for Args {
     type Item = OsString;
 
     fn next(&mut self) -> Option<OsString> {
-        if self.i_forward >= self.count - self.i_back {
-            None
-        } else {
-            let arg = Self::argv(self.i_forward);
-            self.i_forward += 1;
-            Some(arg)
-        }
+        self.iter.next().map(|arg| arg.to_os_string())
     }
 
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        (self.count, Some(self.count))
+        self.iter.size_hint()
     }
-}
 
-impl ExactSizeIterator for Args {
-    fn len(&self) -> usize {
-        self.count
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.len()
+    }
+
+    fn last(self) -> Option<OsString> {
+        self.iter.last().map(|arg| arg.to_os_string())
+    }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.iter.advance_by(n)
     }
 }
 
 impl DoubleEndedIterator for Args {
     fn next_back(&mut self) -> Option<OsString> {
-        if self.i_back >= self.count - self.i_forward {
-            None
-        } else {
-            let arg = Self::argv(self.count - 1 - self.i_back);
-            self.i_back += 1;
-            Some(arg)
-        }
+        self.iter.next_back().map(|arg| arg.to_os_string())
+    }
+
+    #[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()
     }
 }