Add support for panic handler, args, and tests

library/panic_abort/src/lib.rs

@@ -18,6 +18,9 @@
 #[cfg(target_os = "android")]
 mod android;
 
+#[cfg(target_os = "zkvm")]
+mod zkvm;
+
 use core::any::Any;
 use core::panic::BoxMeUp;
 
@@ -33,6 +36,8 @@ pub unsafe fn __rust_start_panic(_payload: *mut &mut dyn BoxMeUp) -> u32 {
     // Android has the ability to attach a message as part of the abort.
     #[cfg(target_os = "android")]
     android::android_set_abort_message(_payload);
+    #[cfg(target_os = "zkvm")]
+    zkvm::zkvm_set_abort_message(_payload);
 
     abort();
 

library/panic_abort/src/zkvm.rs

@@ -0,0 +1,31 @@
+use alloc::string::String;
+use core::panic::BoxMeUp;
+
+// Forward the abort message to libc's android_set_abort_message. We try our best to populate the
+// message but as this function may already be called as part of a failed allocation, it might not be
+// possible to do so.
+//
+// Some methods of core are on purpose avoided (such as try_reserve) as these rely on the correct
+// resolution of rust_eh_personality which is loosely defined in panic_abort.
+//
+// Weakly resolve the symbol for android_set_abort_message. This function is only available
+// for API >= 21.
+pub(crate) unsafe fn zkvm_set_abort_message(payload: *mut &mut dyn BoxMeUp) {
+    let payload = (*payload).get();
+    let msg = match payload.downcast_ref::<&'static str>() {
+        Some(msg) => msg.as_bytes(),
+        None => match payload.downcast_ref::<String>() {
+            Some(msg) => msg.as_bytes(),
+            None => &[],
+        },
+    };
+    if msg.is_empty() {
+        return;
+    }
+
+    extern "C" {
+        fn sys_panic(msg_ptr: *const u8, len: usize) -> !;
+    }
+
+    sys_panic(msg.as_ptr(), msg.len());
+}

library/std/src/sys/zkvm/abi.rs

@@ -46,6 +46,12 @@ extern "C" {
         varname: *const u8,
         varname_len: usize,
     ) -> usize;
+    pub fn sys_argc() -> usize;
+    pub fn sys_argv(
+        out_words: *mut u32,
+        out_nwords: usize,
+        arg_index: usize,
+    ) -> usize;
 
     // Allocate memory from global HEAP.
     pub fn sys_alloc_words(nwords: usize) -> *mut u32;

library/std/src/sys/zkvm/alloc.rs

@@ -1,4 +1,4 @@
-use super::{abi, WORD_SIZE};
+use super::abi;
 use crate::alloc::{GlobalAlloc, Layout, System};
 
 #[stable(feature = "alloc_system_type", since = "1.28.0")]

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

@@ -0,0 +1,79 @@
+use super::{abi, WORD_SIZE};
+use crate::ffi::OsString;
+use crate::fmt;
+use crate::sys_common::FromInner;
+
+pub struct Args {
+    i_forward: usize,
+    i_back: usize,
+    count: usize,
+}
+
+pub fn args() -> Args {
+    let count = unsafe { abi::sys_argc() };
+    Args { i_forward: 0, i_back: 0, count }
+}
+
+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) };
+
+        let arg_len_words = (arg_len + WORD_SIZE - 1) / WORD_SIZE;
+        let words = unsafe { abi::sys_alloc_words(arg_len_words) };
+
+        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(super::os_str::Buf { inner: arg_bytes.to_vec() })
+    }
+}
+
+impl fmt::Debug for Args {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_list().finish()
+    }
+}
+
+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)
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.count, Some(self.count))
+    }
+}
+
+impl ExactSizeIterator for Args {
+    fn len(&self) -> usize {
+        self.count
+    }
+}
+
+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)
+        }
+    }
+}

library/std/src/sys/zkvm/mod.rs

@@ -10,7 +10,7 @@
 const WORD_SIZE: usize = core::mem::size_of::<u32>();
 
 pub mod alloc;
-#[path = "../unsupported/args.rs"]
+#[path = "../zkvm/args.rs"]
 pub mod args;
 #[path = "../unix/cmath.rs"]
 pub mod cmath;

library/test/src/console.rs

@@ -292,7 +292,7 @@ pub fn run_tests_console(opts: &TestOpts, tests: Vec<TestDescAndFn>) -> io::Resu
     // Prevent the usage of `Instant` in some cases:
     // - It's currently not supported for wasm targets.
     // - We disable it for miri because it's not available when isolation is enabled.
-    let is_instant_supported = !cfg!(target_family = "wasm") && !cfg!(miri);
+    let is_instant_supported = !cfg!(target_family = "wasm") && !cfg!(target_os = "zkvm") && !cfg!(miri);
 
     let start_time = is_instant_supported.then(Instant::now);
     run_tests(opts, tests, |x| on_test_event(&x, &mut st, &mut *out))?;

library/test/src/lib.rs

@@ -545,7 +545,7 @@ pub fn run_test(
 
     // Emscripten can catch panics but other wasm targets cannot
     let ignore_because_no_process_support = desc.should_panic != ShouldPanic::No
-        && cfg!(target_family = "wasm")
+        && (cfg!(target_family = "wasm") || cfg!(target_os = "zkvm"))
         && !cfg!(target_os = "emscripten");
 
     if force_ignore || desc.ignore || ignore_because_no_process_support {
@@ -592,7 +592,7 @@ pub fn run_test(
         // If the platform is single-threaded we're just going to run
         // the test synchronously, regardless of the concurrency
         // level.
-        let supports_threads = !cfg!(target_os = "emscripten") && !cfg!(target_family = "wasm");
+        let supports_threads = !cfg!(target_os = "emscripten") && !cfg!(target_family = "wasm") && !cfg!(target_os = "zkvm");
         if supports_threads {
             let cfg = thread::Builder::new().name(name.as_slice().to_owned());
             let mut runtest = Arc::new(Mutex::new(Some(runtest)));