Implement io module for WASIp3

examples/stdio.rs

@@ -0,0 +1,84 @@
+#![cfg_attr(not(target_os = "wasi"), no_main)]
+#![cfg(target_os = "wasi")]
+
+//! Demonstrates async line-oriented stdin, stdout, stderr, and flushing.
+
+use anyhow::Result;
+use wstd::io::{AsyncRead, AsyncWrite};
+
+struct LineReader<R> {
+    reader: R,
+    buffer: Vec<u8>,
+}
+
+impl<R: AsyncRead> LineReader<R> {
+    fn new(reader: R) -> Self {
+        Self {
+            reader,
+            buffer: Vec::new(),
+        }
+    }
+
+    async fn read_line(&mut self) -> std::io::Result<Option<Vec<u8>>> {
+        loop {
+            if let Some(newline) = self.buffer.iter().position(|byte| *byte == b'\n') {
+                return Ok(Some(self.buffer.drain(..=newline).collect()));
+            }
+
+            let mut chunk = [0; 1024];
+            match self.reader.read(&mut chunk).await? {
+                0 if self.buffer.is_empty() => return Ok(None),
+                0 => return Ok(Some(std::mem::take(&mut self.buffer))),
+                len => {
+                    self.buffer.extend_from_slice(&chunk[..len]);
+                }
+            }
+        }
+    }
+}
+
+#[wstd::main]
+async fn main() -> Result<()> {
+    let mut stdin = LineReader::new(wstd::io::stdin());
+    let mut stdout = wstd::io::stdout();
+    let mut stderr = wstd::io::stderr();
+
+    let mut stdin_error = None;
+    while let Some(line) = match stdin.read_line().await {
+        Ok(Some(line)) => Some(line),
+        Ok(None) => {
+            stdin_error = Some(std::io::ErrorKind::UnexpectedEof.into());
+            None
+        }
+        Err(error) => {
+            stdin_error = Some(error);
+            None
+        }
+    } {
+        stdout.write_all(b"stdout: ").await.unwrap();
+        stdout.write_all(&line).await.unwrap();
+        stdout.flush().await.unwrap();
+
+        stderr.write_all(b"stderr: ").await.unwrap();
+        stderr.write_all(&line).await.unwrap();
+        stderr.flush().await.unwrap();
+    }
+    let stdin_error = stdin_error.expect("the loop only exits after a stdin error");
+
+    let stdout_result = match stdout.write_all(b"stdin closed\n").await {
+        Ok(()) => stdout.flush().await,
+        Err(error) => Err(error),
+    };
+
+    stderr
+        .write_all(format!("stdin error: {:?}\n", stdin_error.kind()).as_bytes())
+        .await?;
+    if let Err(error) = stdout_result {
+        stderr
+            .write_all(format!("stdout error: {:?}\n", error.kind()).as_bytes())
+            .await?;
+    }
+    stderr.flush().await?;
+
+    Ok(())
+}

src/io/mod.rs

@@ -6,9 +6,17 @@ mod empty;
 mod read;
 mod seek;
 mod stdio;
-mod streams;
+#[cfg(target_env = "p2")]
+mod streams_p2;
+#[cfg(target_env = "p2")]
+use streams_p2 as streams;
+#[cfg(target_env = "p3")]
+mod streams_p3;
+#[cfg(target_env = "p3")]
+use streams_p3 as streams;
 mod write;
 
+#[cfg(target_env = "p2")]
 pub use crate::runtime::AsyncPollable;
 pub use copy::*;
 pub use cursor::*;

src/io/stdio.rs

@@ -1,16 +1,40 @@
 use super::{AsyncInputStream, AsyncOutputStream, AsyncRead, AsyncWrite, Result};
 use std::cell::LazyCell;
-use wasip2::cli::terminal_input::TerminalInput;
-use wasip2::cli::terminal_output::TerminalOutput;
+#[cfg(target_env = "p3")]
+use std::pin::Pin;
+
+#[cfg(target_env = "p2")]
+use wasip2::cli::{terminal_input::TerminalInput, terminal_output::TerminalOutput};
+#[cfg(target_env = "p3")]
+use wasip3::{
+    cli::{terminal_input::TerminalInput, terminal_output::TerminalOutput, types::ErrorCode},
+    wit_bindgen::FutureRead,
+};
+
+#[cfg(target_env = "p3")]
+type Completion = FutureRead<std::result::Result<(), ErrorCode>>;
+
+#[cfg(target_env = "p3")]
+fn into_io_error(error: ErrorCode) -> std::io::Error {
+    let kind = match error {
+        ErrorCode::Io => std::io::ErrorKind::Other,
+        ErrorCode::IllegalByteSequence => std::io::ErrorKind::InvalidData,
+        ErrorCode::Pipe => std::io::ErrorKind::BrokenPipe,
+    };
+    std::io::Error::new(kind, format!("WASI CLI error: {error:?}"))
+}
 
 /// Use the program's stdin as an `AsyncInputStream`.
-#[derive(Debug)]
+#[cfg_attr(target_env = "p2", derive(Debug))]
 pub struct Stdin {
     stream: AsyncInputStream,
+    #[cfg(target_env = "p3")]
+    completion: Pin<Box<Completion>>,
     terminput: LazyCell<Option<TerminalInput>>,
 }
 
 /// Get the program's stdin for use as an `AsyncInputStream`.
+#[cfg(target_env = "p2")]
 pub fn stdin() -> Stdin {
     let stream = AsyncInputStream::new(wasip2::cli::stdin::get_stdin());
     Stdin {
@@ -19,6 +43,16 @@ pub fn stdin() -> Stdin {
     }
 }
 
+#[cfg(target_env = "p3")]
+pub fn stdin() -> Stdin {
+    let (stream, completion) = wasip3::cli::stdin::read_via_stream();
+    Stdin {
+        stream: AsyncInputStream::new(stream),
+        completion: Box::pin(completion.into_future()),
+        terminput: LazyCell::new(wasip3::cli::terminal_stdin::get_terminal_stdin),
+    }
+}
+
 impl Stdin {
     /// Check if stdin is a terminal.
     pub fn is_terminal(&self) -> bool {
@@ -29,15 +63,39 @@ impl Stdin {
     pub fn into_inner(self) -> AsyncInputStream {
         self.stream
     }
+
+    #[cfg(target_env = "p3")]
+    async fn check_error(&mut self) -> Result<()> {
+        let (stream, completion) = wasip3::cli::stdin::read_via_stream();
+        let old_stream = std::mem::replace(&mut self.stream, AsyncInputStream::new(stream));
+        drop(old_stream);
+        let result = self.completion.as_mut().await.map_err(into_io_error);
+        self.completion = Box::pin(completion.into_future());
+        result
+    }
+}
+
+#[cfg(target_env = "p3")]
+impl std::fmt::Debug for Stdin {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Stdin")
+            .field("stream", &self.stream)
+            .field("terminput", &self.terminput)
+            .finish_non_exhaustive()
+    }
 }
 
 impl AsyncRead for Stdin {
     #[inline]
     async fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
-        self.stream.read(buf).await
+        let read = self.stream.read(buf).await?;
+        #[cfg(target_env = "p3")]
+        if read == 0 && buf.is_empty() {
+            self.check_error().await?;
+        }
+        Ok(read)
     }
 
-    #[inline]
     async fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
         self.stream.read_to_end(buf).await
     }
@@ -49,13 +107,16 @@ impl AsyncRead for Stdin {
 }
 
 /// Use the program's stdout as an `AsyncOutputStream`.
-#[derive(Debug)]
+#[cfg_attr(target_env = "p2", derive(Debug))]
 pub struct Stdout {
     stream: AsyncOutputStream,
+    #[cfg(target_env = "p3")]
+    completion: Pin<Box<Completion>>,
     termoutput: LazyCell<Option<TerminalOutput>>,
 }
 
 /// Get the program's stdout for use as an `AsyncOutputStream`.
+#[cfg(target_env = "p2")]
 pub fn stdout() -> Stdout {
     let stream = AsyncOutputStream::new(wasip2::cli::stdout::get_stdout());
     Stdout {
@@ -64,6 +125,17 @@ pub fn stdout() -> Stdout {
     }
 }
 
+#[cfg(target_env = "p3")]
+pub fn stdout() -> Stdout {
+    let (tx, rx) = wasip3::wit_stream::new();
+    let completion = wasip3::cli::stdout::write_via_stream(rx);
+    Stdout {
+        stream: AsyncOutputStream::new(tx),
+        completion: Box::pin(completion.into_future()),
+        termoutput: LazyCell::new(wasip3::cli::terminal_stdout::get_terminal_stdout),
+    }
+}
+
 impl Stdout {
     /// Check if stdout is a terminal.
     pub fn is_terminal(&self) -> bool {
@@ -74,6 +146,26 @@ impl Stdout {
     pub fn into_inner(self) -> AsyncOutputStream {
         self.stream
     }
+
+    #[cfg(target_env = "p3")]
+    async fn flush(&mut self) -> Result<()> {
+        let Self {
+            stream, completion, ..
+        } = std::mem::replace(self, stdout());
+        drop(stream);
+
+        completion.await.map_err(into_io_error)
+    }
+}
+
+#[cfg(target_env = "p3")]
+impl std::fmt::Debug for Stdout {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Stdout")
+            .field("stream", &self.stream)
+            .field("termoutput", &self.termoutput)
+            .finish_non_exhaustive()
+    }
 }
 
 impl AsyncWrite for Stdout {
@@ -84,7 +176,14 @@ impl AsyncWrite for Stdout {
 
     #[inline]
     async fn flush(&mut self) -> Result<()> {
-        self.stream.flush().await
+        #[cfg(target_env = "p2")]
+        {
+            self.stream.flush().await
+        }
+        #[cfg(target_env = "p3")]
+        {
+            Self::flush(self).await
+        }
     }
 
     #[inline]
@@ -99,13 +198,16 @@ impl AsyncWrite for Stdout {
 }
 
 /// Use the program's stdout as an `AsyncOutputStream`.
-#[derive(Debug)]
+#[cfg_attr(target_env = "p2", derive(Debug))]
 pub struct Stderr {
     stream: AsyncOutputStream,
+    #[cfg(target_env = "p3")]
+    completion: Pin<Box<Completion>>,
     termoutput: LazyCell<Option<TerminalOutput>>,
 }
 
 /// Get the program's stdout for use as an `AsyncOutputStream`.
+#[cfg(target_env = "p2")]
 pub fn stderr() -> Stderr {
     let stream = AsyncOutputStream::new(wasip2::cli::stderr::get_stderr());
     Stderr {
@@ -114,6 +216,17 @@ pub fn stderr() -> Stderr {
     }
 }
 
+#[cfg(target_env = "p3")]
+pub fn stderr() -> Stderr {
+    let (tx, rx) = wasip3::wit_stream::new();
+    let completion = wasip3::cli::stderr::write_via_stream(rx);
+    Stderr {
+        stream: AsyncOutputStream::new(tx),
+        completion: Box::pin(completion.into_future()),
+        termoutput: LazyCell::new(wasip3::cli::terminal_stderr::get_terminal_stderr),
+    }
+}
+
 impl Stderr {
     /// Check if stderr is a terminal.
     pub fn is_terminal(&self) -> bool {
@@ -124,6 +237,26 @@ impl Stderr {
     pub fn into_inner(self) -> AsyncOutputStream {
         self.stream
     }
+
+    #[cfg(target_env = "p3")]
+    async fn flush(&mut self) -> Result<()> {
+        let Self {
+            stream, completion, ..
+        } = std::mem::replace(self, stderr());
+        drop(stream);
+
+        completion.await.map_err(into_io_error)
+    }
+}
+
+#[cfg(target_env = "p3")]
+impl std::fmt::Debug for Stderr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Stderr")
+            .field("stream", &self.stream)
+            .field("termoutput", &self.termoutput)
+            .finish_non_exhaustive()
+    }
 }
 
 impl AsyncWrite for Stderr {
@@ -134,7 +267,14 @@ impl AsyncWrite for Stderr {
 
     #[inline]
     async fn flush(&mut self) -> Result<()> {
-        self.stream.flush().await
+        #[cfg(target_env = "p2")]
+        {
+            self.stream.flush().await
+        }
+        #[cfg(target_env = "p3")]
+        {
+            Self::flush(self).await
+        }
     }
 
     #[inline]
@@ -162,18 +302,20 @@ mod test {
                 .write_all(format!("hello, world! stdout {term} a terminal\n",).as_bytes())
                 .await
                 .unwrap();
+            stdout.flush().await.unwrap();
         })
     }
     #[test]
     // No internal predicate. Run test with --nocapture and inspect output manually.
     fn stderr_println_hello_world() {
         block_on(async {
-            let mut stdout = super::stdout();
-            let term = if stdout.is_terminal() { "is" } else { "is not" };
-            stdout
+            let mut stderr = super::stderr();
+            let term = if stderr.is_terminal() { "is" } else { "is not" };
+            stderr
                 .write_all(format!("hello, world! stderr {term} a terminal\n",).as_bytes())
                 .await
                 .unwrap();
+            stderr.flush().await.unwrap();
         })
     }
 }