| 1 | use bytes::BufMut; |
| 2 | use std::future::Future; |
| 3 | use std::io; |
| 4 | use std::pin::Pin; |
| 5 | use std::task::{Context, Poll}; |
| 6 | use tokio::io::AsyncRead; |
| 7 | |
| 8 | /// Read data from an `AsyncRead` into an implementer of the [`BufMut`] trait. |
| 9 | /// |
| 10 | /// [`BufMut`]: bytes::BufMut |
| 11 | /// |
| 12 | /// # Example |
| 13 | /// |
| 14 | /// ``` |
| 15 | /// use bytes::{Bytes, BytesMut}; |
| 16 | /// use tokio_stream as stream; |
| 17 | /// use tokio::io::Result; |
| 18 | /// use tokio_util::io::{StreamReader, read_buf}; |
| 19 | /// # #[tokio::main] |
| 20 | /// # async fn main() -> std::io::Result<()> { |
| 21 | /// |
| 22 | /// // Create a reader from an iterator. This particular reader will always be |
| 23 | /// // ready. |
| 24 | /// let mut read = StreamReader::new(stream::iter(vec![Result::Ok(Bytes::from_static(&[0, 1, 2, 3]))])); |
| 25 | /// |
| 26 | /// let mut buf = BytesMut::new(); |
| 27 | /// let mut reads = 0; |
| 28 | /// |
| 29 | /// loop { |
| 30 | /// reads += 1; |
| 31 | /// let n = read_buf(&mut read, &mut buf).await?; |
| 32 | /// |
| 33 | /// if n == 0 { |
| 34 | /// break; |
| 35 | /// } |
| 36 | /// } |
| 37 | /// |
| 38 | /// // one or more reads might be necessary. |
| 39 | /// assert!(reads >= 1); |
| 40 | /// assert_eq!(&buf[..], &[0, 1, 2, 3]); |
| 41 | /// # Ok(()) |
| 42 | /// # } |
| 43 | /// ``` |
| 44 | pub async fn read_buf<R, B>(read: &mut R, buf: &mut B) -> io::Result<usize> |
| 45 | where |
| 46 | R: AsyncRead + Unpin, |
| 47 | B: BufMut, |
| 48 | { |
| 49 | return ReadBufFn(read, buf).await; |
| 50 | |
| 51 | struct ReadBufFn<'a, R, B>(&'a mut R, &'a mut B); |
| 52 | |
| 53 | impl<'a, R, B> Future for ReadBufFn<'a, R, B> |
| 54 | where |
| 55 | R: AsyncRead + Unpin, |
| 56 | B: BufMut, |
| 57 | { |
| 58 | type Output = io::Result<usize>; |
| 59 | |
| 60 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
| 61 | let this: &mut ReadBufFn<'a, R, B> = &mut *self; |
| 62 | crate::util::poll_read_buf(io:Pin::new(this.0), cx, buf:this.1) |
| 63 | } |
| 64 | } |
| 65 | } |
| 66 | |