Skip to content

Commit

Permalink
additional io::copy specializations
Browse files Browse the repository at this point in the history
- copying from `&[u8]` and `VecDeque<u8>`
- copying to `Vec<u8>`
  • Loading branch information
the8472 committed Jul 8, 2023
1 parent d1389b9 commit 6f8ba51
Show file tree
Hide file tree
Showing 2 changed files with 116 additions and 3 deletions.
81 changes: 79 additions & 2 deletions library/std/src/io/copy.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
use super::{BorrowedBuf, BufReader, BufWriter, ErrorKind, Read, Result, Write, DEFAULT_BUF_SIZE};
use crate::alloc::Allocator;
use crate::cmp;
use crate::collections::VecDeque;
use crate::io::IoSlice;
use crate::mem::MaybeUninit;

#[cfg(test)]
Expand Down Expand Up @@ -86,7 +90,7 @@ where

/// Specialization of the read-write loop that reuses the internal
/// buffer of a BufReader. If there's no buffer then the writer side
/// should be used intead.
/// should be used instead.
trait BufferedReaderSpec {
fn buffer_size(&self) -> usize;

Expand All @@ -104,7 +108,39 @@ where
}

default fn copy_to(&mut self, _to: &mut (impl Write + ?Sized)) -> Result<u64> {
unimplemented!("only called from specializations");
unreachable!("only called from specializations")
}
}

impl BufferedReaderSpec for &[u8] {
fn buffer_size(&self) -> usize {
// prefer this specialization since the source "buffer" is all we'll ever need,
// even if it's small
usize::MAX
}

fn copy_to(&mut self, to: &mut (impl Write + ?Sized)) -> Result<u64> {
let len = self.len();
to.write_all(self)?;
*self = &self[len..];
Ok(len as u64)
}
}

impl<A: Allocator> BufferedReaderSpec for VecDeque<u8, A> {
fn buffer_size(&self) -> usize {
// prefer this specialization since the source "buffer" is all we'll ever need,
// even if it's small
usize::MAX
}

fn copy_to(&mut self, to: &mut (impl Write + ?Sized)) -> Result<u64> {
let len = self.len();
let (front, back) = self.as_slices();
let bufs = &mut [IoSlice::new(front), IoSlice::new(back)];
to.write_all_vectored(bufs)?;
self.clear();
Ok(len as u64)
}
}

Expand Down Expand Up @@ -218,6 +254,47 @@ impl<I: Write + ?Sized> BufferedWriterSpec for BufWriter<I> {
}
}

impl<A: Allocator> BufferedWriterSpec for Vec<u8, A> {
fn buffer_size(&self) -> usize {
cmp::max(DEFAULT_BUF_SIZE, self.capacity() - self.len())
}

fn copy_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<u64> {
let mut bytes = 0;

// avoid allocating before we have determined that there's anything to read
if self.capacity() == 0 {
bytes = stack_buffer_copy(&mut reader.take(DEFAULT_BUF_SIZE as u64), self)?;
if bytes == 0 {
return Ok(0);
}
}

loop {
self.reserve(DEFAULT_BUF_SIZE);
let mut buf: BorrowedBuf<'_> = self.spare_capacity_mut().into();
match reader.read_buf(buf.unfilled()) {
Ok(()) => {}
Err(e) if e.kind() == ErrorKind::Interrupted => continue,
Err(e) => return Err(e),
};

let read = buf.filled().len();
if read == 0 {
break;
}

// SAFETY: BorrowedBuf guarantees all of its filled bytes are init
// and the number of read bytes can't exceed the spare capacity since
// that's what the buffer is borrowing from.
unsafe { self.set_len(self.len() + read) };
bytes += read as u64;
}

Ok(bytes)
}
}

fn stack_buffer_copy<R: Read + ?Sized, W: Write + ?Sized>(
reader: &mut R,
writer: &mut W,
Expand Down
38 changes: 37 additions & 1 deletion library/std/src/io/copy/tests.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
use crate::cmp::{max, min};
use crate::collections::VecDeque;
use crate::io;
use crate::io::*;

#[test]
Expand All @@ -19,7 +21,7 @@ struct ShortReader {

impl Read for ShortReader {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
let bytes = min(self.cap, self.read_size);
let bytes = min(self.cap, self.read_size).min(buf.len());
self.cap -= bytes;
self.observed_buffer = max(self.observed_buffer, buf.len());
Ok(bytes)
Expand Down Expand Up @@ -78,6 +80,40 @@ fn copy_specializes_bufreader() {
);
}

#[test]
fn copy_specializes_to_vec() {
let cap = 123456;
let mut source = ShortReader { cap, observed_buffer: 0, read_size: 1337 };
let mut sink = Vec::new();
assert_eq!(cap as u64, io::copy(&mut source, &mut sink).unwrap());
assert!(
source.observed_buffer > DEFAULT_BUF_SIZE,
"expected a large buffer to be provided to the reader"
);
}

#[test]
fn copy_specializes_from_vecdeque() {
let mut source = VecDeque::with_capacity(100 * 1024);
for _ in 0..20 * 1024 {
source.push_front(0);
}
for _ in 0..20 * 1024 {
source.push_back(0);
}
let mut sink = WriteObserver { observed_buffer: 0 };
assert_eq!(40 * 1024u64, io::copy(&mut source, &mut sink).unwrap());
assert_eq!(20 * 1024, sink.observed_buffer);
}

#[test]
fn copy_specializes_from_slice() {
let mut source = [1; 60 * 1024].as_slice();
let mut sink = WriteObserver { observed_buffer: 0 };
assert_eq!(60 * 1024u64, io::copy(&mut source, &mut sink).unwrap());
assert_eq!(60 * 1024, sink.observed_buffer);
}

#[cfg(unix)]
mod io_benches {
use crate::fs::File;
Expand Down

0 comments on commit 6f8ba51

Please sign in to comment.