Rollup merge of #137107 - thaliaarchi:io-optional-methods/cursors, r=joboet

Override default `Write` methods for cursor-like types

Override the default `io::Write` methods for cursor-like types to provide more efficient versions.

Writes to resizable containers already write everything, so implement `write_all` and `write_all_vectored` in terms of those. For fixed-sized containers, cut out unnecessary error checking and looping for those same methods.

| `impl Write for T`              | `vectored` | `all` | `all_vectored` | `fmt`   |
| ------------------------------- | ---------- | ----- | -------------- | ------- |
| `&mut [u8]`                     | Y          | Y     | new            |         |
| `Vec<u8>`                       | Y          | Y     | new            | #137762 |
| `VecDeque<u8>`                  | Y          | Y     | new            | #137762 |
| `std::io::Cursor<&mut [u8]>`    | Y          | new   | new            |         |
| `std::io::Cursor<&mut Vec<u8>>` | Y          | new   | new            | #137762 |
| `std::io::Cursor<Vec<u8>>`      | Y          | new   | new            | #137762 |
| `std::io::Cursor<Box<[u8]>>`    | Y          | new   | new            |         |
| `std::io::Cursor<[u8; N]>`      | Y          | new   | new            |         |
| `core::io::BorrowedCursor<'_>`  | new        | new   | new            |         |

Tracked in https://github.com/rust-lang/rust/issues/136756.

# Open questions

Is it guaranteed by `Write::write_all` that the maximal write is performed when not everything can be written? Its documentation describes the behavior of the default implementation, which writes until a 0-length write is encountered, thus implying that a maximal write is expected. In contrast, `Read::read_exact` declares that the contents of the buffer are unspecified for short reads. If it were allowed, these cursor-like types could bail on the write altogether if it has insufficient capacity.

2 files changed, 146 insertions, 14 deletions

diff --git a/library/std/src/io/cursor.rs b/library/std/src/io/cursor.rs
index 08832bbc1e3..d7131e2fe92 100644
--- a/library/std/src/io/cursor.rs
+++ b/library/std/src/io/cursor.rs
@@ -439,6 +439,27 @@ fn slice_write_vectored(
     Ok(nwritten)
 }
 
+#[inline]
+fn slice_write_all(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<()> {
+    let n = slice_write(pos_mut, slice, buf)?;
+    if n < buf.len() { Err(io::Error::WRITE_ALL_EOF) } else { Ok(()) }
+}
+
+#[inline]
+fn slice_write_all_vectored(
+    pos_mut: &mut u64,
+    slice: &mut [u8],
+    bufs: &[IoSlice<'_>],
+) -> io::Result<()> {
+    for buf in bufs {
+        let n = slice_write(pos_mut, slice, buf)?;
+        if n < buf.len() {
+            return Err(io::Error::WRITE_ALL_EOF);
+        }
+    }
+    Ok(())
+}
+
 /// Reserves the required space, and pads the vec with 0s if necessary.
 fn reserve_and_pad<A: Allocator>(
     pos_mut: &mut u64,
@@ -481,9 +502,12 @@ fn reserve_and_pad<A: Allocator>(
     Ok(pos)
 }
 
-/// Writes the slice to the vec without allocating
-/// # Safety: vec must have buf.len() spare capacity
-unsafe fn vec_write_unchecked<A>(pos: usize, vec: &mut Vec<u8, A>, buf: &[u8]) -> usize
+/// Writes the slice to the vec without allocating.
+///
+/// # Safety
+///
+/// `vec` must have `buf.len()` spare capacity.
+unsafe fn vec_write_all_unchecked<A>(pos: usize, vec: &mut Vec<u8, A>, buf: &[u8]) -> usize
 where
     A: Allocator,
 {
@@ -492,7 +516,7 @@ where
     pos + buf.len()
 }
 
-/// Resizing write implementation for [`Cursor`]
+/// Resizing `write_all` implementation for [`Cursor`].
 ///
 /// Cursor is allowed to have a pre-allocated and initialised
 /// vector body, but with a position of 0. This means the [`Write`]
@@ -501,7 +525,7 @@ where
 /// This also allows for the vec body to be empty, but with a position of N.
 /// This means that [`Write`] will pad the vec with 0 initially,
 /// before writing anything from that point
-fn vec_write<A>(pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf: &[u8]) -> io::Result<usize>
+fn vec_write_all<A>(pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf: &[u8]) -> io::Result<usize>
 where
     A: Allocator,
 {
@@ -512,7 +536,7 @@ where
     // Safety: we have ensured that the capacity is available
     // and that all bytes get written up to pos
     unsafe {
-        pos = vec_write_unchecked(pos, vec, buf);
+        pos = vec_write_all_unchecked(pos, vec, buf);
         if pos > vec.len() {
             vec.set_len(pos);
         }
@@ -523,7 +547,7 @@ where
     Ok(buf_len)
 }
 
-/// Resizing write_vectored implementation for [`Cursor`]
+/// Resizing `write_all_vectored` implementation for [`Cursor`].
 ///
 /// Cursor is allowed to have a pre-allocated and initialised
 /// vector body, but with a position of 0. This means the [`Write`]
@@ -532,7 +556,7 @@ where
 /// This also allows for the vec body to be empty, but with a position of N.
 /// This means that [`Write`] will pad the vec with 0 initially,
 /// before writing anything from that point
-fn vec_write_vectored<A>(
+fn vec_write_all_vectored<A>(
     pos_mut: &mut u64,
     vec: &mut Vec<u8, A>,
     bufs: &[IoSlice<'_>],
@@ -550,7 +574,7 @@ where
     // and that all bytes get written up to the last pos
     unsafe {
         for buf in bufs {
-            pos = vec_write_unchecked(pos, vec, buf);
+            pos = vec_write_all_unchecked(pos, vec, buf);
         }
         if pos > vec.len() {
             vec.set_len(pos);
@@ -580,6 +604,16 @@ impl Write for Cursor<&mut [u8]> {
     }
 
     #[inline]
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        slice_write_all(&mut self.pos, self.inner, buf)
+    }
+
+    #[inline]
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        slice_write_all_vectored(&mut self.pos, self.inner, bufs)
+    }
+
+    #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }
@@ -591,11 +625,11 @@ where
     A: Allocator,
 {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        vec_write(&mut self.pos, self.inner, buf)
+        vec_write_all(&mut self.pos, self.inner, buf)
     }
 
     fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
-        vec_write_vectored(&mut self.pos, self.inner, bufs)
+        vec_write_all_vectored(&mut self.pos, self.inner, bufs)
     }
 
     #[inline]
@@ -603,6 +637,16 @@ where
         true
     }
 
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        vec_write_all(&mut self.pos, self.inner, buf)?;
+        Ok(())
+    }
+
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        vec_write_all_vectored(&mut self.pos, self.inner, bufs)?;
+        Ok(())
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
@@ -615,11 +659,11 @@ where
     A: Allocator,
 {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        vec_write(&mut self.pos, &mut self.inner, buf)
+        vec_write_all(&mut self.pos, &mut self.inner, buf)
     }
 
     fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
-        vec_write_vectored(&mut self.pos, &mut self.inner, bufs)
+        vec_write_all_vectored(&mut self.pos, &mut self.inner, bufs)
     }
 
     #[inline]
@@ -627,6 +671,16 @@ where
         true
     }
 
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        vec_write_all(&mut self.pos, &mut self.inner, buf)?;
+        Ok(())
+    }
+
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        vec_write_all_vectored(&mut self.pos, &mut self.inner, bufs)?;
+        Ok(())
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
@@ -654,6 +708,16 @@ where
     }
 
     #[inline]
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        slice_write_all(&mut self.pos, &mut self.inner, buf)
+    }
+
+    #[inline]
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        slice_write_all_vectored(&mut self.pos, &mut self.inner, bufs)
+    }
+
+    #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }
@@ -677,6 +741,16 @@ impl<const N: usize> Write for Cursor<[u8; N]> {
     }
 
     #[inline]
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        slice_write_all(&mut self.pos, &mut self.inner, buf)
+    }
+
+    #[inline]
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        slice_write_all_vectored(&mut self.pos, &mut self.inner, bufs)
+    }
+
+    #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }
diff --git a/library/std/src/io/impls.rs b/library/std/src/io/impls.rs
index 8239b29884e..d0245f3d498 100644
--- a/library/std/src/io/impls.rs
+++ b/library/std/src/io/impls.rs
@@ -455,7 +455,17 @@ impl Write for &mut [u8] {
 
     #[inline]
     fn write_all(&mut self, data: &[u8]) -> io::Result<()> {
-        if self.write(data)? == data.len() { Ok(()) } else { Err(io::Error::WRITE_ALL_EOF) }
+        if self.write(data)? < data.len() { Err(io::Error::WRITE_ALL_EOF) } else { Ok(()) }
+    }
+
+    #[inline]
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        for buf in bufs {
+            if self.write(buf)? < buf.len() {
+                return Err(io::Error::WRITE_ALL_EOF);
+            }
+        }
+        Ok(())
     }
 
     #[inline]
@@ -496,6 +506,12 @@ impl<A: Allocator> Write for Vec<u8, A> {
     }
 
     #[inline]
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        self.write_vectored(bufs)?;
+        Ok(())
+    }
+
+    #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }
@@ -515,6 +531,7 @@ impl<A: Allocator> Read for VecDeque<u8, A> {
         Ok(n)
     }
 
+    #[inline]
     fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
         let (front, back) = self.as_slices();
 
@@ -547,6 +564,7 @@ impl<A: Allocator> Read for VecDeque<u8, A> {
         Ok(())
     }
 
+    #[inline]
     fn read_buf_exact(&mut self, mut cursor: BorrowedCursor<'_>) -> io::Result<()> {
         let len = cursor.capacity();
         let (front, back) = self.as_slices();
@@ -639,6 +657,12 @@ impl<A: Allocator> Write for VecDeque<u8, A> {
     }
 
     #[inline]
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        self.write_vectored(bufs)?;
+        Ok(())
+    }
+
+    #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }
@@ -646,6 +670,7 @@ impl<A: Allocator> Write for VecDeque<u8, A> {
 
 #[unstable(feature = "read_buf", issue = "78485")]
 impl<'a> io::Write for core::io::BorrowedCursor<'a> {
+    #[inline]
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         let amt = cmp::min(buf.len(), self.capacity());
         self.append(&buf[..amt]);
@@ -653,6 +678,39 @@ impl<'a> io::Write for core::io::BorrowedCursor<'a> {
     }
 
     #[inline]
+    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        let mut nwritten = 0;
+        for buf in bufs {
+            let n = self.write(buf)?;
+            nwritten += n;
+            if n < buf.len() {
+                break;
+            }
+        }
+        Ok(nwritten)
+    }
+
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        true
+    }
+
+    #[inline]
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        if self.write(buf)? < buf.len() { Err(io::Error::WRITE_ALL_EOF) } else { Ok(()) }
+    }
+
+    #[inline]
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        for buf in bufs {
+            if self.write(buf)? < buf.len() {
+                return Err(io::Error::WRITE_ALL_EOF);
+            }
+        }
+        Ok(())
+    }
+
+    #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }