additional io::copy specializations

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,

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<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)
@@ -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;