Add new_uninit_slice and assume_init on Box, Rc, and Arc of [T]

3 files changed, 256 insertions, 14 deletions

diff --git a/src/liballoc/boxed.rs b/src/liballoc/boxed.rs
index 5ea7847ca45..a759260bd8f 100644
--- a/src/liballoc/boxed.rs
+++ b/src/liballoc/boxed.rs
@@ -91,6 +91,7 @@ use core::ops::{
     CoerceUnsized, DispatchFromDyn, Deref, DerefMut, Receiver, Generator, GeneratorState
 };
 use core::ptr::{self, NonNull, Unique};
+use core::slice;
 use core::task::{Context, Poll};
 
 use crate::alloc;
@@ -135,7 +136,7 @@ impl<T> Box<T> {
     ///     // Deferred initialization:
     ///     five.as_mut_ptr().write(5);
     ///
-    ///     Box::assume_init(five)
+    ///     five.assume_init()
     /// };
     ///
     /// assert_eq!(*five, 5)
@@ -148,6 +149,35 @@ impl<T> Box<T> {
         Box(unique.cast())
     }
 
+    /// Construct a new boxed slice with uninitialized contents.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    ///
+    /// let mut values = Box::<u32>::new_uninit_slice(3);
+    ///
+    /// let values = unsafe {
+    ///     // Deferred initialization:
+    ///     values[0].as_mut_ptr().write(1);
+    ///     values[1].as_mut_ptr().write(2);
+    ///     values[2].as_mut_ptr().write(3);
+    ///
+    ///     values.assume_init()
+    /// };
+    ///
+    /// assert_eq!(*values, [1, 2, 3])
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "0")]
+    pub fn new_uninit_slice(len: usize) -> Box<[mem::MaybeUninit<T>]> {
+        let layout = alloc::Layout::array::<mem::MaybeUninit<T>>(len).unwrap();
+        let ptr = unsafe { alloc::alloc(layout) };
+        let unique = Unique::new(ptr).unwrap_or_else(|| alloc::handle_alloc_error(layout));
+        let slice = unsafe { slice::from_raw_parts_mut(unique.cast().as_ptr(), len) };
+        Box(Unique::from(slice))
+    }
+
     /// Constructs a new `Pin<Box<T>>`. If `T` does not implement `Unpin`, then
     /// `x` will be pinned in memory and unable to be moved.
     #[stable(feature = "pin", since = "1.33.0")]
@@ -179,15 +209,51 @@ impl<T> Box<mem::MaybeUninit<T>> {
     ///     // Deferred initialization:
     ///     five.as_mut_ptr().write(5);
     ///
-    ///     Box::assume_init(five)
+    ///     five.assume_init()
     /// };
     ///
     /// assert_eq!(*five, 5)
     /// ```
     #[unstable(feature = "new_uninit", issue = "0")]
     #[inline]
-    pub unsafe fn assume_init(this: Self) -> Box<T> {
-        Box(Box::into_unique(this).cast())
+    pub unsafe fn assume_init(self) -> Box<T> {
+        Box(Box::into_unique(self).cast())
+    }
+}
+
+impl<T> Box<[mem::MaybeUninit<T>]> {
+    /// Convert to `Box<[T]>`.
+    ///
+    /// # Safety
+    ///
+    /// As with [`MaybeUninit::assume_init`],
+    /// it is up to the caller to guarantee that the values
+    /// really are in an initialized state.
+    /// Calling this when the content is not yet fully initialized
+    /// causes immediate undefined behavior.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    ///
+    /// let mut values = Box::<u32>::new_uninit_slice(3);
+    ///
+    /// let values = unsafe {
+    ///     // Deferred initialization:
+    ///     values[0].as_mut_ptr().write(1);
+    ///     values[1].as_mut_ptr().write(2);
+    ///     values[2].as_mut_ptr().write(3);
+    ///
+    ///     values.assume_init()
+    /// };
+    ///
+    /// assert_eq!(*values, [1, 2, 3])
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "0")]
+    #[inline]
+    pub unsafe fn assume_init(self) -> Box<[T]> {
+        Box(Unique::new_unchecked(Box::into_raw(self) as _))
     }
 }
 
diff --git a/src/liballoc/rc.rs b/src/liballoc/rc.rs
index c902580354d..404c86a24da 100644
--- a/src/liballoc/rc.rs
+++ b/src/liballoc/rc.rs
@@ -343,7 +343,7 @@ impl<T> Rc<T> {
     ///     // Deferred initialization:
     ///     Rc::get_mut_unchecked(&mut five).as_mut_ptr().write(5);
     ///
-    ///     Rc::assume_init(five)
+    ///     five.assume_init()
     /// };
     ///
     /// assert_eq!(*five, 5)
@@ -364,6 +364,50 @@ impl<T> Rc<T> {
         }
     }
 
+    /// Construct a new reference-counted slice with uninitialized contents.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    /// #![feature(get_mut_unchecked)]
+    ///
+    /// use std::rc::Rc;
+    ///
+    /// let mut values = Rc::<u32>::new_uninit_slice(3);
+    ///
+    /// let values = unsafe {
+    ///     // Deferred initialization:
+    ///     Rc::get_mut_unchecked(&mut values)[0].as_mut_ptr().write(1);
+    ///     Rc::get_mut_unchecked(&mut values)[1].as_mut_ptr().write(2);
+    ///     Rc::get_mut_unchecked(&mut values)[2].as_mut_ptr().write(3);
+    ///
+    ///     values.assume_init()
+    /// };
+    ///
+    /// assert_eq!(*values, [1, 2, 3])
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "0")]
+    pub fn new_uninit_slice(len: usize) -> Rc<[mem::MaybeUninit<T>]> {
+        let data_layout = Layout::array::<mem::MaybeUninit<T>>(len).unwrap();
+        let (layout, offset) = Layout::new::<RcBox<()>>().extend(data_layout).unwrap();
+        unsafe {
+            let allocated_ptr = Global.alloc(layout)
+                .unwrap_or_else(|_| handle_alloc_error(layout))
+                .as_ptr();
+            let data_ptr = allocated_ptr.add(offset) as *mut mem::MaybeUninit<T>;
+            let slice: *mut [mem::MaybeUninit<T>] = from_raw_parts_mut(data_ptr, len);
+            let wide_ptr = slice as *mut RcBox<[mem::MaybeUninit<T>]>;
+            let wide_ptr = set_data_ptr(wide_ptr, allocated_ptr);
+            ptr::write(&mut (*wide_ptr).strong, Cell::new(1));
+            ptr::write(&mut (*wide_ptr).weak, Cell::new(1));
+            Rc {
+                ptr: NonNull::new_unchecked(wide_ptr),
+                phantom: PhantomData,
+            }
+        }
+    }
+
     /// Constructs a new `Pin<Rc<T>>`. If `T` does not implement `Unpin`, then
     /// `value` will be pinned in memory and unable to be moved.
     #[stable(feature = "pin", since = "1.33.0")]
@@ -439,16 +483,60 @@ impl<T> Rc<mem::MaybeUninit<T>> {
     ///     // Deferred initialization:
     ///     Rc::get_mut_unchecked(&mut five).as_mut_ptr().write(5);
     ///
-    ///     Rc::assume_init(five)
+    ///     five.assume_init()
     /// };
     ///
     /// assert_eq!(*five, 5)
     /// ```
     #[unstable(feature = "new_uninit", issue = "0")]
     #[inline]
-    pub unsafe fn assume_init(this: Self) -> Rc<T> {
-        let ptr = this.ptr.cast();
-        mem::forget(this);
+    pub unsafe fn assume_init(self) -> Rc<T> {
+        let ptr = self.ptr.cast();
+        mem::forget(self);
+        Rc {
+            ptr,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<T> Rc<[mem::MaybeUninit<T>]> {
+    /// Convert to `Rc<[T]>`.
+    ///
+    /// # Safety
+    ///
+    /// As with [`MaybeUninit::assume_init`],
+    /// it is up to the caller to guarantee that the value
+    /// really is in an initialized state.
+    /// Calling this when the content is not yet fully initialized
+    /// causes immediate undefined behavior.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    /// #![feature(get_mut_unchecked)]
+    ///
+    /// use std::rc::Rc;
+    ///
+    /// let mut values = Rc::<u32>::new_uninit_slice(3);
+    ///
+    /// let values = unsafe {
+    ///     // Deferred initialization:
+    ///     Rc::get_mut_unchecked(&mut values)[0].as_mut_ptr().write(1);
+    ///     Rc::get_mut_unchecked(&mut values)[1].as_mut_ptr().write(2);
+    ///     Rc::get_mut_unchecked(&mut values)[2].as_mut_ptr().write(3);
+    ///
+    ///     values.assume_init()
+    /// };
+    ///
+    /// assert_eq!(*values, [1, 2, 3])
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "0")]
+    #[inline]
+    pub unsafe fn assume_init(self) -> Rc<[T]> {
+        let ptr = NonNull::new_unchecked(self.ptr.as_ptr() as _);
+        mem::forget(self);
         Rc {
             ptr,
             phantom: PhantomData,
diff --git a/src/liballoc/sync.rs b/src/liballoc/sync.rs
index 8c63c81ee27..3f47b9e6c7d 100644
--- a/src/liballoc/sync.rs
+++ b/src/liballoc/sync.rs
@@ -327,7 +327,7 @@ impl<T> Arc<T> {
     ///     // Deferred initialization:
     ///     Arc::get_mut_unchecked(&mut five).as_mut_ptr().write(5);
     ///
-    ///     Arc::assume_init(five)
+    ///     five.assume_init()
     /// };
     ///
     /// assert_eq!(*five, 5)
@@ -348,6 +348,50 @@ impl<T> Arc<T> {
         }
     }
 
+    /// Construct a new reference-counted slice with uninitialized contents.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    /// #![feature(get_mut_unchecked)]
+    ///
+    /// use std::sync::Arc;
+    ///
+    /// let mut values = Arc::<u32>::new_uninit_slice(3);
+    ///
+    /// let values = unsafe {
+    ///     // Deferred initialization:
+    ///     Arc::get_mut_unchecked(&mut values)[0].as_mut_ptr().write(1);
+    ///     Arc::get_mut_unchecked(&mut values)[1].as_mut_ptr().write(2);
+    ///     Arc::get_mut_unchecked(&mut values)[2].as_mut_ptr().write(3);
+    ///
+    ///     values.assume_init()
+    /// };
+    ///
+    /// assert_eq!(*values, [1, 2, 3])
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "0")]
+    pub fn new_uninit_slice(len: usize) -> Arc<[mem::MaybeUninit<T>]> {
+        let data_layout = Layout::array::<mem::MaybeUninit<T>>(len).unwrap();
+        let (layout, offset) = Layout::new::<ArcInner<()>>().extend(data_layout).unwrap();
+        unsafe {
+            let allocated_ptr = Global.alloc(layout)
+                .unwrap_or_else(|_| handle_alloc_error(layout))
+                .as_ptr();
+            let data_ptr = allocated_ptr.add(offset) as *mut mem::MaybeUninit<T>;
+            let slice: *mut [mem::MaybeUninit<T>] = from_raw_parts_mut(data_ptr, len);
+            let wide_ptr = slice as *mut ArcInner<[mem::MaybeUninit<T>]>;
+            let wide_ptr = set_data_ptr(wide_ptr, allocated_ptr);
+            ptr::write(&mut (*wide_ptr).strong, atomic::AtomicUsize::new(1));
+            ptr::write(&mut (*wide_ptr).weak, atomic::AtomicUsize::new(1));
+            Arc {
+                ptr: NonNull::new_unchecked(wide_ptr),
+                phantom: PhantomData,
+            }
+        }
+    }
+
     /// Constructs a new `Pin<Arc<T>>`. If `T` does not implement `Unpin`, then
     /// `data` will be pinned in memory and unable to be moved.
     #[stable(feature = "pin", since = "1.33.0")]
@@ -423,16 +467,60 @@ impl<T> Arc<mem::MaybeUninit<T>> {
     ///     // Deferred initialization:
     ///     Arc::get_mut_unchecked(&mut five).as_mut_ptr().write(5);
     ///
-    ///     Arc::assume_init(five)
+    ///     five.assume_init()
     /// };
     ///
     /// assert_eq!(*five, 5)
     /// ```
     #[unstable(feature = "new_uninit", issue = "0")]
     #[inline]
-    pub unsafe fn assume_init(this: Self) -> Arc<T> {
-        let ptr = this.ptr.cast();
-        mem::forget(this);
+    pub unsafe fn assume_init(self) -> Arc<T> {
+        let ptr = self.ptr.cast();
+        mem::forget(self);
+        Arc {
+            ptr,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<T> Arc<[mem::MaybeUninit<T>]> {
+    /// Convert to `Arc<[T]>`.
+    ///
+    /// # Safety
+    ///
+    /// As with [`MaybeUninit::assume_init`],
+    /// it is up to the caller to guarantee that the value
+    /// really is in an initialized state.
+    /// Calling this when the content is not yet fully initialized
+    /// causes immediate undefined behavior.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    /// #![feature(get_mut_unchecked)]
+    ///
+    /// use std::sync::Arc;
+    ///
+    /// let mut values = Arc::<u32>::new_uninit_slice(3);
+    ///
+    /// let values = unsafe {
+    ///     // Deferred initialization:
+    ///     Arc::get_mut_unchecked(&mut values)[0].as_mut_ptr().write(1);
+    ///     Arc::get_mut_unchecked(&mut values)[1].as_mut_ptr().write(2);
+    ///     Arc::get_mut_unchecked(&mut values)[2].as_mut_ptr().write(3);
+    ///
+    ///     values.assume_init()
+    /// };
+    ///
+    /// assert_eq!(*values, [1, 2, 3])
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "0")]
+    #[inline]
+    pub unsafe fn assume_init(self) -> Arc<[T]> {
+        let ptr = NonNull::new_unchecked(self.ptr.as_ptr() as _);
+        mem::forget(self);
         Arc {
             ptr,
             phantom: PhantomData,