wasmcloud-component 0.2.0

use std::io::{Read, Write};

pub struct InputStreamReader<'a> {
    stream: &'a mut ::wasi::io::streams::InputStream,
}

impl<'a> From<&'a mut ::wasi::io::streams::InputStream> for InputStreamReader<'a> {
    fn from(stream: &'a mut ::wasi::io::streams::InputStream) -> Self {
        Self { stream }
    }
}

impl std::io::Read for InputStreamReader<'_> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.stream.read(buf)
    }
}

pub struct OutputStreamWriter<'a> {
    stream: &'a mut ::wasi::io::streams::OutputStream,
}

impl<'a> From<&'a mut ::wasi::io::streams::OutputStream> for OutputStreamWriter<'a> {
    fn from(stream: &'a mut ::wasi::io::streams::OutputStream) -> Self {
        Self { stream }
    }
}

impl std::io::Write for OutputStreamWriter<'_> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.stream.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.stream.flush()
    }
}

pub struct StdioStream<'a> {
    stdin: std::io::StdinLock<'a>,
    stdout: std::io::StdoutLock<'a>,
}

impl StdioStream<'_> {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

impl Read for StdioStream<'_> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.stdin.read(buf)
    }
}

impl Write for StdioStream<'_> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.stdout.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.stdout.flush()
    }
}

impl Default for StdioStream<'_> {
    fn default() -> Self {
        Self {
            stdin: std::io::stdin().lock(),
            stdout: std::io::stdout().lock(),
        }
    }
}

/// Similar to [`crate::wasi::io::poll::poll`], but polls all `pollables` until they are all ready.
///
/// Poll for completion on a set of pollables.
///
/// This function takes a list of pollables, which identify I/O sources of interest, and waits until all of the events are ready for I/O.
pub fn join(pollables: &[&crate::wasi::io::poll::Pollable]) {
    let mut pollables = pollables.to_vec();
    while !pollables.is_empty() {
        let ready_indices = crate::wasi::io::poll::poll(&pollables);
        ready_indices.iter().rev().for_each(|&i| {
            pollables.swap_remove(i as usize);
        });
    }
}

#[cfg(feature = "futures")]
impl futures::AsyncRead for StdioStream<'_> {
    fn poll_read(
        mut self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
        buf: &mut [u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        std::task::Poll::Ready(self.stdin.read(buf))
    }
}

#[cfg(feature = "futures")]
impl futures::AsyncWrite for StdioStream<'_> {
    fn poll_write(
        mut self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        std::task::Poll::Ready(self.stdout.write(buf))
    }

    fn poll_flush(
        mut self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        std::task::Poll::Ready(self.stdout.flush())
    }

    fn poll_close(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        self.poll_flush(cx)
    }
}

#[cfg(feature = "tokio")]
impl tokio::io::AsyncRead for StdioStream<'_> {
    fn poll_read(
        mut self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        let mut fill = vec![0; buf.capacity()];
        std::task::Poll::Ready({
            let n = self.stdin.read(&mut fill)?;
            buf.put_slice(&fill[..n]);
            Ok(())
        })
    }
}

#[cfg(feature = "tokio")]
impl tokio::io::AsyncWrite for StdioStream<'_> {
    fn poll_write(
        mut self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<Result<usize, std::io::Error>> {
        std::task::Poll::Ready(self.stdout.write(buf))
    }

    fn poll_flush(
        mut self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        std::task::Poll::Ready(self.stdout.flush())
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        self.poll_flush(cx)
    }
}