Skip to content

Commit

Permalink
Change read_to_{end, string} to return the total number of bytes read
Browse files Browse the repository at this point in the history
  • Loading branch information
taiki-e authored and cramertj committed Jul 8, 2019
1 parent a762a0d commit 6e4eddb
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 15 deletions.
10 changes: 8 additions & 2 deletions futures-util/src/io/mod.rs
Expand Up @@ -219,6 +219,8 @@ pub trait AsyncReadExt: AsyncRead {

/// Creates a future which will read all the bytes from this `AsyncRead`.
///
/// On success the total number of bytes read is returned.
///
/// # Examples
///
/// ```
Expand All @@ -230,8 +232,9 @@ pub trait AsyncReadExt: AsyncRead {
/// let mut reader = Cursor::new([1, 2, 3, 4]);
/// let mut output = Vec::with_capacity(4);
///
/// reader.read_to_end(&mut output).await?;
/// let bytes = reader.read_to_end(&mut output).await?;
///
/// assert_eq!(bytes, 4);
/// assert_eq!(output, vec![1, 2, 3, 4]);
/// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap();
/// ```
Expand All @@ -246,6 +249,8 @@ pub trait AsyncReadExt: AsyncRead {

/// Creates a future which will read all the bytes from this `AsyncRead`.
///
/// On success the total number of bytes read is returned.
///
/// # Examples
///
/// ```
Expand All @@ -257,8 +262,9 @@ pub trait AsyncReadExt: AsyncRead {
/// let mut reader = Cursor::new(&b"1234"[..]);
/// let mut buffer = String::with_capacity(4);
///
/// reader.read_to_string(&mut buffer).await?;
/// let bytes = reader.read_to_string(&mut buffer).await?;
///
/// assert_eq!(bytes, 4);
/// assert_eq!(buffer, String::from("1234"));
/// # Ok::<(), Box<dyn std::error::Error>>(()) }).unwrap();
/// ```
Expand Down
17 changes: 12 additions & 5 deletions futures-util/src/io/read_to_end.rs
Expand Up @@ -11,13 +11,19 @@ use std::vec::Vec;
pub struct ReadToEnd<'a, R: ?Sized + Unpin> {
reader: &'a mut R,
buf: &'a mut Vec<u8>,
start_len: usize,
}

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

impl<'a, R: AsyncRead + ?Sized + Unpin> ReadToEnd<'a, R> {
pub(super) fn new(reader: &'a mut R, buf: &'a mut Vec<u8>) -> Self {
ReadToEnd { reader, buf }
let start_len = buf.len();
Self {
reader,
buf,
start_len,
}
}
}

Expand All @@ -42,7 +48,8 @@ pub(super) fn read_to_end_internal<R: AsyncRead + ?Sized>(
mut rd: Pin<&mut R>,
cx: &mut Context<'_>,
buf: &mut Vec<u8>,
) -> Poll<io::Result<()>> {
start_len: usize,
) -> Poll<io::Result<usize>> {
let mut g = Guard { len: buf.len(), buf };
let ret;
loop {
Expand All @@ -57,7 +64,7 @@ pub(super) fn read_to_end_internal<R: AsyncRead + ?Sized>(

match ready!(rd.as_mut().poll_read(cx, &mut g.buf[g.len..])) {
Ok(0) => {
ret = Poll::Ready(Ok(()));
ret = Poll::Ready(Ok(g.len - start_len));
break;
}
Ok(n) => g.len += n,
Expand All @@ -74,10 +81,10 @@ pub(super) fn read_to_end_internal<R: AsyncRead + ?Sized>(
impl<A> Future for ReadToEnd<'_, A>
where A: AsyncRead + ?Sized + Unpin,
{
type Output = io::Result<()>;
type Output = io::Result<usize>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
read_to_end_internal(Pin::new(&mut this.reader), cx, this.buf)
read_to_end_internal(Pin::new(&mut this.reader), cx, this.buf, this.start_len)
}
}
14 changes: 9 additions & 5 deletions futures-util/src/io/read_to_string.rs
Expand Up @@ -13,16 +13,19 @@ 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> {}

impl<'a, R: AsyncRead + ?Sized + Unpin> ReadToString<'a, R> {
pub(super) fn new(reader: &'a mut R, buf: &'a mut String) -> Self {
let start_len = buf.len();
Self {
reader,
bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) },
buf,
start_len,
}
}
}
Expand All @@ -32,8 +35,9 @@ fn read_to_string_internal<R: AsyncRead + ?Sized>(
cx: &mut Context<'_>,
buf: &mut String,
bytes: &mut Vec<u8>,
) -> Poll<io::Result<()>> {
let ret = ready!(read_to_end_internal(reader, cx, bytes));
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(
Expand All @@ -53,10 +57,10 @@ impl<A> Future for ReadToString<'_, A>
where
A: AsyncRead + ?Sized + Unpin,
{
type Output = io::Result<()>;
type Output = io::Result<usize>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self { reader, buf, bytes } = &mut *self;
read_to_string_internal(Pin::new(reader), cx, buf, bytes)
let Self { reader, buf, bytes, start_len } = &mut *self;
read_to_string_internal(Pin::new(reader), cx, buf, bytes, *start_len)
}
}
6 changes: 3 additions & 3 deletions futures/tests/io_read_to_string.rs
Expand Up @@ -11,12 +11,12 @@ use std::io::Cursor;
fn read_to_string() {
let mut c = Cursor::new(&b""[..]);
let mut v = String::new();
assert!(block_on(c.read_to_string(&mut v)).is_ok());
assert_eq!(block_on(c.read_to_string(&mut v)).unwrap(), 0);
assert_eq!(v, "");

let mut c = Cursor::new(&b"1"[..]);
let mut v = String::new();
assert!(block_on(c.read_to_string(&mut v)).is_ok());
assert_eq!(block_on(c.read_to_string(&mut v)).unwrap(), 1);
assert_eq!(v, "1");

let mut c = Cursor::new(&b"\xff"[..]);
Expand All @@ -41,6 +41,6 @@ fn interleave_pending() {
.interleave_pending();

let mut v = String::new();
assert!(run(buf.read_to_string(&mut v)).is_ok());
assert_eq!(run(buf.read_to_string(&mut v)).unwrap(), 5);
assert_eq!(v, "12333");
}

0 comments on commit 6e4eddb

Please sign in to comment.