Rollup merge of #92663 - cuviper:generic-write-cursor, r=dtolnay

Implement `Write for Cursor<[u8; N]>`, plus `A: Allocator` cursor support

This implements `Write for Cursor<[u8; N]>`, and also adds support for generic `A: Allocator` in `Box` and `Vec` cursors.

This was inspired by a user questioning why they couldn't write a `Cursor<[u8; N]>`:
https://users.rust-lang.org/t/why-vec-and-not-u8-makes-cursor-have-write/68210

Related history:
- #27197 switched `AsRef<[u8]>` for reading and seeking
- #67415 tried to use `AsMut<[u8]>` for writing, but did not specialize `Vec`.

2 files changed, 86 insertions, 51 deletions

diff --git a/library/std/src/io/cursor.rs b/library/std/src/io/cursor.rs
index fc19704bece..57f1d628f6a 100644
--- a/library/std/src/io/cursor.rs
+++ b/library/std/src/io/cursor.rs
@@ -3,6 +3,7 @@ mod tests;
 
 use crate::io::prelude::*;
 
+use crate::alloc::Allocator;
 use crate::cmp;
 use crate::io::{self, ErrorKind, IoSlice, IoSliceMut, ReadBuf, SeekFrom};
 
@@ -398,7 +399,10 @@ fn slice_write_vectored(
 }
 
 // Resizing write implementation
-fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usize> {
+fn vec_write<A>(pos_mut: &mut u64, vec: &mut Vec<u8, A>, buf: &[u8]) -> io::Result<usize>
+where
+    A: Allocator,
+{
     let pos: usize = (*pos_mut).try_into().map_err(|_| {
         io::const_io_error!(
             ErrorKind::InvalidInput,
@@ -426,11 +430,14 @@ fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usi
     Ok(buf.len())
 }
 
-fn vec_write_vectored(
+fn vec_write_vectored<A>(
     pos_mut: &mut u64,
-    vec: &mut Vec<u8>,
+    vec: &mut Vec<u8, A>,
     bufs: &[IoSlice<'_>],
-) -> io::Result<usize> {
+) -> io::Result<usize>
+where
+    A: Allocator,
+{
     let mut nwritten = 0;
     for buf in bufs {
         nwritten += vec_write(pos_mut, vec, buf)?;
@@ -462,7 +469,10 @@ impl Write for Cursor<&mut [u8]> {
 }
 
 #[stable(feature = "cursor_mut_vec", since = "1.25.0")]
-impl Write for Cursor<&mut Vec<u8>> {
+impl<A> Write for Cursor<&mut Vec<u8, A>>
+where
+    A: Allocator,
+{
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         vec_write(&mut self.pos, self.inner, buf)
     }
@@ -483,7 +493,10 @@ impl Write for Cursor<&mut Vec<u8>> {
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl Write for Cursor<Vec<u8>> {
+impl<A> Write for Cursor<Vec<u8, A>>
+where
+    A: Allocator,
+{
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         vec_write(&mut self.pos, &mut self.inner, buf)
     }
@@ -504,7 +517,33 @@ impl Write for Cursor<Vec<u8>> {
 }
 
 #[stable(feature = "cursor_box_slice", since = "1.5.0")]
-impl Write for Cursor<Box<[u8]>> {
+impl<A> Write for Cursor<Box<[u8], A>>
+where
+    A: Allocator,
+{
+    #[inline]
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        slice_write(&mut self.pos, &mut self.inner, buf)
+    }
+
+    #[inline]
+    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        slice_write_vectored(&mut self.pos, &mut self.inner, bufs)
+    }
+
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        true
+    }
+
+    #[inline]
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+#[stable(feature = "cursor_array", since = "1.61.0")]
+impl<const N: usize> Write for Cursor<[u8; N]> {
     #[inline]
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         slice_write(&mut self.pos, &mut self.inner, buf)
diff --git a/library/std/src/io/cursor/tests.rs b/library/std/src/io/cursor/tests.rs
index 5da31ce0ba7..f1ee177b7f3 100644
--- a/library/std/src/io/cursor/tests.rs
+++ b/library/std/src/io/cursor/tests.rs
@@ -50,9 +50,11 @@ fn test_mem_mut_writer() {
     assert_eq!(&writer.get_ref()[..], b);
 }
 
-#[test]
-fn test_box_slice_writer() {
-    let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
+fn test_slice_writer<T>(writer: &mut Cursor<T>)
+where
+    T: AsRef<[u8]>,
+    Cursor<T>: Write,
+{
     assert_eq!(writer.position(), 0);
     assert_eq!(writer.write(&[0]).unwrap(), 1);
     assert_eq!(writer.position(), 1);
@@ -65,12 +67,14 @@ fn test_box_slice_writer() {
     assert_eq!(writer.write(&[8, 9]).unwrap(), 1);
     assert_eq!(writer.write(&[10]).unwrap(), 0);
     let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
-    assert_eq!(&**writer.get_ref(), b);
+    assert_eq!(writer.get_ref().as_ref(), b);
 }
 
-#[test]
-fn test_box_slice_writer_vectored() {
-    let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
+fn test_slice_writer_vectored<T>(writer: &mut Cursor<T>)
+where
+    T: AsRef<[u8]>,
+    Cursor<T>: Write,
+{
     assert_eq!(writer.position(), 0);
     assert_eq!(writer.write_vectored(&[IoSlice::new(&[0])]).unwrap(), 1);
     assert_eq!(writer.position(), 1);
@@ -85,53 +89,45 @@ fn test_box_slice_writer_vectored() {
     assert_eq!(writer.write_vectored(&[IoSlice::new(&[8, 9])]).unwrap(), 1);
     assert_eq!(writer.write_vectored(&[IoSlice::new(&[10])]).unwrap(), 0);
     let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
-    assert_eq!(&**writer.get_ref(), b);
+    assert_eq!(writer.get_ref().as_ref(), b);
+}
+
+#[test]
+fn test_box_slice_writer() {
+    let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
+    test_slice_writer(&mut writer);
+}
+
+#[test]
+fn test_box_slice_writer_vectored() {
+    let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
+    test_slice_writer_vectored(&mut writer);
+}
+
+#[test]
+fn test_array_writer() {
+    let mut writer = Cursor::new([0u8; 9]);
+    test_slice_writer(&mut writer);
+}
+
+#[test]
+fn test_array_writer_vectored() {
+    let mut writer = Cursor::new([0u8; 9]);
+    test_slice_writer_vectored(&mut writer);
 }
 
 #[test]
 fn test_buf_writer() {
     let mut buf = [0 as u8; 9];
-    {
-        let mut writer = Cursor::new(&mut buf[..]);
-        assert_eq!(writer.position(), 0);
-        assert_eq!(writer.write(&[0]).unwrap(), 1);
-        assert_eq!(writer.position(), 1);
-        assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
-        assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
-        assert_eq!(writer.position(), 8);
-        assert_eq!(writer.write(&[]).unwrap(), 0);
-        assert_eq!(writer.position(), 8);
-
-        assert_eq!(writer.write(&[8, 9]).unwrap(), 1);
-        assert_eq!(writer.write(&[10]).unwrap(), 0);
-    }
-    let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
-    assert_eq!(buf, b);
+    let mut writer = Cursor::new(&mut buf[..]);
+    test_slice_writer(&mut writer);
 }
 
 #[test]
 fn test_buf_writer_vectored() {
     let mut buf = [0 as u8; 9];
-    {
-        let mut writer = Cursor::new(&mut buf[..]);
-        assert_eq!(writer.position(), 0);
-        assert_eq!(writer.write_vectored(&[IoSlice::new(&[0])]).unwrap(), 1);
-        assert_eq!(writer.position(), 1);
-        assert_eq!(
-            writer
-                .write_vectored(&[IoSlice::new(&[1, 2, 3]), IoSlice::new(&[4, 5, 6, 7])],)
-                .unwrap(),
-            7,
-        );
-        assert_eq!(writer.position(), 8);
-        assert_eq!(writer.write_vectored(&[]).unwrap(), 0);
-        assert_eq!(writer.position(), 8);
-
-        assert_eq!(writer.write_vectored(&[IoSlice::new(&[8, 9])]).unwrap(), 1);
-        assert_eq!(writer.write_vectored(&[IoSlice::new(&[10])]).unwrap(), 0);
-    }
-    let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
-    assert_eq!(buf, b);
+    let mut writer = Cursor::new(&mut buf[..]);
+    test_slice_writer_vectored(&mut writer);
 }
 
 #[test]