kona_std_fpvm/
channel.rs

1//! This module contains a rudamentary channel between two file descriptors, using [crate::io]
2//! for reading and writing from the file descriptors.
3
4use crate::{FileDescriptor, io};
5use alloc::boxed::Box;
6use async_trait::async_trait;
7use core::{
8    cell::RefCell,
9    cmp::Ordering,
10    future::Future,
11    pin::Pin,
12    task::{Context, Poll},
13};
14use kona_preimage::{
15    Channel,
16    errors::{ChannelError, ChannelResult},
17};
18
19/// [FileChannel] is a handle for one end of a bidirectional channel.
20#[derive(Debug, Clone, Copy)]
21pub struct FileChannel {
22    /// File descriptor to read from
23    read_handle: FileDescriptor,
24    /// File descriptor to write to
25    write_handle: FileDescriptor,
26}
27
28impl FileChannel {
29    /// Create a new [FileChannel] from two file descriptors.
30    pub const fn new(read_handle: FileDescriptor, write_handle: FileDescriptor) -> Self {
31        Self { read_handle, write_handle }
32    }
33
34    /// Returns the a copy of the [FileDescriptor] used for the read end of the channel.
35    pub const fn read_handle(&self) -> FileDescriptor {
36        self.read_handle
37    }
38
39    /// Returns the a copy of the [FileDescriptor] used for the write end of the channel.
40    pub const fn write_handle(&self) -> FileDescriptor {
41        self.write_handle
42    }
43}
44
45#[async_trait]
46impl Channel for FileChannel {
47    async fn read(&self, buf: &mut [u8]) -> ChannelResult<usize> {
48        io::read(self.read_handle, buf).map_err(|_| ChannelError::Closed)
49    }
50
51    async fn read_exact(&self, buf: &mut [u8]) -> ChannelResult<usize> {
52        ReadFuture::new(*self, buf).await.map_err(|_| ChannelError::Closed)
53    }
54
55    async fn write(&self, buf: &[u8]) -> ChannelResult<usize> {
56        WriteFuture::new(*self, buf).await.map_err(|_| ChannelError::Closed)
57    }
58}
59
60/// A future that reads from a channel, returning [Poll::Ready] when the buffer is full.
61struct ReadFuture<'a> {
62    /// The channel to read from
63    channel: FileChannel,
64    /// The buffer to read into
65    buf: RefCell<&'a mut [u8]>,
66    /// The number of bytes read so far
67    read: usize,
68}
69
70impl<'a> ReadFuture<'a> {
71    /// Create a new [ReadFuture] from a channel and a buffer.
72    #[allow(clippy::missing_const_for_fn)]
73    fn new(channel: FileChannel, buf: &'a mut [u8]) -> Self {
74        Self { channel, buf: RefCell::new(buf), read: 0 }
75    }
76}
77
78impl Future for ReadFuture<'_> {
79    type Output = ChannelResult<usize>;
80
81    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
82        let mut buf = self.buf.borrow_mut();
83        let buf_len = buf.len();
84        let chunk_read = io::read(self.channel.read_handle, &mut buf[self.read..])
85            .map_err(|_| ChannelError::Closed)?;
86
87        // Drop the borrow on self.
88        drop(buf);
89
90        self.read += chunk_read;
91
92        match self.read.cmp(&buf_len) {
93            Ordering::Greater | Ordering::Equal => Poll::Ready(Ok(self.read)),
94            Ordering::Less => {
95                // Register the current task to be woken up when it can make progress
96                ctx.waker().wake_by_ref();
97                Poll::Pending
98            }
99        }
100    }
101}
102
103/// A future that writes to a channel, returning [Poll::Ready] when the full buffer has been
104/// written.
105struct WriteFuture<'a> {
106    /// The channel to write to
107    channel: FileChannel,
108    /// The buffer to write
109    buf: &'a [u8],
110    /// The number of bytes written so far
111    written: usize,
112}
113
114impl<'a> WriteFuture<'a> {
115    /// Create a new [WriteFuture] from a channel and a buffer.
116    const fn new(channel: FileChannel, buf: &'a [u8]) -> Self {
117        Self { channel, buf, written: 0 }
118    }
119}
120
121impl Future for WriteFuture<'_> {
122    type Output = ChannelResult<usize>;
123
124    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
125        match io::write(self.channel.write_handle(), &self.buf[self.written..]) {
126            Ok(0) => Poll::Ready(Ok(self.written)), // Finished writing
127            Ok(n) => {
128                self.written += n;
129
130                if self.written >= self.buf.len() {
131                    return Poll::Ready(Ok(self.written));
132                }
133
134                // Register the current task to be woken up when it can make progress
135                ctx.waker().wake_by_ref();
136                Poll::Pending
137            }
138            Err(_) => Poll::Ready(Err(ChannelError::Closed)),
139        }
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn test_get_read_handle() {
149        let read_handle = FileDescriptor::StdIn;
150        let write_handle = FileDescriptor::StdOut;
151        let chan = FileChannel::new(read_handle, write_handle);
152        let ref_read_handle = chan.read_handle();
153        assert_eq!(read_handle, ref_read_handle);
154    }
155}
kona_std_fpvm/channel.rs

kona_std_fpvm/
channel.rs
