mathieu/granite-rust
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Implement MaybeUninit::fill{,_with,_from}

Browse source 
ACP: rust-lang/libs-team#156

Signed-off-by: Andrew Wock <ajwock@gmail.com>
This commit is contained in:
Andrew Wock 2024-02-18 18:06:16 -05:00
parent c7beecf3e3
commit 0a0074980f
3 changed files with 401 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

202
library/core/src/mem/maybe_uninit.rs
View file

@ -1125,22 +1125,6 @@ impl<T> MaybeUninit<T> {
// unlike copy_from_slice this does not call clone_from_slice on the slice
// this is because `MaybeUninit<T: Clone>` does not implement Clone.
struct Guard<'a, T> {
slice: &'a mut [MaybeUninit<T>],
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<T> MaybeUninit<T> {
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<T>], 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::<i32>::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<T>], 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<T>],
it: I,
) -> (&'a mut [T], &'a mut [MaybeUninit<T>])
where
I: IntoIterator<Item = T>,
{
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<T, const N: usize> [MaybeUninit<T>; N] {
unsafe { intrinsics::transmute_unchecked(self) }
}
}
struct Guard<'a, T> {
slice: &'a mut [MaybeUninit<T>],
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<T> {
fn spec_fill(&mut self, value: T);
}
impl<T: Clone> SpecFill<T> for [MaybeUninit<T>] {
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<T: Copy> SpecFill<T> for [MaybeUninit<T>] {
fn spec_fill(&mut self, value: T) {
self.fill(MaybeUninit::new(value));
}
}

1
library/core/tests/lib.rs
View file

@ -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)]

223
library/core/tests/mem.rs
View file

@ -308,17 +308,17 @@ 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;
impl Drop for Bomb {
fn drop(&mut self) {
panic!("dropped a bomb! kaboom")
}
}
let mut dst = [MaybeUninit::uninit()];
let src = [Bomb];
@ -327,6 +327,211 @@ fn uninit_write_slice_cloned_no_drop() {
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<()>,
}
#[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<Self::Item> {
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::fill_from(&mut dst, src.iter());
forget(src);
}
#[test]
fn uninit_const_assume_init_read() {
const FOO: u32 = unsafe { MaybeUninit::new(42).assume_init_read() };