async_hal/io/
copy_buf.rs

1use super::{AsyncBufRead, AsyncWrite};
2use core::{
3    pin::Pin,
4    task::{Context, Poll},
5};
6use futures::{ready, Future};
7
8/// A future that asynchronously copies the entire contents of a reader into a
9/// writer.
10///
11/// This struct is generally created by calling [`copy_buf`][copy_buf]. Please
12/// see the documentation of `copy_buf()` for more details.
13///
14/// [copy_buf]: copy_buf()
15#[derive(Debug)]
16#[must_use = "futures do nothing unless you `.await` or poll them"]
17struct CopyBuf<'a, R: ?Sized, W: ?Sized> {
18    reader: &'a mut R,
19    writer: &'a mut W,
20    amt: u64,
21}
22
23/// Asynchronously copies the entire contents of a reader into a writer.
24///
25/// This function returns a future that will continuously read data from
26/// `reader` and then write it into `writer` in a streaming fashion until
27/// `reader` returns EOF or fails.
28///
29/// On success, the total number of bytes that were copied from `reader` to
30/// `writer` is returned.
31///
32/// This is a [`async_hal::io::copy`] alternative for [`AsyncBufRead`] readers
33/// with no extra buffer allocation, since [`AsyncBufRead`] allow access
34/// to the reader's inner buffer.
35///
36/// [`async_hal::io::copy`]: crate::io::copy
37/// [`AsyncBufRead`]: crate::io::AsyncBufRead
38///
39/// # Errors
40///
41/// The returned future will finish with an error will return an error
42/// immediately if any call to `poll_fill_buf` or `poll_write` returns an
43/// error.
44///
45/// # Examples
46///
47/// ```
48/// use async_hal::io;
49///
50/// # let task = async {
51/// let mut reader: &[u8] = b"hello";
52/// let mut writer = [0; 5];
53///
54/// io::copy_buf(&mut reader, &mut writer.as_mut()).await?;
55///
56/// assert_eq!(b"hello", &writer[..]);
57/// # Ok::<_, void::Void>(())
58/// # };
59/// # futures::pin_mut!(task);
60/// # async_hal::block_on(task, || {}).unwrap();
61/// ```
62pub async fn copy_buf<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> Result<u64, R::Error>
63where
64    R: AsyncBufRead + Unpin + ?Sized,
65    W: AsyncWrite + Unpin + ?Sized,
66    R::Error: From<W::Error>,
67{
68    CopyBuf {
69        reader,
70        writer,
71        amt: 0,
72    }
73    .await
74}
75
76impl<R, W> Future for CopyBuf<'_, R, W>
77where
78    R: AsyncBufRead + Unpin + ?Sized,
79    W: AsyncWrite + Unpin + ?Sized,
80    R::Error: From<W::Error>,
81{
82    type Output = Result<u64, R::Error>;
83
84    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
85        loop {
86            let me = &mut *self;
87            let buffer = ready!(Pin::new(&mut *me.reader).poll_fill_buf(cx))?;
88            if buffer.is_empty() {
89                ready!(Pin::new(&mut *me.writer).poll_flush(cx))?;
90                return Poll::Ready(Ok(self.amt));
91            }
92
93            let i = ready!(Pin::new(&mut *me.writer).poll_write(cx, buffer))?;
94            if i == 0 {
95                todo!()
96                //return Poll::Ready(Err(std::io::ErrorKind::WriteZero.into()));
97            }
98            self.amt += i as u64;
99            Pin::new(&mut *self.reader).consume(i);
100        }
101    }
102}