Repository: tokio-rs/tokio
File: tokio/src/io/util/read_to_end.rs (Rust)

Code before:
    // amount of data to return. When the vector is full with its starting
    // capacity, we first try to read into a small buffer to see if we reached
    // an EOF. This avoids the unnecessary allocation that we attempt before
    // reading into the vector.

    const NUM_BYTES: usize = 32;
    let try_small_read = buf.try_small_read_first(NUM_BYTES);

    // Get a ReadBuf into the vector.
    let mut read_buf = buf.get_read_buf();
    let filled_before = read_buf.filled().len();
    let poll_result;

    let filled_after = if try_small_read {
        let mut small_buf: [MaybeUninit<u8>; NUM_BYTES] = [MaybeUninit::uninit(); NUM_BYTES];
        let mut small_read_buf = ReadBuf::uninit(&mut small_buf);
        poll_result = read.poll_read(cx, &mut small_read_buf);
        let filled = small_read_buf.filled();
        if filled.len() > 0 {
            buf.reserve(NUM_BYTES);
            read_buf = buf.get_read_buf();
            read_buf.put_slice(filled);
        }
        filled_before + filled.len()
    } else {
        buf.reserve(NUM_BYTES);
        read_buf = buf.get_read_buf();
        poll_result = read.poll_read(cx, &mut read_buf);
        read_buf.filled().len()
    };

    // Update the length of the vector using the result of poll_read.
    let read_buf_parts = into_read_buf_parts(read_buf);
    buf.apply_read_buf(read_buf_parts);
    let n = filled_after - filled_before;

    match poll_result {
        Poll::Pending => {
            // In this case, nothing should have been read. However we still
            // update the vector in case the poll_read call initialized parts of
            // the vector's unused capacity.
            debug_assert_eq!(filled_before, filled_after);
            Poll::Pending
        }
        Poll::Ready(Err(err)) => {
            debug_assert_eq!(filled_before, filled_after);
            Poll::Ready(Err(err))
        }
        Poll::Ready(Ok(())) => Poll::Ready(Ok(n)),
    }
}

Review comment (bug):
I'm pretty sure that this test would fail if you have a `reserve` here.
```
#[tokio::test]
async fn two_small_reads() {
    let bytes = b"imlargerthan32bytessoIcanhelpwiththetest";
    let mut mock = Builder::new().read(&bytes[..20]).read(&bytes[20..]).build();
    let mut buf = Vec::with_capacity(bytes.len());
    AsyncReadExt::read_to_end(&mut mock, &mut buf)
        .await
        .unwrap();
    assert_eq!(bytes.len(), buf.capacity());
}
```

Code after:
    // amount of data to return. When the vector is full with its starting
    // capacity, we first try to read into a small buffer to see if we reached
    // an EOF. This only happens when the starting capacity is >= NUM_BYTES, since
    // we allocate at least NUM_BYTES each time. This avoids the unnecessary
    // allocation that we attempt before reading into the vector.

    const NUM_BYTES: usize = 32;
    let try_small_read = buf.try_small_read_first(NUM_BYTES);

    // Get a ReadBuf into the vector.
    let mut read_buf;
    let poll_result;

    let n = if try_small_read {
        // Read some bytes using a small read.
        let mut small_buf: [MaybeUninit<u8>; NUM_BYTES] = [MaybeUninit::uninit(); NUM_BYTES];
        let mut small_read_buf = ReadBuf::uninit(&mut small_buf);
        poll_result = read.poll_read(cx, &mut small_read_buf);
        let to_write = small_read_buf.filled();

        // Ensure we have enough space to fill our vector with what we read.
        read_buf = buf.get_read_buf();
        if to_write.len() > read_buf.remaining() {
            buf.reserve(NUM_BYTES);
            read_buf = buf.get_read_buf();
        }
        read_buf.put_slice(to_write);

        to_write.len()
    } else {
        // Ensure we have enough space for reading.
        buf.reserve(NUM_BYTES);
        read_buf = buf.get_read_buf();

        // Read data directly into vector.
        let filled_before = read_buf.filled().len();
        poll_result = read.poll_read(cx, &mut read_buf);

        // Compute the number of bytes read.
        read_buf.filled().len() - filled_before
    };

    // Update the length of the vector using the result of poll_read.
    let read_buf_parts = into_read_buf_parts(read_buf);
    buf.apply_read_buf(read_buf_parts);

    match poll_result {
        Poll::Pending => {
            // In this case, nothing should have been read. However we still
            // update the vector in case the poll_read call initialized parts of
            // the vector's unused capacity.