Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

tokio: add AsyncReadExt::read_to_string #1326

Merged
merged 1 commit into from Jul 19, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
25 changes: 25 additions & 0 deletions tokio/src/io/async_read_ext.rs
Expand Up @@ -2,6 +2,7 @@ use crate::io::copy::{copy, Copy};
use crate::io::read::{read, Read};
use crate::io::read_exact::{read_exact, ReadExact};
use crate::io::read_to_end::{read_to_end, ReadToEnd};
use crate::io::read_to_string::{read_to_string, ReadToString};

use tokio_io::{AsyncRead, AsyncWrite};

Expand Down Expand Up @@ -63,12 +64,36 @@ pub trait AsyncReadExt: AsyncRead {
}

/// Read all bytes until EOF in this source, placing them into `dst`.
///
/// On success the total number of bytes read is returned.
///
/// # Examples
///
/// ```
/// unimplemented!();
/// ```
fn read_to_end<'a>(&'a mut self, dst: &'a mut Vec<u8>) -> ReadToEnd<'a, Self>
where
Self: Unpin,
{
read_to_end(self, dst)
}

/// Read all bytes until EOF in this source, placing them into `dst`.
///
/// On success the total number of bytes read is returned.
///
/// # Examples
///
/// ```
/// unimplemented!();
/// ```
fn read_to_string<'a>(&'a mut self, dst: &'a mut String) -> ReadToString<'a, Self>
where
Self: Unpin,
{
read_to_string(self, dst)
}
}

impl<R: AsyncRead + ?Sized> AsyncReadExt for R {}
1 change: 1 addition & 0 deletions tokio/src/io/mod.rs
Expand Up @@ -45,6 +45,7 @@ mod read;
mod read_exact;
mod read_line;
mod read_to_end;
mod read_to_string;
mod read_until;
mod write;
mod write_all;
Expand Down
72 changes: 72 additions & 0 deletions tokio/src/io/read_to_string.rs
@@ -0,0 +1,72 @@
use super::read_to_end::read_to_end_internal;
use futures_core::ready;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{io, mem, str};
use tokio_io::AsyncRead;

/// Future for the [`read_to_string`](super::AsyncReadExt::read_to_string) method.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct ReadToString<'a, R: ?Sized + Unpin> {
reader: &'a mut R,
buf: &'a mut String,
bytes: Vec<u8>,
start_len: usize,
}

impl<R: ?Sized + Unpin> Unpin for ReadToString<'_, R> {}

pub(crate) fn read_to_string<'a, R>(reader: &'a mut R, buf: &'a mut String) -> ReadToString<'a, R>
where
R: AsyncRead + ?Sized + Unpin,
{
let start_len = buf.len();
ReadToString {
reader,
bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) },
buf,
start_len,
}
}

fn read_to_string_internal<R: AsyncRead + ?Sized>(
reader: Pin<&mut R>,
cx: &mut Context<'_>,
buf: &mut String,
bytes: &mut Vec<u8>,
start_len: usize,
) -> Poll<io::Result<usize>> {
let ret = ready!(read_to_end_internal(reader, cx, bytes, start_len));
if str::from_utf8(&bytes).is_err() {
Poll::Ready(ret.and_then(|_| {
Err(io::Error::new(
io::ErrorKind::InvalidData,
"stream did not contain valid UTF-8",
))
}))
} else {
debug_assert!(buf.is_empty());
// Safety: `bytes` is a valid UTF-8 because `str::from_utf8` returned `Ok`.
mem::swap(unsafe { buf.as_mut_vec() }, bytes);
Poll::Ready(ret)
}
}

impl<A> Future for ReadToString<'_, A>
where
A: AsyncRead + ?Sized + Unpin,
{
type Output = io::Result<usize>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self {
reader,
buf,
bytes,
start_len,
} = &mut *self;
read_to_string_internal(Pin::new(reader), cx, buf, bytes, *start_len)
}
}
40 changes: 40 additions & 0 deletions tokio/tests/io_read_to_string.rs
@@ -0,0 +1,40 @@
#![deny(warnings, rust_2018_idioms)]
#![feature(async_await)]

use tokio::io::{AsyncRead, AsyncReadExt};
use tokio_test::assert_ok;

use std::pin::Pin;
use std::task::{Context, Poll};
use std::{cmp, io};

#[tokio::test]
async fn read_to_string() {
struct Rd {
val: &'static [u8],
}

impl AsyncRead for Rd {
fn poll_read(
mut self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
let me = &mut *self;
let len = cmp::min(buf.len(), me.val.len());

buf[..len].copy_from_slice(&me.val[..len]);
me.val = &me.val[len..];
Poll::Ready(Ok(len))
}
}

let mut buf = String::new();
let mut rd = Rd {
val: b"hello world",
};

let n = assert_ok!(rd.read_to_string(&mut buf).await);
assert_eq!(n, 11);
assert_eq!(buf[..], "hello world"[..]);
}