Auto merge of rust-lang#74850 - TimDiekmann:remove-in-place-alloc, r=Amanieu

Remove in-place allocation and revert to separate methods for zeroed allocations

closes rust-lang/wg-allocators#58

Author: bors

library/alloc/src/alloc.rs

@@ -164,25 +164,34 @@ pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
 #[unstable(feature = "allocator_api", issue = "32838")]
 unsafe impl AllocRef for Global {
     #[inline]
-    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
-        unsafe {
-            let size = layout.size();
-            if size == 0 {
-                Ok(MemoryBlock { ptr: layout.dangling(), size: 0 })
-            } else {
-                let raw_ptr = match init {
-                    AllocInit::Uninitialized => alloc(layout),
-                    AllocInit::Zeroed => alloc_zeroed(layout),
-                };
-                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
-                Ok(MemoryBlock { ptr, size })
-            }
-        }
+    fn alloc(&mut self, layout: Layout) -> Result<MemoryBlock, AllocErr> {
+        let size = layout.size();
+        let ptr = if size == 0 {
+            layout.dangling()
+        } else {
+            // SAFETY: `layout` is non-zero in size,
+            unsafe { NonNull::new(alloc(layout)).ok_or(AllocErr)? }
+        };
+        Ok(MemoryBlock { ptr, size })
+    }
+
+    #[inline]
+    fn alloc_zeroed(&mut self, layout: Layout) -> Result<MemoryBlock, AllocErr> {
+        let size = layout.size();
+        let ptr = if size == 0 {
+            layout.dangling()
+        } else {
+            // SAFETY: `layout` is non-zero in size,
+            unsafe { NonNull::new(alloc_zeroed(layout)).ok_or(AllocErr)? }
+        };
+        Ok(MemoryBlock { ptr, size })
     }
 
     #[inline]
     unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
         if layout.size() != 0 {
+            // SAFETY: `layout` is non-zero in size,
+            // other conditions must be upheld by the caller
             unsafe { dealloc(ptr.as_ptr(), layout) }
         }
     }
@@ -193,38 +202,55 @@ unsafe impl AllocRef for Global {
         ptr: NonNull<u8>,
         layout: Layout,
         new_size: usize,
-        placement: ReallocPlacement,
-        init: AllocInit,
     ) -> Result<MemoryBlock, AllocErr> {
-        let size = layout.size();
         debug_assert!(
-            new_size >= size,
-            "`new_size` must be greater than or equal to `memory.size()`"
+            new_size >= layout.size(),
+            "`new_size` must be greater than or equal to `layout.size()`"
         );
 
-        if size == new_size {
-            return Ok(MemoryBlock { ptr, size });
+        // SAFETY: `new_size` must be non-zero, which is checked in the match expression.
+        // Other conditions must be upheld by the caller
+        unsafe {
+            match layout.size() {
+                old_size if old_size == new_size => Ok(MemoryBlock { ptr, size: new_size }),
+                0 => self.alloc(Layout::from_size_align_unchecked(new_size, layout.align())),
+                old_size => {
+                    // `realloc` probably checks for `new_size > size` or something similar.
+                    intrinsics::assume(new_size > old_size);
+                    let raw_ptr = realloc(ptr.as_ptr(), layout, new_size);
+                    let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                    Ok(MemoryBlock { ptr, size: new_size })
+                }
+            }
         }
+    }
 
-        match placement {
-            ReallocPlacement::InPlace => Err(AllocErr),
-            ReallocPlacement::MayMove if layout.size() == 0 => {
-                let new_layout =
-                    unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
-                self.alloc(new_layout, init)
-            }
-            ReallocPlacement::MayMove => {
-                // `realloc` probably checks for `new_size > size` or something similar.
-                let ptr = unsafe {
-                    intrinsics::assume(new_size > size);
-                    realloc(ptr.as_ptr(), layout, new_size)
-                };
-                let memory =
-                    MemoryBlock { ptr: NonNull::new(ptr).ok_or(AllocErr)?, size: new_size };
-                unsafe {
-                    init.init_offset(memory, size);
+    #[inline]
+    unsafe fn grow_zeroed(
+        &mut self,
+        ptr: NonNull<u8>,
+        layout: Layout,
+        new_size: usize,
+    ) -> Result<MemoryBlock, AllocErr> {
+        debug_assert!(
+            new_size >= layout.size(),
+            "`new_size` must be greater than or equal to `layout.size()`"
+        );
+
+        // SAFETY: `new_size` must be non-zero, which is checked in the match expression.
+        // Other conditions must be upheld by the caller
+        unsafe {
+            match layout.size() {
+                old_size if old_size == new_size => Ok(MemoryBlock { ptr, size: new_size }),
+                0 => self.alloc_zeroed(Layout::from_size_align_unchecked(new_size, layout.align())),
+                old_size => {
+                    // `realloc` probably checks for `new_size > size` or something similar.
+                    intrinsics::assume(new_size > old_size);
+                    let raw_ptr = realloc(ptr.as_ptr(), layout, new_size);
+                    raw_ptr.add(old_size).write_bytes(0, new_size - old_size);
+                    let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                    Ok(MemoryBlock { ptr, size: new_size })
                 }
-                Ok(memory)
             }
         }
     }
@@ -235,35 +261,34 @@ unsafe impl AllocRef for Global {
         ptr: NonNull<u8>,
         layout: Layout,
         new_size: usize,
-        placement: ReallocPlacement,
     ) -> Result<MemoryBlock, AllocErr> {
-        let size = layout.size();
+        let old_size = layout.size();
         debug_assert!(
-            new_size <= size,
-            "`new_size` must be smaller than or equal to `memory.size()`"
+            new_size <= old_size,
+            "`new_size` must be smaller than or equal to `layout.size()`"
         );
 
-        if size == new_size {
-            return Ok(MemoryBlock { ptr, size });
-        }
-
-        match placement {
-            ReallocPlacement::InPlace => Err(AllocErr),
-            ReallocPlacement::MayMove if new_size == 0 => {
-                unsafe {
-                    self.dealloc(ptr, layout);
-                }
-                Ok(MemoryBlock { ptr: layout.dangling(), size: 0 })
+        let ptr = if new_size == old_size {
+            ptr
+        } else if new_size == 0 {
+            // SAFETY: `layout` is non-zero in size as `old_size` != `new_size`
+            // Other conditions must be upheld by the caller
+            unsafe {
+                self.dealloc(ptr, layout);
             }
-            ReallocPlacement::MayMove => {
-                // `realloc` probably checks for `new_size < size` or something similar.
-                let ptr = unsafe {
-                    intrinsics::assume(new_size < size);
-                    realloc(ptr.as_ptr(), layout, new_size)
-                };
-                Ok(MemoryBlock { ptr: NonNull::new(ptr).ok_or(AllocErr)?, size: new_size })
-            }
-        }
+            layout.dangling()
+        } else {
+            // SAFETY: new_size is not zero,
+            // Other conditions must be upheld by the caller
+            let raw_ptr = unsafe {
+                // `realloc` probably checks for `new_size < old_size` or something similar.
+                intrinsics::assume(new_size < old_size);
+                realloc(ptr.as_ptr(), layout, new_size)
+            };
+            NonNull::new(raw_ptr).ok_or(AllocErr)?
+        };
+
+        Ok(MemoryBlock { ptr, size: new_size })
     }
 }
 
@@ -274,7 +299,7 @@ unsafe impl AllocRef for Global {
 #[inline]
 unsafe fn exchange_malloc(size: usize, align: usize) -> *mut u8 {
     let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
-    match Global.alloc(layout, AllocInit::Uninitialized) {
+    match Global.alloc(layout) {
         Ok(memory) => memory.ptr.as_ptr(),
         Err(_) => handle_alloc_error(layout),
     }

library/alloc/src/alloc/tests.rs

@@ -8,9 +8,8 @@ use test::Bencher;
 fn allocate_zeroed() {
     unsafe {
         let layout = Layout::from_size_align(1024, 1).unwrap();
-        let memory = Global
-            .alloc(layout.clone(), AllocInit::Zeroed)
-            .unwrap_or_else(|_| handle_alloc_error(layout));
+        let memory =
+            Global.alloc_zeroed(layout.clone()).unwrap_or_else(|_| handle_alloc_error(layout));
 
         let mut i = memory.ptr.cast::<u8>().as_ptr();
         let end = i.add(layout.size());

library/alloc/src/boxed.rs

@@ -146,7 +146,7 @@ use core::pin::Pin;
 use core::ptr::{self, Unique};
 use core::task::{Context, Poll};
 
-use crate::alloc::{self, AllocInit, AllocRef, Global};
+use crate::alloc::{self, AllocRef, Global};
 use crate::borrow::Cow;
 use crate::raw_vec::RawVec;
 use crate::str::from_boxed_utf8_unchecked;
@@ -197,11 +197,8 @@ impl<T> Box<T> {
     #[unstable(feature = "new_uninit", issue = "63291")]
     pub fn new_uninit() -> Box<mem::MaybeUninit<T>> {
         let layout = alloc::Layout::new::<mem::MaybeUninit<T>>();
-        let ptr = Global
-            .alloc(layout, AllocInit::Uninitialized)
-            .unwrap_or_else(|_| alloc::handle_alloc_error(layout))
-            .ptr
-            .cast();
+        let ptr =
+            Global.alloc(layout).unwrap_or_else(|_| alloc::handle_alloc_error(layout)).ptr.cast();
         unsafe { Box::from_raw(ptr.as_ptr()) }
     }
 
@@ -227,7 +224,7 @@ impl<T> Box<T> {
     pub fn new_zeroed() -> Box<mem::MaybeUninit<T>> {
         let layout = alloc::Layout::new::<mem::MaybeUninit<T>>();
         let ptr = Global
-            .alloc(layout, AllocInit::Zeroed)
+            .alloc_zeroed(layout)
             .unwrap_or_else(|_| alloc::handle_alloc_error(layout))
             .ptr
             .cast();

library/alloc/src/raw_vec.rs

@@ -8,18 +8,20 @@ use core::ops::Drop;
 use core::ptr::{NonNull, Unique};
 use core::slice;
 
-use crate::alloc::{
-    handle_alloc_error,
-    AllocInit::{self, *},
-    AllocRef, Global, Layout,
-    ReallocPlacement::{self, *},
-};
+use crate::alloc::{handle_alloc_error, AllocRef, Global, Layout};
 use crate::boxed::Box;
 use crate::collections::TryReserveError::{self, *};
 
 #[cfg(test)]
 mod tests;
 
+enum AllocInit {
+    /// The contents of the new memory are uninitialized.
+    Uninitialized,
+    /// The new memory is guaranteed to be zeroed.
+    Zeroed,
+}
+
 /// A low-level utility for more ergonomically allocating, reallocating, and deallocating
 /// a buffer of memory on the heap without having to worry about all the corner cases
 /// involved. This type is excellent for building your own data structures like Vec and VecDeque.
@@ -156,14 +158,14 @@ impl<T, A: AllocRef> RawVec<T, A> {
     /// allocator for the returned `RawVec`.
     #[inline]
     pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
-        Self::allocate_in(capacity, Uninitialized, alloc)
+        Self::allocate_in(capacity, AllocInit::Uninitialized, alloc)
     }
 
     /// Like `with_capacity_zeroed`, but parameterized over the choice
     /// of allocator for the returned `RawVec`.
     #[inline]
     pub fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self {
-        Self::allocate_in(capacity, Zeroed, alloc)
+        Self::allocate_in(capacity, AllocInit::Zeroed, alloc)
     }
 
     fn allocate_in(capacity: usize, init: AllocInit, mut alloc: A) -> Self {
@@ -180,7 +182,11 @@ impl<T, A: AllocRef> RawVec<T, A> {
                 Ok(_) => {}
                 Err(_) => capacity_overflow(),
             }
-            let memory = match alloc.alloc(layout, init) {
+            let result = match init {
+                AllocInit::Uninitialized => alloc.alloc(layout),
+                AllocInit::Zeroed => alloc.alloc_zeroed(layout),
+            };
+            let memory = match result {
                 Ok(memory) => memory,
                 Err(_) => handle_alloc_error(layout),
             };
@@ -358,7 +364,7 @@ impl<T, A: AllocRef> RawVec<T, A> {
     ///
     /// Aborts on OOM.
     pub fn shrink_to_fit(&mut self, amount: usize) {
-        match self.shrink(amount, MayMove) {
+        match self.shrink(amount) {
             Err(CapacityOverflow) => capacity_overflow(),
             Err(AllocError { layout, .. }) => handle_alloc_error(layout),
             Ok(()) => { /* yay */ }
@@ -450,22 +456,16 @@ impl<T, A: AllocRef> RawVec<T, A> {
         Ok(())
     }
 
-    fn shrink(
-        &mut self,
-        amount: usize,
-        placement: ReallocPlacement,
-    ) -> Result<(), TryReserveError> {
+    fn shrink(&mut self, amount: usize) -> Result<(), TryReserveError> {
         assert!(amount <= self.capacity(), "Tried to shrink to a larger capacity");
 
         let (ptr, layout) = if let Some(mem) = self.current_memory() { mem } else { return Ok(()) };
         let new_size = amount * mem::size_of::<T>();
 
         let memory = unsafe {
-            self.alloc.shrink(ptr, layout, new_size, placement).map_err(|_| {
-                TryReserveError::AllocError {
-                    layout: Layout::from_size_align_unchecked(new_size, layout.align()),
-                    non_exhaustive: (),
-                }
+            self.alloc.shrink(ptr, layout, new_size).map_err(|_| TryReserveError::AllocError {
+                layout: Layout::from_size_align_unchecked(new_size, layout.align()),
+                non_exhaustive: (),
             })?
         };
         self.set_memory(memory);
@@ -492,9 +492,9 @@ where
 
     let memory = if let Some((ptr, old_layout)) = current_memory {
         debug_assert_eq!(old_layout.align(), new_layout.align());
-        unsafe { alloc.grow(ptr, old_layout, new_layout.size(), MayMove, Uninitialized) }
+        unsafe { alloc.grow(ptr, old_layout, new_layout.size()) }
     } else {
-        alloc.alloc(new_layout, Uninitialized)
+        alloc.alloc(new_layout)
     }
     .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })?;
 

library/alloc/src/raw_vec/tests.rs

@@ -20,12 +20,12 @@ fn allocator_param() {
         fuel: usize,
     }
     unsafe impl AllocRef for BoundedAlloc {
-        fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
+        fn alloc(&mut self, layout: Layout) -> Result<MemoryBlock, AllocErr> {
             let size = layout.size();
             if size > self.fuel {
                 return Err(AllocErr);
             }
-            match Global.alloc(layout, init) {
+            match Global.alloc(layout) {
                 ok @ Ok(_) => {
                     self.fuel -= size;
                     ok

library/alloc/src/rc.rs

@@ -250,7 +250,7 @@ use core::pin::Pin;
 use core::ptr::{self, NonNull};
 use core::slice::from_raw_parts_mut;
 
-use crate::alloc::{box_free, handle_alloc_error, AllocInit, AllocRef, Global, Layout};
+use crate::alloc::{box_free, handle_alloc_error, AllocRef, Global, Layout};
 use crate::borrow::{Cow, ToOwned};
 use crate::string::String;
 use crate::vec::Vec;
@@ -928,9 +928,7 @@ impl<T: ?Sized> Rc<T> {
         let layout = Layout::new::<RcBox<()>>().extend(value_layout).unwrap().0.pad_to_align();
 
         // Allocate for the layout.
-        let mem = Global
-            .alloc(layout, AllocInit::Uninitialized)
-            .unwrap_or_else(|_| handle_alloc_error(layout));
+        let mem = Global.alloc(layout).unwrap_or_else(|_| handle_alloc_error(layout));
 
         // Initialize the RcBox
         let inner = mem_to_rcbox(mem.ptr.as_ptr());