4 files changed, 252 insertions, 17 deletions
diff --git a/library/std/src/sys/stdio/mod.rs b/library/std/src/sys/stdio/mod.rs
index 336d4c8527d..404ac877926 100644
--- a/library/std/src/sys/stdio/mod.rs
+++ b/library/std/src/sys/stdio/mod.rs
@@ -1,40 +1,55 @@
 #![forbid(unsafe_op_in_unsafe_fn)]
 
-cfg_if::cfg_if! {
-    if #[cfg(any(
-        target_family = "unix",
-        target_os = "hermit"
-    ))] {
+cfg_select! {
+    any(target_family = "unix", target_os = "hermit") => {
         mod unix;
         pub use unix::*;
-    } else if #[cfg(target_os = "windows")] {
+    }
+    target_os = "windows" => {
         mod windows;
         pub use windows::*;
-    } else if #[cfg(all(target_vendor = "fortanix", target_env = "sgx"))] {
+    }
+    all(target_vendor = "fortanix", target_env = "sgx") => {
         mod sgx;
         pub use sgx::*;
-    } else if #[cfg(target_os = "solid_asp3")] {
+    }
+    target_os = "solid_asp3" => {
         mod solid;
         pub use solid::*;
-    } else if #[cfg(target_os = "teeos")] {
+    }
+    target_os = "teeos" => {
         mod teeos;
         pub use teeos::*;
-    } else if #[cfg(target_os = "trusty")] {
+    }
+    target_os = "trusty" => {
         mod trusty;
         pub use trusty::*;
-    } else if #[cfg(target_os = "uefi")] {
+    }
+    target_os = "uefi" => {
         mod uefi;
         pub use uefi::*;
-    } else if #[cfg(target_os = "wasi")] {
-        mod wasi;
-        pub use wasi::*;
-    } else if #[cfg(target_os = "xous")] {
+    }
+    target_os = "vexos" => {
+        mod vexos;
+        pub use vexos::*;
+    }
+    all(target_os = "wasi", target_env = "p1") => {
+        mod wasip1;
+        pub use wasip1::*;
+    }
+    all(target_os = "wasi", target_env = "p2") => {
+        mod wasip2;
+        pub use wasip2::*;
+    }
+    target_os = "xous" => {
         mod xous;
         pub use xous::*;
-    } else if #[cfg(target_os = "zkvm")] {
+    }
+    target_os = "zkvm" => {
         mod zkvm;
         pub use zkvm::*;
-    } else {
+    }
+    _ => {
         mod unsupported;
         pub use unsupported::*;
     }
diff --git a/library/std/src/sys/stdio/vexos.rs b/library/std/src/sys/stdio/vexos.rs
new file mode 100644
index 00000000000..9a391feb7a8
--- /dev/null
+++ b/library/std/src/sys/stdio/vexos.rs
@@ -0,0 +1,100 @@
+use crate::io;
+
+pub struct Stdin;
+pub struct Stdout;
+pub type Stderr = Stdout;
+
+pub const STDIO_CHANNEL: u32 = 1;
+
+impl Stdin {
+    pub const fn new() -> Stdin {
+        Stdin
+    }
+}
+
+impl io::Read for Stdin {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let mut count = 0;
+
+        for out_byte in buf.iter_mut() {
+            let byte = unsafe { vex_sdk::vexSerialReadChar(STDIO_CHANNEL) };
+            if byte < 0 {
+                break;
+            }
+
+            *out_byte = byte as u8;
+            count += 1;
+        }
+
+        Ok(count)
+    }
+}
+
+impl Stdout {
+    pub const fn new() -> Stdout {
+        Stdout
+    }
+}
+
+impl io::Write for Stdout {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let mut written = 0;
+
+        // HACK: VEXos holds an internal ringbuffer for serial writes that is flushed to USB1
+        // roughly every millisecond by `vexTasksRun`. For writes larger than 2048 bytes, we
+        // must block until that buffer is flushed to USB1 before writing the rest of `buf`.
+        //
+        // This is fairly nonstandard for a `write` implementation, but it avoids a guaranteed
+        // recursive panic when using macros such as `print!` to write large amounts of data
+        // (buf.len() > 2048) to stdout at once.
+        for chunk in buf.chunks(STDOUT_BUF_SIZE) {
+            if unsafe { vex_sdk::vexSerialWriteFree(STDIO_CHANNEL) as usize } < chunk.len() {
+                self.flush().unwrap();
+            }
+
+            let count: usize = unsafe {
+                vex_sdk::vexSerialWriteBuffer(STDIO_CHANNEL, chunk.as_ptr(), chunk.len() as u32)
+            }
+            .try_into()
+            .map_err(|_| {
+                io::const_error!(io::ErrorKind::Uncategorized, "internal write error occurred")
+            })?;
+
+            written += count;
+
+            // This is a sanity check to ensure that we don't end up with non-contiguous
+            // buffer writes. e.g. a chunk gets only partially written, but we continue
+            // attempting to write the remaining chunks.
+            //
+            // In practice, this should never really occur since the previous flush ensures
+            // enough space in FIFO to write the entire chunk to vexSerialWriteBuffer.
+            if count != chunk.len() {
+                break;
+            }
+        }
+
+        Ok(written)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        // This may block for up to a millisecond.
+        unsafe {
+            while (vex_sdk::vexSerialWriteFree(STDIO_CHANNEL) as usize) != STDOUT_BUF_SIZE {
+                vex_sdk::vexTasksRun();
+            }
+        }
+
+        Ok(())
+    }
+}
+
+pub const STDIN_BUF_SIZE: usize = 4096;
+pub const STDOUT_BUF_SIZE: usize = 2048;
+
+pub fn is_ebadf(_err: &io::Error) -> bool {
+    false
+}
+
+pub fn panic_output() -> Option<impl io::Write> {
+    Some(Stdout::new())
+}
diff --git a/library/std/src/sys/stdio/wasi.rs b/library/std/src/sys/stdio/wasip1.rs
index b70efd026f9..b70efd026f9 100644
--- a/library/std/src/sys/stdio/wasi.rs
+++ b/library/std/src/sys/stdio/wasip1.rs
diff --git a/library/std/src/sys/stdio/wasip2.rs b/library/std/src/sys/stdio/wasip2.rs
new file mode 100644
index 00000000000..1fcb49a083d
--- /dev/null
+++ b/library/std/src/sys/stdio/wasip2.rs
@@ -0,0 +1,120 @@
+use wasip2::cli;
+use wasip2::io::streams::{Error, InputStream, OutputStream, StreamError};
+
+use crate::io::{self, BorrowedBuf, BorrowedCursor};
+
+pub struct Stdin(Option<InputStream>);
+pub struct Stdout(Option<OutputStream>);
+pub struct Stderr(Option<OutputStream>);
+
+fn error_to_io(err: Error) -> io::Error {
+    // There exists a function in `wasi:filesystem` to optionally acquire an
+    // error code from an error, but the streams in use in this module are
+    // exclusively used with stdio meaning that a filesystem error is not
+    // possible here.
+    //
+    // In lieu of an error code, which WASIp2 does not specify, this instead
+    // carries along the `to_debug_string` implementation that the host
+    // supplies. If this becomes too expensive in the future this could also
+    // become `io::Error::from_raw_os_error(libc::EIO)` or similar.
+    io::Error::new(io::ErrorKind::Other, err.to_debug_string())
+}
+
+impl Stdin {
+    pub const fn new() -> Stdin {
+        Stdin(None)
+    }
+
+    fn stream(&mut self) -> &InputStream {
+        self.0.get_or_insert_with(cli::stdin::get_stdin)
+    }
+}
+
+impl io::Read for Stdin {
+    fn read(&mut self, data: &mut [u8]) -> io::Result<usize> {
+        let mut buf = BorrowedBuf::from(data);
+        self.read_buf(buf.unfilled())?;
+        Ok(buf.len())
+    }
+
+    fn read_buf(&mut self, mut buf: BorrowedCursor<'_>) -> io::Result<()> {
+        match self.stream().blocking_read(u64::try_from(buf.capacity()).unwrap()) {
+            Ok(result) => {
+                buf.append(&result);
+                Ok(())
+            }
+            Err(StreamError::Closed) => Ok(()),
+            Err(StreamError::LastOperationFailed(e)) => Err(error_to_io(e)),
+        }
+    }
+}
+
+impl Stdout {
+    pub const fn new() -> Stdout {
+        Stdout(None)
+    }
+
+    fn stream(&mut self) -> &OutputStream {
+        self.0.get_or_insert_with(cli::stdout::get_stdout)
+    }
+}
+
+fn write(stream: &OutputStream, buf: &[u8]) -> io::Result<usize> {
+    // WASIp2's `blocking_write_and_flush` function is defined as accepting no
+    // more than 4096 bytes. Larger writes can be issued by manually using
+    // `check_write`, `write`, and `blocking_flush`, but for now just go ahead
+    // and use `blocking_write_and_flush` and report a short write and let a
+    // higher level loop over the result.
+    const MAX: usize = 4096;
+    let buf = &buf[..buf.len().min(MAX)];
+    match stream.blocking_write_and_flush(buf) {
+        Ok(()) => Ok(buf.len()),
+        Err(StreamError::Closed) => Ok(0),
+        Err(StreamError::LastOperationFailed(e)) => Err(error_to_io(e)),
+    }
+}
+
+impl io::Write for Stdout {
+    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
+        write(self.stream(), data)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        // Note that `OutputStream` has a `flush` function but for stdio all
+        // writes are accompanied with a flush which means that this flush
+        // doesn't need to do anything.
+        Ok(())
+    }
+}
+
+impl Stderr {
+    pub const fn new() -> Stderr {
+        Stderr(None)
+    }
+
+    fn stream(&mut self) -> &OutputStream {
+        self.0.get_or_insert_with(cli::stderr::get_stderr)
+    }
+}
+
+impl io::Write for Stderr {
+    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
+        write(self.stream(), data)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        // See `Stdout::flush` for why this is a noop.
+        Ok(())
+    }
+}
+
+pub const STDIN_BUF_SIZE: usize = crate::sys::io::DEFAULT_BUF_SIZE;
+
+pub fn is_ebadf(_err: &io::Error) -> bool {
+    // WASIp2 stdio streams are always available so ebadf never shows up.
+    false
+}
+
+pub fn panic_output() -> Option<impl io::Write> {
+    Some(Stderr::new())
+}