zkvm: add partial std support

Co-authored-by: Frank Laub <flaub@risc0.com>
Co-authored-by: nils <nils@risc0.com>
Co-authored-by: Victor Graf <victor@risczero.com>
Co-authored-by: weikengchen <w.k@berkeley.edu>

11 files changed, 487 insertions, 5 deletions

diff --git a/library/std/build.rs b/library/std/build.rs
index 0f5068b59b7..60c097db2f4 100644
--- a/library/std/build.rs
+++ b/library/std/build.rs
@@ -35,6 +35,7 @@ fn main() {
         || target.contains("hurd")
         || target.contains("uefi")
         || target.contains("teeos")
+        || target.contains("zkvm")
         // See src/bootstrap/synthetic_targets.rs
         || env::var("RUSTC_BOOTSTRAP_SYNTHETIC_TARGET").is_ok()
     {
diff --git a/library/std/src/sys/pal/common/alloc.rs b/library/std/src/sys/pal/common/alloc.rs
index b7357460f39..8cf9ef68047 100644
--- a/library/std/src/sys/pal/common/alloc.rs
+++ b/library/std/src/sys/pal/common/alloc.rs
@@ -16,7 +16,7 @@ use crate::ptr;
     target_arch = "sparc",
     target_arch = "wasm32",
     target_arch = "hexagon",
-    all(target_arch = "riscv32", not(target_os = "espidf")),
+    all(target_arch = "riscv32", not(any(target_os = "espidf", target_os = "zkvm"))),
     all(target_arch = "xtensa", not(target_os = "espidf")),
 ))]
 pub const MIN_ALIGN: usize = 8;
@@ -32,11 +32,11 @@ pub const MIN_ALIGN: usize = 8;
     target_arch = "wasm64",
 ))]
 pub const MIN_ALIGN: usize = 16;
-// The allocator on the esp-idf platform guarantees 4 byte alignment.
-#[cfg(any(
-    all(target_arch = "riscv32", target_os = "espidf"),
+// The allocator on the esp-idf and zkvm platforms guarantee 4 byte alignment.
+#[cfg(all(any(
+    all(target_arch = "riscv32", any(target_os = "espidf", target_os = "zkvm")),
     all(target_arch = "xtensa", target_os = "espidf"),
-))]
+)))]
 pub const MIN_ALIGN: usize = 4;
 
 pub unsafe fn realloc_fallback(
diff --git a/library/std/src/sys/pal/mod.rs b/library/std/src/sys/pal/mod.rs
index 66b2a4b8885..041b7c35582 100644
--- a/library/std/src/sys/pal/mod.rs
+++ b/library/std/src/sys/pal/mod.rs
@@ -55,6 +55,9 @@ cfg_if::cfg_if! {
     } else if #[cfg(target_os = "teeos")] {
         mod teeos;
         pub use self::teeos::*;
+    } else if #[cfg(target_os = "zkvm")] {
+        mod zkvm;
+        pub use self::zkvm::*;
     } else {
         mod unsupported;
         pub use self::unsupported::*;
diff --git a/library/std/src/sys/pal/zkvm/abi.rs b/library/std/src/sys/pal/zkvm/abi.rs
new file mode 100644
index 00000000000..53332d90e02
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/abi.rs
@@ -0,0 +1,55 @@
+//! ABI definitions for symbols exported by risc0-zkvm-platform.
+
+// Included here so we don't have to depend on risc0-zkvm-platform.
+//
+// FIXME: Should we move this to the "libc" crate?  It seems like other
+// architectures put a lot of this kind of stuff there.  But there's
+// currently no risc0 fork of the libc crate, so we'd either have to
+// fork it or upstream it.
+
+#![allow(dead_code)]
+pub const DIGEST_WORDS: usize = 8;
+
+/// Standard IO file descriptors for use with sys_read and sys_write.
+pub mod fileno {
+    pub const STDIN: u32 = 0;
+    pub const STDOUT: u32 = 1;
+    pub const STDERR: u32 = 2;
+    pub const JOURNAL: u32 = 3;
+}
+
+extern "C" {
+    // Wrappers around syscalls provided by risc0-zkvm-platform:
+    pub fn sys_halt();
+    pub fn sys_output(output_id: u32, output_value: u32);
+    pub fn sys_sha_compress(
+        out_state: *mut [u32; DIGEST_WORDS],
+        in_state: *const [u32; DIGEST_WORDS],
+        block1_ptr: *const [u32; DIGEST_WORDS],
+        block2_ptr: *const [u32; DIGEST_WORDS],
+    );
+    pub fn sys_sha_buffer(
+        out_state: *mut [u32; DIGEST_WORDS],
+        in_state: *const [u32; DIGEST_WORDS],
+        buf: *const u8,
+        count: u32,
+    );
+    pub fn sys_rand(recv_buf: *mut u32, words: usize);
+    pub fn sys_panic(msg_ptr: *const u8, len: usize) -> !;
+    pub fn sys_log(msg_ptr: *const u8, len: usize);
+    pub fn sys_cycle_count() -> usize;
+    pub fn sys_read(fd: u32, recv_buf: *mut u8, nrequested: usize) -> usize;
+    pub fn sys_write(fd: u32, write_buf: *const u8, nbytes: usize);
+    pub fn sys_getenv(
+        recv_buf: *mut u32,
+        words: usize,
+        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;
+    pub fn sys_alloc_aligned(nwords: usize, align: usize) -> *mut u8;
+}
diff --git a/library/std/src/sys/pal/zkvm/alloc.rs b/library/std/src/sys/pal/zkvm/alloc.rs
new file mode 100644
index 00000000000..fd333f12151
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/alloc.rs
@@ -0,0 +1,15 @@
+use super::abi;
+use crate::alloc::{GlobalAlloc, Layout, System};
+
+#[stable(feature = "alloc_system_type", since = "1.28.0")]
+unsafe impl GlobalAlloc for System {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        abi::sys_alloc_aligned(layout.size(), layout.align())
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
+        // this allocator never deallocates memory
+    }
+}
diff --git a/library/std/src/sys/pal/zkvm/args.rs b/library/std/src/sys/pal/zkvm/args.rs
new file mode 100644
index 00000000000..7753cf63840
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/args.rs
@@ -0,0 +1,80 @@
+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)
+        }
+    }
+}
diff --git a/library/std/src/sys/pal/zkvm/env.rs b/library/std/src/sys/pal/zkvm/env.rs
new file mode 100644
index 00000000000..b85153642b1
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/env.rs
@@ -0,0 +1,9 @@
+pub mod os {
+    pub const FAMILY: &str = "";
+    pub const OS: &str = "";
+    pub const DLL_PREFIX: &str = "";
+    pub const DLL_SUFFIX: &str = ".elf";
+    pub const DLL_EXTENSION: &str = "elf";
+    pub const EXE_SUFFIX: &str = ".elf";
+    pub const EXE_EXTENSION: &str = "elf";
+}
diff --git a/library/std/src/sys/pal/zkvm/mod.rs b/library/std/src/sys/pal/zkvm/mod.rs
new file mode 100644
index 00000000000..7f221dc4fd9
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/mod.rs
@@ -0,0 +1,93 @@
+//! System bindings for the risc0 zkvm platform
+//!
+//! This module contains the facade (aka platform-specific) implementations of
+//! OS level functionality for zkvm.
+//!
+//! This is all super highly experimental and not actually intended for
+//! wide/production use yet, it's still all in the experimental category. This
+//! will likely change over time.
+
+const WORD_SIZE: usize = core::mem::size_of::<u32>();
+
+pub mod alloc;
+#[path = "../zkvm/args.rs"]
+pub mod args;
+#[path = "../unix/cmath.rs"]
+pub mod cmath;
+pub mod env;
+#[path = "../unsupported/fs.rs"]
+pub mod fs;
+#[path = "../unsupported/io.rs"]
+pub mod io;
+#[path = "../unsupported/net.rs"]
+pub mod net;
+#[path = "../unsupported/once.rs"]
+pub mod once;
+pub mod os;
+#[path = "../unix/os_str.rs"]
+pub mod os_str;
+#[path = "../unix/path.rs"]
+pub mod path;
+#[path = "../unsupported/pipe.rs"]
+pub mod pipe;
+#[path = "../unsupported/process.rs"]
+pub mod process;
+pub mod stdio;
+pub mod thread_local_key;
+#[path = "../unsupported/time.rs"]
+pub mod time;
+
+#[path = "../unsupported/locks/mod.rs"]
+pub mod locks;
+#[path = "../unsupported/thread.rs"]
+pub mod thread;
+
+#[path = "../unsupported/thread_parking.rs"]
+pub mod thread_parking;
+
+mod abi;
+
+use crate::io as std_io;
+
+pub mod memchr {
+    pub use core::slice::memchr::{memchr, memrchr};
+}
+
+// SAFETY: must be called only once during runtime initialization.
+// NOTE: this is not guaranteed to run, for example when Rust code is called externally.
+pub unsafe fn init(_argc: isize, _argv: *const *const u8, _sigpipe: u8) {}
+
+// SAFETY: must be called only once during runtime cleanup.
+// NOTE: this is not guaranteed to run, for example when the program aborts.
+pub unsafe fn cleanup() {}
+
+pub fn unsupported<T>() -> std_io::Result<T> {
+    Err(unsupported_err())
+}
+
+pub fn unsupported_err() -> std_io::Error {
+    std_io::const_io_error!(
+        std_io::ErrorKind::Unsupported,
+        "operation not supported on this platform",
+    )
+}
+
+pub fn is_interrupted(_code: i32) -> bool {
+    false
+}
+
+pub fn decode_error_kind(_code: i32) -> crate::io::ErrorKind {
+    crate::io::ErrorKind::Uncategorized
+}
+
+pub fn abort_internal() -> ! {
+    core::intrinsics::abort();
+}
+
+pub fn hashmap_random_keys() -> (u64, u64) {
+    let mut buf = [0u32; 4];
+    unsafe {
+        abi::sys_rand(buf.as_mut_ptr(), 4);
+    };
+    ((buf[0] as u64) << 32 + buf[1] as u64, (buf[2] as u64) << 32 + buf[3] as u64)
+}
diff --git a/library/std/src/sys/pal/zkvm/os.rs b/library/std/src/sys/pal/zkvm/os.rs
new file mode 100644
index 00000000000..d8739ee3824
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/os.rs
@@ -0,0 +1,139 @@
+use super::{abi, unsupported, WORD_SIZE};
+use crate::error::Error as StdError;
+use crate::ffi::{OsStr, OsString};
+use crate::fmt;
+use crate::io;
+use crate::marker::PhantomData;
+use crate::path::{self, PathBuf};
+use crate::sys_common::FromInner;
+
+pub fn errno() -> i32 {
+    0
+}
+
+pub fn error_string(_errno: i32) -> String {
+    "operation successful".to_string()
+}
+
+pub fn getcwd() -> io::Result<PathBuf> {
+    unsupported()
+}
+
+pub fn chdir(_: &path::Path) -> io::Result<()> {
+    unsupported()
+}
+
+pub struct SplitPaths<'a>(!, PhantomData<&'a ()>);
+
+pub fn split_paths(_unparsed: &OsStr) -> SplitPaths<'_> {
+    panic!("unsupported")
+}
+
+impl<'a> Iterator for SplitPaths<'a> {
+    type Item = PathBuf;
+    fn next(&mut self) -> Option<PathBuf> {
+        self.0
+    }
+}
+
+#[derive(Debug)]
+pub struct JoinPathsError;
+
+pub fn join_paths<I, T>(_paths: I) -> Result<OsString, JoinPathsError>
+where
+    I: Iterator<Item = T>,
+    T: AsRef<OsStr>,
+{
+    Err(JoinPathsError)
+}
+
+impl fmt::Display for JoinPathsError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        "not supported on this platform yet".fmt(f)
+    }
+}
+
+impl StdError for JoinPathsError {
+    #[allow(deprecated)]
+    fn description(&self) -> &str {
+        "not supported on this platform yet"
+    }
+}
+
+pub fn current_exe() -> io::Result<PathBuf> {
+    unsupported()
+}
+
+pub struct Env(!);
+
+impl Iterator for Env {
+    type Item = (OsString, OsString);
+    fn next(&mut self) -> Option<(OsString, OsString)> {
+        self.0
+    }
+}
+
+pub fn env() -> Env {
+    panic!("not supported on this platform")
+}
+
+impl Env {
+    pub fn str_debug(&self) -> impl fmt::Debug + '_ {
+        let Self(inner) = self;
+        match *inner {}
+    }
+}
+
+impl fmt::Debug for Env {
+    fn fmt(&self, _: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let Self(inner) = self;
+        match *inner {}
+    }
+}
+
+pub fn getenv(varname: &OsStr) -> Option<OsString> {
+    let varname = varname.as_encoded_bytes();
+    let nbytes =
+        unsafe { abi::sys_getenv(crate::ptr::null_mut(), 0, varname.as_ptr(), varname.len()) };
+    if nbytes == usize::MAX {
+        return None;
+    }
+
+    let nwords = (nbytes + WORD_SIZE - 1) / WORD_SIZE;
+    let words = unsafe { abi::sys_alloc_words(nwords) };
+
+    let nbytes2 = unsafe { abi::sys_getenv(words, nwords, varname.as_ptr(), varname.len()) };
+    debug_assert_eq!(nbytes, nbytes2);
+
+    // 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 u8s: &[u8] = unsafe { crate::slice::from_raw_parts(words.cast() as *const u8, nbytes) };
+    Some(OsString::from_inner(super::os_str::Buf { inner: u8s.to_vec() }))
+}
+
+pub fn setenv(_: &OsStr, _: &OsStr) -> io::Result<()> {
+    Err(io::const_io_error!(io::ErrorKind::Unsupported, "cannot set env vars on this platform"))
+}
+
+pub fn unsetenv(_: &OsStr) -> io::Result<()> {
+    Err(io::const_io_error!(io::ErrorKind::Unsupported, "cannot unset env vars on this platform"))
+}
+
+pub fn temp_dir() -> PathBuf {
+    panic!("no filesystem on this platform")
+}
+
+pub fn home_dir() -> Option<PathBuf> {
+    None
+}
+
+pub fn exit(_code: i32) -> ! {
+    crate::intrinsics::abort()
+}
+
+pub fn getpid() -> u32 {
+    panic!("no pids on this platform")
+}
diff --git a/library/std/src/sys/pal/zkvm/stdio.rs b/library/std/src/sys/pal/zkvm/stdio.rs
new file mode 100644
index 00000000000..e771ed0de28
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/stdio.rs
@@ -0,0 +1,64 @@
+use super::{abi, abi::fileno};
+use crate::io;
+
+pub struct Stdin;
+pub struct Stdout;
+pub struct Stderr;
+
+impl Stdin {
+    pub const fn new() -> Stdin {
+        Stdin
+    }
+}
+
+impl io::Read for Stdin {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        Ok(unsafe { abi::sys_read(fileno::STDIN, buf.as_mut_ptr(), buf.len()) })
+    }
+}
+
+impl Stdout {
+    pub const fn new() -> Stdout {
+        Stdout
+    }
+}
+
+impl io::Write for Stdout {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        unsafe { abi::sys_write(fileno::STDOUT, buf.as_ptr(), buf.len()) }
+
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+impl Stderr {
+    pub const fn new() -> Stderr {
+        Stderr
+    }
+}
+
+impl io::Write for Stderr {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        unsafe { abi::sys_write(fileno::STDERR, buf.as_ptr(), buf.len()) }
+
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+pub const STDIN_BUF_SIZE: usize = crate::sys_common::io::DEFAULT_BUF_SIZE;
+
+pub fn is_ebadf(_err: &io::Error) -> bool {
+    true
+}
+
+pub fn panic_output() -> Option<impl io::Write> {
+    Some(Stderr::new())
+}
diff --git a/library/std/src/sys/pal/zkvm/thread_local_key.rs b/library/std/src/sys/pal/zkvm/thread_local_key.rs
new file mode 100644
index 00000000000..3ffe6247344
--- /dev/null
+++ b/library/std/src/sys/pal/zkvm/thread_local_key.rs
@@ -0,0 +1,23 @@
+use crate::alloc::{alloc, Layout};
+
+pub type Key = usize;
+
+#[inline]
+pub unsafe fn create(_dtor: Option<unsafe extern "C" fn(*mut u8)>) -> Key {
+    alloc(Layout::new::<*mut u8>()) as _
+}
+
+#[inline]
+pub unsafe fn set(key: Key, value: *mut u8) {
+    let key: *mut *mut u8 = core::ptr::from_exposed_addr_mut(key);
+    *key = value;
+}
+
+#[inline]
+pub unsafe fn get(key: Key) -> *mut u8 {
+    let key: *mut *mut u8 = core::ptr::from_exposed_addr_mut(key);
+    *key
+}
+
+#[inline]
+pub unsafe fn destroy(_key: Key) {}