diff options
Diffstat (limited to 'library/std/src/io/copy.rs')
-rw-r--r-- | library/std/src/io/copy.rs | 81 |
1 files changed, 79 insertions, 2 deletions
diff --git a/library/std/src/io/copy.rs b/library/std/src/io/copy.rs index ef1f4031e..3322940d2 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<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) } } @@ -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, |