From a22c8ed6311979d292e835e5a5498d0e3cc95dc0 Mon Sep 17 00:00:00 2001 From: Thalia Archibald Date: Sat, 12 Apr 2025 16:33:38 -0700 Subject: [PATCH 1/2] zkVM: Fix env::ArgsOs 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. --- library/std/src/sys/args/zkvm.rs | 71 ++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 17 deletions(-) diff --git a/library/std/src/sys/args/zkvm.rs b/library/std/src/sys/args/zkvm.rs index 194ba7159d459..e4b1ad540ac04 100644 --- a/library/std/src/sys/args/zkvm.rs +++ b/library/std/src/sys/args/zkvm.rs @@ -1,18 +1,19 @@ use crate::ffi::OsString; use crate::fmt; +use crate::num::NonZero; use crate::sys::os_str; use crate::sys::pal::{WORD_SIZE, abi}; use crate::sys_common::FromInner; +#[derive(Clone)] pub struct Args { - i_forward: usize, - i_back: usize, - count: usize, + front: usize, + back: usize, } pub fn args() -> Args { let count = unsafe { abi::sys_argc() }; - Args { i_forward: 0, i_back: 0, count } + Args { front: 0, back: count } } impl Args { @@ -38,44 +39,80 @@ impl Args { } } +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() + f.debug_list().entries(self.clone()).finish() } } impl Iterator for Args { type Item = OsString; + #[inline] fn next(&mut self) -> Option { - if self.i_forward >= self.count - self.i_back { + if self.front == self.back { None } else { - let arg = Self::argv(self.i_forward); - self.i_forward += 1; + let arg = Self::argv(self.front); + self.front += 1; Some(arg) } } + #[inline] fn size_hint(&self) -> (usize, Option) { - (self.count, Some(self.count)) + let len = self.len(); + (len, Some(len)) } -} -impl ExactSizeIterator for Args { - fn len(&self) -> usize { - self.count + #[inline] + fn count(self) -> usize { + self.len() + } + + #[inline] + fn last(mut self) -> Option { + self.next_back() + } + + #[inline] + fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { + let step_size = self.len().min(n); + self.front += step_size; + NonZero::new(n - step_size).map_or(Ok(()), Err) } } impl DoubleEndedIterator for Args { + #[inline] fn next_back(&mut self) -> Option { - if self.i_back >= self.count - self.i_forward { + if self.back == self.front { None } else { - let arg = Self::argv(self.count - 1 - self.i_back); - self.i_back += 1; - Some(arg) + self.back -= 1; + Some(Self::argv(self.back)) } } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero> { + let step_size = self.len().min(n); + self.back -= step_size; + NonZero::new(n - step_size).map_or(Ok(()), Err) + } +} + +impl ExactSizeIterator for Args { + #[inline] + fn len(&self) -> usize { + self.back - self.front + } + + #[inline] + fn is_empty(&self) -> bool { + self.front == self.back + } } From 00095982a0d2f6c8bfb5cb2706338ad2b14304f5 Mon Sep 17 00:00:00 2001 From: Thalia Archibald Date: Tue, 15 Apr 2025 02:41:51 -0700 Subject: [PATCH 2/2] zkVM: Cache args from host Retrieve argc/argv from the host once, on demand when the first `env::ArgsOs` is constructed, and globally cache it. Copy each argument to an `OsString` while iterating. --- library/std/src/sys/args/zkvm.rs | 86 ++++++++++++-------------------- 1 file changed, 31 insertions(+), 55 deletions(-) diff --git a/library/std/src/sys/args/zkvm.rs b/library/std/src/sys/args/zkvm.rs index e4b1ad540ac04..d26bf1eaff91f 100644 --- a/library/std/src/sys/args/zkvm.rs +++ b/library/std/src/sys/args/zkvm.rs @@ -1,26 +1,20 @@ -use crate::ffi::OsString; -use crate::fmt; +use crate::ffi::{OsStr, OsString}; use crate::num::NonZero; -use crate::sys::os_str; +use crate::sync::OnceLock; use crate::sys::pal::{WORD_SIZE, abi}; -use crate::sys_common::FromInner; - -#[derive(Clone)] -pub struct Args { - front: usize, - back: usize, -} +use crate::{fmt, ptr, slice}; pub fn args() -> Args { - let count = unsafe { abi::sys_argc() }; - Args { front: 0, back: 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) }; @@ -28,15 +22,16 @@ 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> = OnceLock::new(); + +pub struct Args { + iter: slice::Iter<'static, &'static OsStr>, } impl !Send for Args {} @@ -44,75 +39,56 @@ impl !Sync for Args {} impl fmt::Debug for Args { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list().entries(self.clone()).finish() + self.iter.as_slice().fmt(f) } } impl Iterator for Args { type Item = OsString; - #[inline] fn next(&mut self) -> Option { - if self.front == self.back { - None - } else { - let arg = Self::argv(self.front); - self.front += 1; - Some(arg) - } + self.iter.next().map(|arg| arg.to_os_string()) } #[inline] fn size_hint(&self) -> (usize, Option) { - let len = self.len(); - (len, Some(len)) + self.iter.size_hint() } #[inline] fn count(self) -> usize { - self.len() + self.iter.len() } - #[inline] - fn last(mut self) -> Option { - self.next_back() + fn last(self) -> Option { + self.iter.last().map(|arg| arg.to_os_string()) } #[inline] fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { - let step_size = self.len().min(n); - self.front += step_size; - NonZero::new(n - step_size).map_or(Ok(()), Err) + self.iter.advance_by(n) } } impl DoubleEndedIterator for Args { - #[inline] fn next_back(&mut self) -> Option { - if self.back == self.front { - None - } else { - self.back -= 1; - Some(Self::argv(self.back)) - } + self.iter.next_back().map(|arg| arg.to_os_string()) } #[inline] fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero> { - let step_size = self.len().min(n); - self.back -= step_size; - NonZero::new(n - step_size).map_or(Ok(()), Err) + self.iter.advance_back_by(n) } } impl ExactSizeIterator for Args { #[inline] fn len(&self) -> usize { - self.back - self.front + self.iter.len() } #[inline] fn is_empty(&self) -> bool { - self.front == self.back + self.iter.is_empty() } }