Rollup merge of #158548 - bushrat011899:alloc_io_copy_internals, r=clarfonthey

Move `std::io::copy` to `alloc::io`

ACP: rust-lang/libs-team#755
Tracking issue: #154046
Split From: #156527
~~Blocked On: #158547

## Description

Moves `std::io::copy` into `alloc::io`. Blocked on #158547.

This relies on specialization to allow `std` to provide optimised copy implementations for its types where appropriate. The exact technique involves defining a new trait, `alloc::io::SpecCopy`:

```rust
#[doc(hidden)]
#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
#[rustc_specialization_trait]
pub trait SpecCopy: Read {
    /// Attempt to copy from this reader to the provided writer using a specialized
    /// process.
    fn copy<R: Read + ?Sized, W: Write + ?Sized>(
        _reader: &mut R,
        _writer: &mut W,
    ) -> Result<CopyState>;
}
```

Since optimised copying requires both the reader and writer to support the operation between each other, we can choose one of them to be the implementer of the copy algorithm, and delegate specialization to it. In this case, I've chosen the reader to be the provider of the specialized copy implementation arbitrarily. Note that the `SpecCopy::copy` function is generic over the reader specifically to allow wrappers like `Take<R>` to be visible to the implementation of `copy`.

Because this introduces a new layer of specialization to `io::copy`, I think this PR should be benchmarked to make sure performance characteristics aren't too different. I am expecting compilation time to be slightly worse, since there's just more specialization happening, but the actual code run _should_ be the same.

---

## Notes

* No AI tooling of any kind was used during the creation of this PR.
* Please see #154046 (comment) for a review order and broader context for this PR.

Author: JonathanBrouwer

library/alloc/src/io/copy.rs

@@ -0,0 +1,90 @@
+mod generic;
+mod specialization;
+
+use self::generic::generic_copy;
+#[doc(hidden)]
+#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
+pub use self::specialization::SpecCopy;
+use self::specialization::SpecCopyInner;
+use crate::io::{Read, Result, Write};
+
+/// Used as a part of [copy specialization](SpecCopy) to communicate how many bytes
+/// were copied, and whether copying is done.
+///
+/// * [`Ended(n)`](CopyState::Ended) indicates copying completed, moving a total
+///   of `n` bytes.
+/// * [`Fallback(n)`](CopyState::Fallback) indicates copying is _might_ not be
+///   complete, and so far `n` bytes have been copied using specialization.
+///   The remaining must be copied using a fallback implementation.
+///
+/// If a particular `Read` and `Write` combination do not implement a specialized
+/// copy routine, the specialized function will return `Fallback(0)`.
+#[derive(Debug)]
+#[doc(hidden)]
+#[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
+pub enum CopyState {
+    Ended(u64),
+    Fallback(u64),
+}
+
+/// Copies the entire contents of a reader into a writer.
+///
+/// This function will continuously read data from `reader` and then
+/// write it into `writer` in a streaming fashion until `reader`
+/// returns EOF.
+///
+/// On success, the total number of bytes that were copied from
+/// `reader` to `writer` is returned.
+///
+/// If you want to copy the contents of one file to another and you’re
+/// working with filesystem paths, see the [`fs::copy`] function.
+///
+// FIXME(#74481): Hard-links required to link from `alloc` to `std`
+/// [`fs::copy`]: ../../std/fs/fn.copy.html
+///
+/// # Errors
+///
+/// This function will return an error immediately if any call to [`read`] or
+/// [`write`] returns an error. All instances of [`ErrorKind::Interrupted`] are
+/// handled by this function and the underlying operation is retried.
+///
+/// [`read`]: Read::read
+/// [`write`]: Write::write
+/// [`ErrorKind::Interrupted`]: crate::io::ErrorKind::Interrupted
+///
+/// # Examples
+///
+/// ```
+/// use std::io;
+///
+/// fn main() -> io::Result<()> {
+///     let mut reader: &[u8] = b"hello";
+///     let mut writer: Vec<u8> = vec![];
+///
+///     io::copy(&mut reader, &mut writer)?;
+///
+///     assert_eq!(&b"hello"[..], &writer[..]);
+///     Ok(())
+/// }
+/// ```
+///
+/// # Platform-specific behavior
+///
+/// On Linux (including Android), this function uses `copy_file_range(2)`,
+/// `sendfile(2)` or `splice(2)` syscalls to move data directly between file
+/// descriptors if possible.
+///
+/// Note that platform-specific behavior may change in the future.
+#[stable(feature = "rust1", since = "1.0.0")]
+pub fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> Result<u64>
+where
+    R: Read,
+    W: Write,
+{
+    match SpecCopyInner::copy(reader, writer)? {
+        CopyState::Ended(copied) => Ok(copied),
+        CopyState::Fallback(copied) => {
+            generic_copy(reader, writer).map(|additional| copied + additional)
+        }
+    }
+}

library/alloc/src/io/copy/generic.rs

@@ -0,0 +1,256 @@
+use core::mem::MaybeUninit;
+
+#[cfg(not(no_global_oom_handling))]
+use crate::collections::VecDeque;
+use crate::io::{BorrowedBuf, BufReader, BufWriter, DEFAULT_BUF_SIZE, Read, Result, Write};
+use crate::vec::Vec;
+#[cfg_attr(
+    no_global_oom_handling,
+    expect(unused_imports, reason = "only required for VecDeque specialization")
+)]
+use crate::{alloc::Allocator, io::IoSlice};
+
+/// The userspace read-write-loop implementation of `io::copy` that is used when
+/// OS-specific specializations for copy offloading are not available or not applicable.
+///
+/// This function is able to perform a mild amount of specialization based on
+/// the size of the `reader` and `writer` buffers, if they have any.
+///
+/// * If `reader`'s buffer is large enough ([`>=DEFAULT_BUF_SIZE`](DEFAULT_BUF_SIZE)),
+///   _and_ it is larger than `writer`'s buffer, copying will be controlled by `R`.
+/// * Otherwise, copying will be controlled by `writer`.
+///
+/// Currently, `[u8]`, `Vec<u8>`, `VecDeque<u8>`, `BufReader<T>` and `BufWriter<T>`
+/// are specialized with.
+pub(super) fn generic_copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> Result<u64>
+where
+    R: Read,
+    W: Write,
+{
+    let read_priority = BufferedReaderSpec::buffer_priority(reader);
+    let write_priority = BufferedWriterSpec::buffer_priority(writer);
+
+    if read_priority >= DEFAULT_BUF_SIZE && read_priority >= write_priority {
+        return BufferedReaderSpec::copy_to(reader, writer);
+    }
+
+    BufferedWriterSpec::copy_from(writer, reader)
+}
+
+/// This is used by [`generic_copy`] to decide whether to use [`BufferedReaderSpec::copy_to`]
+/// or [`BufferedWriterSpec::copy_from`].
+type BufferPriority = usize;
+
+/// Unbuffered readers and writers have the lowest priority.
+const UNBUFFERED: BufferPriority = BufferPriority::MIN;
+
+/// Readers and writers with their entire contents have the highest priority.
+const IN_MEMORY: BufferPriority = BufferPriority::MAX;
+
+/// Specialization of the read-write loop in [`generic_copy`] that reuses the
+/// internal buffer of a [`BufReader`]. If there's no buffer then the writer side
+/// should be used instead.
+trait BufferedReaderSpec {
+    fn buffer_priority(&self) -> BufferPriority;
+
+    fn copy_to(&mut self, to: &mut (impl Write + ?Sized)) -> Result<u64>;
+}
+
+impl<T> BufferedReaderSpec for T
+where
+    Self: Read,
+    T: ?Sized,
+{
+    #[inline]
+    default fn buffer_priority(&self) -> BufferPriority {
+        UNBUFFERED
+    }
+
+    default fn copy_to(&mut self, _to: &mut (impl Write + ?Sized)) -> Result<u64> {
+        unreachable!("only called from specializations")
+    }
+}
+
+impl BufferedReaderSpec for &[u8] {
+    fn buffer_priority(&self) -> BufferPriority {
+        IN_MEMORY
+    }
+
+    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)
+    }
+}
+
+#[cfg(not(no_global_oom_handling))]
+impl<A: Allocator> BufferedReaderSpec for VecDeque<u8, A> {
+    fn buffer_priority(&self) -> BufferPriority {
+        IN_MEMORY
+    }
+
+    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)
+    }
+}
+
+impl<I> BufferedReaderSpec for BufReader<I>
+where
+    Self: Read,
+    I: ?Sized,
+{
+    fn buffer_priority(&self) -> BufferPriority {
+        self.capacity()
+    }
+
+    fn copy_to(&mut self, to: &mut (impl Write + ?Sized)) -> Result<u64> {
+        let mut len = 0;
+
+        loop {
+            // Hack: this relies on `impl Read for BufReader` always calling fill_buf
+            // if the buffer is empty, even for empty slices.
+            // It can't be called directly here since specialization prevents us
+            // from adding I: Read
+            match self.read(&mut []) {
+                Ok(_) => {}
+                Err(e) if e.is_interrupted() => continue,
+                Err(e) => return Err(e),
+            }
+            let buf = self.buffer();
+            if self.buffer().len() == 0 {
+                return Ok(len);
+            }
+
+            // In case the writer side is a BufWriter then its write_all
+            // implements an optimization that passes through large
+            // buffers to the underlying writer. That code path is #[cold]
+            // but we're still avoiding redundant memcopies when doing
+            // a copy between buffered inputs and outputs.
+            to.write_all(buf)?;
+            len += buf.len() as u64;
+            self.discard_buffer();
+        }
+    }
+}
+
+/// Specialization of the read-write loop in `generic_copy` that either uses a
+/// stack buffer or reuses the internal buffer of a [`BufWriter`].
+trait BufferedWriterSpec: Write {
+    fn buffer_priority(&self) -> BufferPriority;
+
+    fn copy_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<u64>;
+}
+
+impl<W: Write + ?Sized> BufferedWriterSpec for W {
+    #[inline]
+    default fn buffer_priority(&self) -> BufferPriority {
+        UNBUFFERED
+    }
+
+    default fn copy_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<u64> {
+        // Unlike `BufferedReaderSpec::copy_to`, this _will_ be called as the fallback
+        // when both the reader and writer provide no specialization.
+        stack_buffer_copy(reader, self)
+    }
+}
+
+impl<I: Write + ?Sized> BufferedWriterSpec for BufWriter<I> {
+    fn buffer_priority(&self) -> BufferPriority {
+        self.capacity()
+    }
+
+    fn copy_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<u64> {
+        if self.capacity() < DEFAULT_BUF_SIZE {
+            // Since neither this buffer nor the reader's buffer are large enough,
+            // fall back to the unspecialized implementation.
+            return stack_buffer_copy(reader, self);
+        }
+
+        let mut len = 0;
+        let mut init = false;
+
+        loop {
+            let buf = self.buffer_mut();
+            let mut read_buf: BorrowedBuf<'_, u8> = buf.spare_capacity_mut().into();
+
+            if init {
+                // SAFETY: `init` is only true after `reader` initializes
+                // `read_buf`. See the comment about `flush_buf` below.
+                unsafe { read_buf.set_init() };
+            }
+
+            if read_buf.capacity() >= DEFAULT_BUF_SIZE {
+                let mut cursor = read_buf.unfilled();
+                match reader.read_buf(cursor.reborrow()) {
+                    Ok(()) => {
+                        let bytes_read = cursor.written();
+
+                        if bytes_read == 0 {
+                            return Ok(len);
+                        }
+
+                        init = read_buf.is_init();
+                        len += bytes_read as u64;
+
+                        // SAFETY: BorrowedBuf guarantees all of its filled bytes are init
+                        unsafe { buf.set_len(buf.len() + bytes_read) };
+
+                        // Read again if the buffer still has enough capacity, as BufWriter itself would do
+                        // This will occur if the reader returns short reads
+                    }
+                    Err(ref e) if e.is_interrupted() => {}
+                    Err(e) => return Err(e),
+                }
+            } else {
+                // SAFETY: `flush_buf` will not de-initialize any elements of
+                // the spare capacity so we can remember `init` across this.
+                self.flush_buf()?;
+            }
+        }
+    }
+}
+
+impl BufferedWriterSpec for Vec<u8> {
+    fn buffer_priority(&self) -> BufferPriority {
+        self.capacity() - self.len()
+    }
+
+    fn copy_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<u64> {
+        reader.read_to_end(self).map(|bytes| u64::try_from(bytes).expect("usize overflowed u64"))
+    }
+}
+
+/// Copies from `reader` to `writer` using a stack-allocated buffer (a fixed sized array).
+fn stack_buffer_copy<R: Read + ?Sized, W: Write + ?Sized>(
+    reader: &mut R,
+    writer: &mut W,
+) -> Result<u64> {
+    let buf: &mut [_] = &mut [MaybeUninit::uninit(); DEFAULT_BUF_SIZE];
+    let mut buf: BorrowedBuf<'_, u8> = buf.into();
+
+    let mut len = 0;
+
+    loop {
+        match reader.read_buf(buf.unfilled()) {
+            Ok(()) => {}
+            Err(e) if e.is_interrupted() => continue,
+            Err(e) => return Err(e),
+        };
+
+        if buf.filled().is_empty() {
+            break;
+        }
+
+        len += buf.filled().len() as u64;
+        writer.write_all(buf.filled())?;
+        buf.clear();
+    }
+
+    Ok(len)
+}