diff --git a/library/std/src/io/impls.rs b/library/std/src/io/impls.rs index d0245f3d4984c..4fe29a6e6cc5d 100644 --- a/library/std/src/io/impls.rs +++ b/library/std/src/io/impls.rs @@ -474,6 +474,65 @@ impl Write for &mut [u8] { } } +/// Write is implemented for `&mut [MaybeUninit]` by copying into the slice, +/// initialising its data. +/// +/// Note that writing updates the slice to point to the yet unwritten part. +/// The slice will be empty when it has been completely overwritten. +/// +/// If the number of bytes to be written exceeds the size of the slice, write +/// operations will return short writes: ultimately, `Ok(0)`; in this situation, +/// `write_all` returns an error of kind `ErrorKind::WriteZero`. +#[stable(feature = "maybe_uninit_slice_io_write", since = "1.87.0")] +impl Write for &mut [mem::MaybeUninit] { + #[inline] + fn write(&mut self, data: &[u8]) -> io::Result { + let amt = cmp::min(data.len(), self.len()); + let (a, b) = mem::take(self).split_at_mut(amt); + a.write_copy_of_slice(&data[..amt]); + *self = b; + Ok(amt) + } + + #[inline] + fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result { + let mut nwritten = 0; + for buf in bufs { + nwritten += self.write(buf)?; + if self.is_empty() { + break; + } + } + + Ok(nwritten) + } + + #[inline] + fn is_write_vectored(&self) -> bool { + true + } + + #[inline] + fn write_all(&mut self, data: &[u8]) -> io::Result<()> { + 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] + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + /// Write is implemented for `Vec` by appending to the vector. /// The vector will grow as needed. #[stable(feature = "rust1", since = "1.0.0")] diff --git a/library/std/src/io/tests.rs b/library/std/src/io/tests.rs index b22988d4a8a9d..99aa23fc8b2ec 100644 --- a/library/std/src/io/tests.rs +++ b/library/std/src/io/tests.rs @@ -945,3 +945,20 @@ fn try_oom_error() { let io_err = io::Error::from(reserve_err); assert_eq!(io::ErrorKind::OutOfMemory, io_err.kind()); } + +#[test] +fn write_maybe_uninit_slice() { + let mut buf = [MaybeUninit::uninit(); 16]; + + let mut cursor = &mut buf[..]; + assert_eq!(3, cursor.write(b"foo").unwrap()); + assert_eq!(6, cursor.write_vectored(&[IoSlice::new(b"bar"), IoSlice::new(b"baz"),]).unwrap()); + assert_eq!(7, cursor.len()); + cursor.write_all(b"qux").unwrap(); + cursor.write_all(b"toolong").unwrap_err(); + assert_eq!(0, cursor.len()); + + // SAFETY: We’ve initialised it using Write interface. + let buf = unsafe { buf.assume_init_ref() }; + assert_eq!(b"foobarbazquxtool", buf); +}