diff --git a/library/std/src/io/copy.rs b/library/std/src/io/copy.rs index ef1f4031ef202..3322940d2452f 100644 --- a/library/std/src/io/copy.rs +++ b/library/std/src/io/copy.rs @@ -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)] @@ -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; @@ -104,7 +108,39 @@ where } default fn copy_to(&mut self, _to: &mut (impl Write + ?Sized)) -> Result { - 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 { + let len = self.len(); + to.write_all(self)?; + *self = &self[len..]; + Ok(len as u64) + } +} + +impl BufferedReaderSpec for VecDeque { + 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 { + 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) } } @@ -218,6 +254,47 @@ impl BufferedWriterSpec for BufWriter { } } +impl BufferedWriterSpec for Vec { + fn buffer_size(&self) -> usize { + cmp::max(DEFAULT_BUF_SIZE, self.capacity() - self.len()) + } + + fn copy_from(&mut self, reader: &mut R) -> Result { + 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( reader: &mut R, writer: &mut W, diff --git a/library/std/src/io/copy/tests.rs b/library/std/src/io/copy/tests.rs index 8c816af15b808..af137eaf856f5 100644 --- a/library/std/src/io/copy/tests.rs +++ b/library/std/src/io/copy/tests.rs @@ -1,4 +1,6 @@ use crate::cmp::{max, min}; +use crate::collections::VecDeque; +use crate::io; use crate::io::*; #[test] @@ -19,7 +21,7 @@ struct ShortReader { impl Read for ShortReader { fn read(&mut self, buf: &mut [u8]) -> Result { - 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) @@ -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;