diff --git a/library/std/src/io/impls.rs b/library/std/src/io/impls.rs index 64d2457bce159..42f39ee49a576 100644 --- a/library/std/src/io/impls.rs +++ b/library/std/src/io/impls.rs @@ -7,7 +7,7 @@ use crate::fmt; use crate::io::{ self, BufRead, ErrorKind, IoSlice, IoSliceMut, Read, ReadBuf, Seek, SeekFrom, Write, }; -use crate::mem; +use crate::{mem, ptr}; // ============================================================================= // Forwarding implementations @@ -410,3 +410,59 @@ impl Write for Vec { Ok(()) } } + +/// Write is implemented for `&mut [MaybeUninit]` by copying into the slice, overwriting +/// its data. After a write, the sucessfully written elements are guaranteed to be initialized. +/// +/// 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 = "write_maybeuninit", since = "1.60.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::replace(self, &mut []).split_at_mut(amt); + // SAFETY: since self is a mutable slice, the aliasing rules prevent overlap + unsafe { + ptr::copy_nonoverlapping(data.as_ptr(), a.as_mut_ptr() as *mut u8, 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() { + Ok(()) + } else { + Err(io::const_io_error!(ErrorKind::WriteZero, "failed to write whole buffer")) + } + } + + #[inline] + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} diff --git a/library/std/src/io/impls/tests.rs b/library/std/src/io/impls/tests.rs index d1cd84a67ada5..fabec898040c0 100644 --- a/library/std/src/io/impls/tests.rs +++ b/library/std/src/io/impls/tests.rs @@ -1,4 +1,5 @@ use crate::io::prelude::*; +use crate::mem::MaybeUninit; #[bench] fn bench_read_slice(b: &mut test::Bencher) { @@ -44,7 +45,21 @@ fn bench_read_vec(b: &mut test::Bencher) { #[bench] fn bench_write_vec(b: &mut test::Bencher) { - let mut buf = Vec::with_capacity(1024); + let mut buf = Vec::::with_capacity(1024); + let src = [5; 128]; + + b.iter(|| { + let mut wr = &mut buf[..]; + for _ in 0..8 { + let _ = wr.write_all(&src); + test::black_box(&wr); + } + }) +} + +#[bench] +fn bench_write_maybeuninit(b: &mut test::Bencher) { + let mut buf = [MaybeUninit::::uninit(); 1024]; let src = [5; 128]; b.iter(|| {