From 0d8b5551b9ee16d62f2dffc325777ac839260cae Mon Sep 17 00:00:00 2001 From: John Millikin Date: Tue, 31 Oct 2023 11:54:41 +0900 Subject: [PATCH] Implement `MaybeUninit::fill{,_cloned,_mut,_with,_from}` ACP: https://github.com/rust-lang/libs-team/issues/156 --- library/core/src/mem/maybe_uninit.rs | 230 +++++++++++++++++++++++++-- library/core/tests/lib.rs | 1 + library/core/tests/mem.rs | 219 ++++++++++++++++++++++++- 3 files changed, 426 insertions(+), 24 deletions(-) diff --git a/library/core/src/mem/maybe_uninit.rs b/library/core/src/mem/maybe_uninit.rs index 855bb1675c59d..1d1228b9c6650 100644 --- a/library/core/src/mem/maybe_uninit.rs +++ b/library/core/src/mem/maybe_uninit.rs @@ -1123,22 +1123,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 @@ -1147,7 +1131,7 @@ impl MaybeUninit { let src = &src[..len]; // guard is needed b/c panic might happen during a clone - let mut guard = Guard { slice: this, initialized: 0 }; + let mut guard = CloneGuard { slice: this, initialized: 0 }; for i in 0..len { guard.slice[i].write(src[i].clone()); @@ -1160,6 +1144,202 @@ impl MaybeUninit { unsafe { MaybeUninit::slice_assume_init_mut(this) } } + /// Fills `this` with elements by copying `value`, returning a reference to + /// the now initialized contents of `this`. + /// + /// This is similar to [`slice::fill`] but is restricted to `Copy` values. + /// Use [`MaybeUninit::fill_cloned`] to initialize from a `Clone` value. + /// + /// # Examples + /// + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut dst = [MaybeUninit::uninit(); 5]; + /// let init = MaybeUninit::fill(&mut dst, 0u8); + /// + /// assert_eq!(init, &[0, 0, 0, 0, 0]); + /// ``` + #[doc(alias = "memset")] + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill<'a>(this: &'a mut [MaybeUninit], value: T) -> &'a [T] + where + T: Copy, + { + this.fill(MaybeUninit::new(value)); + // SAFETY: Valid elements have just been copied into `this` so it is initialized + unsafe { MaybeUninit::slice_assume_init_ref(this) } + } + + /// Fills `this` with elements by copying `value`, returning a mutable + /// reference to the now initialized contents of `this`. + /// + /// This is similar to [`slice::fill`] but is restricted to `Copy` values. + /// Use [`MaybeUninit::fill_cloned`] to initialize from a `Clone` value. + /// + /// # Examples + /// + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut dst = [MaybeUninit::uninit(); 5]; + /// let init = MaybeUninit::fill_mut(&mut dst, 0u8); + /// init[4] = 123; + /// + /// assert_eq!(init, &[0, 0, 0, 0, 123]); + /// ``` + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill_mut<'a>(this: &'a mut [MaybeUninit], value: T) -> &'a mut [T] + where + T: Copy, + { + this.fill(MaybeUninit::new(value)); + // SAFETY: Valid elements have just been copied into `this` so it is initialized + unsafe { MaybeUninit::slice_assume_init_mut(this) } + } + + /// Fills `this` with elements by cloning `value`, returning a reference to + /// the now initialized contents of `this`. Any already initialized elements + /// will not be dropped. + /// + /// This is similar to [`slice::fill`] but does not drop existing elements. + /// + /// # Panics + /// + /// This function will panic if the implementation of `Clone` panics. + /// + /// If there is a panic, the already initialized elements will be dropped. + /// + /// # Examples + /// + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut dst = [ + /// MaybeUninit::uninit(), + /// MaybeUninit::uninit(), + /// MaybeUninit::uninit(), + /// ]; + /// let msg = String::from("hello"); + /// let init = MaybeUninit::fill_cloned(&mut dst, &msg); + /// + /// assert_eq!(init, &["hello", "hello", "hello"]); + /// ``` + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill_cloned<'a>(this: &'a mut [MaybeUninit], value: &T) -> &'a [T] + where + T: Clone, + { + let len = this.len(); + + // guard is needed b/c panic might happen during a clone + let mut guard = CloneGuard { slice: this, initialized: 0 }; + + for i in 0..len { + guard.slice[i].write(value.clone()); + guard.initialized += 1; + } + + super::forget(guard); + + // SAFETY: Valid elements have just been written into `this` so it is initialized + unsafe { MaybeUninit::slice_assume_init_ref(this) } + } + + /// Fills `this` with elements returned by calling a closure repeatedly, + /// returning a reference to the now initialized contents of `this`. Any + /// already initialized elements will not be dropped. + /// + /// This is similar to [`slice::fill_with`] but does not drop existing + /// elements. + /// + /// # Panics + /// + /// This function will panic if the closure panics. + /// + /// If there is a panic, the already initialized elements will be dropped. + /// + /// # Examples + /// + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut dst = [MaybeUninit::uninit(); 5]; + /// let mut next = 0; + /// let init = MaybeUninit::fill_with(&mut dst, || { next += 1; next }); + /// + /// assert_eq!(init, &[1, 2, 3, 4, 5]); + /// ``` + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill_with<'a, F>(this: &'a mut [MaybeUninit], mut f: F) -> &'a [T] + where + F: FnMut() -> T, + { + let len = this.len(); + let mut guard = CloneGuard { slice: this, initialized: 0 }; + + for i in 0..len { + guard.slice[i].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_ref(this) } + } + + /// Fills `this` with the contents of an iterator, returning a reference to + /// the now initialized contents of `this`. Any already initialized elements + /// will not be dropped. + /// + /// # Panics + /// + /// This function will panic if the iterator panics. + /// + /// If there is a panic, the already initialized elements will be dropped. + /// + /// # Examples + /// + /// ``` + /// #![feature(maybe_uninit_fill)] + /// use std::mem::MaybeUninit; + /// + /// let mut dst = [MaybeUninit::uninit(); 5]; + /// let mut iter = [1, 2, 3].into_iter().cycle(); + /// let init = MaybeUninit::fill_from(&mut dst, iter); + /// + /// assert_eq!(init, &[1, 2, 3, 1, 2]); + /// ``` + #[unstable(feature = "maybe_uninit_fill", issue = "117428")] + pub fn fill_from<'a, I>(this: &'a mut [MaybeUninit], mut iter: I) -> &'a [T] + where + I: Iterator, + { + let len = this.len(); + let mut guard = CloneGuard { slice: this, initialized: 0 }; + + for i in 0..len { + match iter.next() { + Some(value) => { + guard.slice[i].write(value); + guard.initialized += 1; + } + None => break, + } + } + + let init_len = guard.initialized; + super::forget(guard); + + // SAFETY: Valid elements have just been written into `this` so it is initialized + unsafe { MaybeUninit::slice_assume_init_ref(&mut this[..init_len]) } + } + /// 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 @@ -1313,3 +1493,19 @@ impl [MaybeUninit; N] { unsafe { intrinsics::transmute_unchecked(self) } } } + +struct CloneGuard<'a, T> { + slice: &'a mut [MaybeUninit], + initialized: usize, +} + +impl<'a, T> Drop for CloneGuard<'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)); + } + } +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index df7b34ce73b42..eca3df8527924 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -49,6 +49,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 5c2e18745ea21..4a115c626dd46 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 DropBomb; + +impl Drop for DropBomb { + 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 = [DropBomb]; + + MaybeUninit::write_slice_cloned(&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); +} + +#[test] +fn uninit_fill_mut() { + let mut dst = [MaybeUninit::new(255); 64]; + let expect = [0; 64]; + + let init: &mut [u8] = MaybeUninit::fill_mut(&mut dst, 0); + assert_eq!(init, &expect); +} - impl Drop for Bomb { - fn drop(&mut self) { - panic!("dropped a bomb! kaboom") +#[test] +fn uninit_fill_cloned() { + let mut dst = [MaybeUninit::new(255); 64]; + let expect = [0; 64]; + + assert_eq!(MaybeUninit::fill_cloned(&mut dst, &0), &expect); +} + +struct IncrementUntilPanic { + limit: usize, + rc: Rc<()>, +} + +impl Clone for IncrementUntilPanic { + 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_cloned_mid_panic() { + use std::panic; + + let rc = Rc::new(()); + + let mut dst = [ + MaybeUninit::uninit(), + MaybeUninit::uninit(), + MaybeUninit::uninit(), + ]; + + let src = IncrementUntilPanic { limit: 3, rc: rc.clone() }; + let err = panic::catch_unwind(panic::AssertUnwindSafe(|| { + MaybeUninit::fill_cloned(&mut dst, &src); + })); + + 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_cloned_no_drop() { let mut dst = [MaybeUninit::uninit()]; - let src = [Bomb]; + let src = DropBomb; - MaybeUninit::write_slice_cloned(&mut dst, &src); + MaybeUninit::fill_cloned(&mut dst, &src); + + forget(src); +} + +#[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 = IncrementUntilPanic { 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 = DropBomb; + + 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]; + + assert_eq!(MaybeUninit::fill_from(&mut dst, src.into_iter()), &src); +} + +#[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 = [DropBomb]; + + MaybeUninit::fill_from(&mut dst, src.iter()); forget(src); }