From 0a0074980fb7802ac621ceac92bf98bdc640bd78 Mon Sep 17 00:00:00 2001 From: Andrew Wock Date: Sun, 18 Feb 2024 18:06:16 -0500 Subject: [PATCH] Implement MaybeUninit::fill{,_with,_from} ACP: rust-lang/libs-team#156 Signed-off-by: Andrew Wock --- library/core/src/mem/maybe_uninit.rs | 202 +++++++++++++++++++++++-- library/core/tests/lib.rs | 1 + library/core/tests/mem.rs | 217 ++++++++++++++++++++++++++- 3 files changed, 398 insertions(+), 22 deletions(-) diff --git a/library/core/src/mem/maybe_uninit.rs b/library/core/src/mem/maybe_uninit.rs index c19b5791562ce..026e21586d403 100644 --- a/library/core/src/mem/maybe_uninit.rs +++ b/library/core/src/mem/maybe_uninit.rs @@ -1125,22 +1125,6 @@ impl MaybeUninit { // unlike copy_from_slice this does not call clone_from_slice on the slice // this is because `MaybeUninit` does not implement Clone. - struct Guard<'a, T> { - slice: &'a mut [MaybeUninit], - initialized: usize, - } - - impl<'a, T> Drop for Guard<'a, T> { - fn drop(&mut self) { - let initialized_part = &mut self.slice[..self.initialized]; - // SAFETY: this raw slice will contain only initialized objects - // that's why, it is allowed to drop it. - unsafe { - crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(initialized_part)); - } - } - } - assert_eq!(this.len(), src.len(), "destination and source slices have different lengths"); // NOTE: We need to explicitly slice them to the same length // for bounds checking to be elided, and the optimizer will @@ -1162,6 +1146,151 @@ impl MaybeUninit { unsafe { MaybeUninit::slice_assume_init_mut(this) } } + /// Fills `this` with elements by cloning `value`, returning a mutable reference to the now + /// initialized contents of `this`. + /// Any previously initialized elements will not be dropped. + /// + /// This is similar to [`slice::fill`]. + /// + /// # Panics + /// + /// This function will panic if any call to `Clone` panics. + /// + /// If such a panic occurs, any elements previously initialized during this operation will be + /// dropped. + /// + /// # Examples + /// + /// Fill an uninit vec with 1. + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut buf = vec![MaybeUninit::uninit(); 10]; + /// let initialized = MaybeUninit::fill(buf.as_mut_slice(), 1); + /// assert_eq!(initialized, &mut [1; 10]); + /// ``` + #[doc(alias = "memset")] + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill<'a>(this: &'a mut [MaybeUninit], value: T) -> &'a mut [T] + where + T: Clone, + { + SpecFill::spec_fill(this, value); + // SAFETY: Valid elements have just been filled into `this` so it is initialized + unsafe { MaybeUninit::slice_assume_init_mut(this) } + } + + /// Fills `this` with elements returned by calling a closure repeatedly. + /// + /// This method uses a closure to create new values. If you'd rather `Clone` a given value, use + /// [`MaybeUninit::fill`]. If you want to use the `Default` trait to generate values, you can + /// pass [`Default::default`] as the argument. + /// + /// # Panics + /// + /// This function will panic if any call to the provided closure panics. + /// + /// If such a panic occurs, any elements previously initialized during this operation will be + /// dropped. + /// + /// # Examples + /// + /// Fill an uninit vec with the default value. + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut buf = vec![MaybeUninit::::uninit(); 10]; + /// let initialized = MaybeUninit::fill_with(buf.as_mut_slice(), Default::default); + /// assert_eq!(initialized, &mut [0; 10]); + /// ``` + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill_with<'a, F>(this: &'a mut [MaybeUninit], mut f: F) -> &'a mut [T] + where + F: FnMut() -> T, + { + let mut guard = Guard { slice: this, initialized: 0 }; + + for element in guard.slice.iter_mut() { + element.write(f()); + guard.initialized += 1; + } + + super::forget(guard); + + // SAFETY: Valid elements have just been written into `this` so it is initialized + unsafe { MaybeUninit::slice_assume_init_mut(this) } + } + + /// Fills `this` with elements yielded by an iterator until either all elements have been + /// initialized or the iterator is empty. + /// + /// Returns two slices. The first slice contains the initialized portion of the original slice. + /// The second slice is the still-uninitialized remainder of the original slice. + /// + /// # Panics + /// + /// This function panics if the iterator's `next` function panics. + /// + /// If such a panic occurs, any elements previously initialized during this operation will be + /// dropped. + /// + /// # Examples + /// + /// Fill an uninit vec with a cycling iterator. + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut buf = vec![MaybeUninit::uninit(); 5]; + /// + /// let iter = [1, 2, 3].into_iter().cycle(); + /// let (initialized, remainder) = MaybeUninit::fill_from(&mut buf, iter); + /// + /// assert_eq!(initialized, &mut [1, 2, 3, 1, 2]); + /// assert_eq!(0, remainder.len()); + /// ``` + /// + /// Fill an uninit vec, but not completely. + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut buf = vec![MaybeUninit::uninit(); 5]; + /// let iter = [1, 2]; + /// let (initialized, remainder) = MaybeUninit::fill_from(&mut buf, iter); + /// + /// assert_eq!(initialized, &mut [1, 2]); + /// assert_eq!(remainder.len(), 3); + /// ``` + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill_from<'a, I>( + this: &'a mut [MaybeUninit], + it: I, + ) -> (&'a mut [T], &'a mut [MaybeUninit]) + where + I: IntoIterator, + { + let iter = it.into_iter(); + let mut guard = Guard { slice: this, initialized: 0 }; + + for (element, val) in guard.slice.iter_mut().zip(iter) { + element.write(val); + guard.initialized += 1; + } + + let initialized_len = guard.initialized; + super::forget(guard); + + // SAFETY: guard.initialized <= this.len() + let (initted, remainder) = unsafe { this.split_at_mut_unchecked(initialized_len) }; + + // SAFETY: Valid elements have just been written into `init`, so that portion + // of `this` is initialized. + (unsafe { MaybeUninit::slice_assume_init_mut(initted) }, remainder) + } + /// Returns the contents of this `MaybeUninit` as a slice of potentially uninitialized bytes. /// /// Note that even if the contents of a `MaybeUninit` have been initialized, the value may still @@ -1315,3 +1444,44 @@ impl [MaybeUninit; N] { unsafe { intrinsics::transmute_unchecked(self) } } } + +struct Guard<'a, T> { + slice: &'a mut [MaybeUninit], + initialized: usize, +} + +impl<'a, T> Drop for Guard<'a, T> { + fn drop(&mut self) { + let initialized_part = &mut self.slice[..self.initialized]; + // SAFETY: this raw sub-slice will contain only initialized objects. + unsafe { + crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(initialized_part)); + } + } +} + +trait SpecFill { + fn spec_fill(&mut self, value: T); +} + +impl SpecFill for [MaybeUninit] { + default fn spec_fill(&mut self, value: T) { + let mut guard = Guard { slice: self, initialized: 0 }; + + if let Some((last, elems)) = guard.slice.split_last_mut() { + for el in elems { + el.write(value.clone()); + guard.initialized += 1; + } + + last.write(value); + } + super::forget(guard); + } +} + +impl SpecFill for [MaybeUninit] { + fn spec_fill(&mut self, value: T) { + self.fill(MaybeUninit::new(value)); + } +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index fa0e9a979d060..c5a7e87c4aa4f 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -54,6 +54,7 @@ #![feature(slice_from_ptr_range)] #![feature(slice_split_once)] #![feature(split_as_slice)] +#![feature(maybe_uninit_fill)] #![feature(maybe_uninit_uninit_array)] #![feature(maybe_uninit_write_slice)] #![feature(maybe_uninit_uninit_array_transpose)] diff --git a/library/core/tests/mem.rs b/library/core/tests/mem.rs index 0f7fde747690a..e388800f400df 100644 --- a/library/core/tests/mem.rs +++ b/library/core/tests/mem.rs @@ -308,21 +308,226 @@ fn uninit_write_slice_cloned_mid_panic() { } } +#[derive(Clone)] +struct Bomb; + +impl Drop for Bomb { + fn drop(&mut self) { + panic!("dropped a bomb! kaboom!") + } +} + #[test] fn uninit_write_slice_cloned_no_drop() { - #[derive(Clone)] - struct Bomb; + let mut dst = [MaybeUninit::uninit()]; + let src = [Bomb]; + + MaybeUninit::clone_from_slice(&mut dst, &src); + + forget(src); +} + +#[test] +fn uninit_fill() { + let mut dst = [MaybeUninit::new(255); 64]; + let expect = [0; 64]; + + assert_eq!(MaybeUninit::fill(&mut dst, 0), &expect); +} + +#[cfg(panic = "unwind")] +struct CloneUntilPanic { + limit: usize, + rc: Rc<()>, +} - impl Drop for Bomb { - fn drop(&mut self) { - panic!("dropped a bomb! kaboom") +#[cfg(panic = "unwind")] +impl Clone for CloneUntilPanic { + fn clone(&self) -> Self { + if Rc::strong_count(&self.rc) >= self.limit { + panic!("expected panic on clone"); } + Self { limit: self.limit, rc: self.rc.clone() } } +} + +#[test] +#[cfg(panic = "unwind")] +fn uninit_fill_clone_panic_drop() { + use std::panic; + + let rc = Rc::new(()); + + let mut dst = [MaybeUninit::uninit(), MaybeUninit::uninit(), MaybeUninit::uninit()]; + + let src = CloneUntilPanic { limit: 3, rc: rc.clone() }; + let err = panic::catch_unwind(panic::AssertUnwindSafe(|| { + MaybeUninit::fill(&mut dst, src); + })); + + match err { + Ok(_) => unreachable!(), + Err(payload) => { + payload + .downcast::<&'static str>() + .and_then(|s| if *s == "expected panic on clone" { Ok(s) } else { Err(s) }) + .unwrap_or_else(|p| panic::resume_unwind(p)); + assert_eq!(Rc::strong_count(&rc), 1) + } + } +} + +#[test] +#[cfg(panic = "unwind")] +fn uninit_fill_clone_no_drop_clones() { + let mut dst = [MaybeUninit::uninit(), MaybeUninit::uninit(), MaybeUninit::uninit()]; + + MaybeUninit::fill(&mut dst, Bomb); +} + +#[test] +fn uninit_fill_with() { + let mut dst = [MaybeUninit::new(255); 64]; + let expect = [0; 64]; + + assert_eq!(MaybeUninit::fill_with(&mut dst, || 0), &expect); +} + +#[test] +#[cfg(panic = "unwind")] +fn uninit_fill_with_mid_panic() { + use std::panic; + + let rc = Rc::new(()); + + let mut dst = [MaybeUninit::uninit(), MaybeUninit::uninit(), MaybeUninit::uninit()]; + + let src = CloneUntilPanic { limit: 3, rc: rc.clone() }; + let err = panic::catch_unwind(panic::AssertUnwindSafe(|| { + MaybeUninit::fill_with(&mut dst, || src.clone()); + })); + + drop(src); + + match err { + Ok(_) => unreachable!(), + Err(payload) => { + payload + .downcast::<&'static str>() + .and_then(|s| if *s == "expected panic on clone" { Ok(s) } else { Err(s) }) + .unwrap_or_else(|p| panic::resume_unwind(p)); + assert_eq!(Rc::strong_count(&rc), 1) + } + } +} + +#[test] +#[cfg(panic = "unwind")] +fn uninit_fill_with_no_drop() { + let mut dst = [MaybeUninit::uninit()]; + let src = Bomb; + + MaybeUninit::fill_with(&mut dst, || src.clone()); + + forget(src); +} + +#[test] +fn uninit_fill_from() { + let mut dst = [MaybeUninit::new(255); 64]; + let src = [0; 64]; + + let (initted, remainder) = MaybeUninit::fill_from(&mut dst, src.into_iter()); + assert_eq!(initted, &src); + assert_eq!(remainder.len(), 0); +} + +#[test] +fn uninit_fill_from_partial() { + let mut dst = [MaybeUninit::new(255); 64]; + let src = [0; 48]; + + let (initted, remainder) = MaybeUninit::fill_from(&mut dst, src.into_iter()); + assert_eq!(initted, &src); + assert_eq!(remainder.len(), 16); +} + +#[test] +fn uninit_over_fill() { + let mut dst = [MaybeUninit::new(255); 64]; + let src = [0; 72]; + + let (initted, remainder) = MaybeUninit::fill_from(&mut dst, src.into_iter()); + assert_eq!(initted, &src[0..64]); + assert_eq!(remainder.len(), 0); +} + +#[test] +fn uninit_empty_fill() { + let mut dst = [MaybeUninit::new(255); 64]; + let src = [0; 0]; + + let (initted, remainder) = MaybeUninit::fill_from(&mut dst, src.into_iter()); + assert_eq!(initted, &src[0..0]); + assert_eq!(remainder.len(), 64); +} + +#[test] +#[cfg(panic = "unwind")] +fn uninit_fill_from_mid_panic() { + use std::panic; + + struct IterUntilPanic { + limit: usize, + rc: Rc<()>, + } + + impl Iterator for IterUntilPanic { + type Item = Rc<()>; + fn next(&mut self) -> Option { + if Rc::strong_count(&self.rc) >= self.limit { + panic!("expected panic on next"); + } + Some(self.rc.clone()) + } + } + + let rc = Rc::new(()); + + let mut dst = [ + MaybeUninit::uninit(), + MaybeUninit::uninit(), + MaybeUninit::uninit(), + MaybeUninit::uninit(), + ]; + + let src = IterUntilPanic { limit: 3, rc: rc.clone() }; + + let err = panic::catch_unwind(panic::AssertUnwindSafe(|| { + MaybeUninit::fill_from(&mut dst, src); + })); + + match err { + Ok(_) => unreachable!(), + Err(payload) => { + payload + .downcast::<&'static str>() + .and_then(|s| if *s == "expected panic on next" { Ok(s) } else { Err(s) }) + .unwrap_or_else(|p| panic::resume_unwind(p)); + + assert_eq!(Rc::strong_count(&rc), 1) + } + } +} + +#[test] +#[cfg(panic = "unwind")] +fn uninit_fill_from_no_drop() { let mut dst = [MaybeUninit::uninit()]; let src = [Bomb]; - MaybeUninit::clone_from_slice(&mut dst, &src); + MaybeUninit::fill_from(&mut dst, src.iter()); forget(src); }