Auto merge of #75657 - TimDiekmann:allocref-cleanup, r=Amanieu

Clean up AllocRef implementation and documentation

r? @Amanieu

library/alloc/src/alloc.rs

@@ -161,30 +161,69 @@ pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
     unsafe { __rust_alloc_zeroed(layout.size(), layout.align()) }
 }
 
+impl Global {
+    #[inline]
+    fn alloc_impl(&mut self, layout: Layout, zeroed: bool) -> Result<NonNull<[u8]>, AllocErr> {
+        match layout.size() {
+            0 => Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0)),
+            // SAFETY: `layout` is non-zero in size,
+            size => unsafe {
+                let raw_ptr = if zeroed { alloc_zeroed(layout) } else { alloc(layout) };
+                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                Ok(NonNull::slice_from_raw_parts(ptr, size))
+            },
+        }
+    }
+
+    // Safety: Same as `AllocRef::grow`
+    #[inline]
+    unsafe fn grow_impl(
+        &mut self,
+        ptr: NonNull<u8>,
+        layout: Layout,
+        new_size: usize,
+        zeroed: bool,
+    ) -> Result<NonNull<[u8]>, AllocErr> {
+        debug_assert!(
+            new_size >= layout.size(),
+            "`new_size` must be greater than or equal to `layout.size()`"
+        );
+
+        match layout.size() {
+            // SAFETY: the caller must ensure that the `new_size` does not overflow.
+            // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid for a Layout.
+            0 => unsafe {
+                let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
+                self.alloc_impl(new_layout, zeroed)
+            },
+
+            // SAFETY: `new_size` is non-zero as `old_size` is greater than or equal to `new_size`
+            // as required by safety conditions. Other conditions must be upheld by the caller
+            old_size => unsafe {
+                // `realloc` probably checks for `new_size >= size` or something similar.
+                intrinsics::assume(new_size >= layout.size());
+
+                let raw_ptr = realloc(ptr.as_ptr(), layout, new_size);
+                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                if zeroed {
+                    raw_ptr.add(old_size).write_bytes(0, new_size - old_size);
+                }
+                Ok(NonNull::slice_from_raw_parts(ptr, new_size))
+            },
+        }
+    }
+}
+
 #[unstable(feature = "allocator_api", issue = "32838")]
 unsafe impl AllocRef for Global {
     #[inline]
     fn alloc(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
-        let size = layout.size();
+        self.alloc_impl(layout, false)
-        let ptr = if size == 0 {
-            layout.dangling()
-        } else {
-            // SAFETY: `layout` is non-zero in size,
-            unsafe { NonNull::new(alloc(layout)).ok_or(AllocErr)? }
-        };
-        Ok(NonNull::slice_from_raw_parts(ptr, size))
     }
 
     #[inline]
     fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
-        let size = layout.size();
+        self.alloc_impl(layout, true)
-        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(NonNull::slice_from_raw_parts(ptr, size))
     }
 
     #[inline]
@@ -203,26 +242,8 @@ unsafe impl AllocRef for Global {
         layout: Layout,
         new_size: usize,
     ) -> Result<NonNull<[u8]>, AllocErr> {
-        debug_assert!(
+        // SAFETY: all conditions must be upheld by the caller
-            new_size >= layout.size(),
+        unsafe { self.grow_impl(ptr, layout, new_size, false) }
-            "`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.
-        // If `new_size` is zero, then `old_size` has to be zero as well.
-        // Other conditions must be upheld by the caller
-        unsafe {
-            match layout.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(NonNull::slice_from_raw_parts(ptr, new_size))
-                }
-            }
-        }
     }
 
     #[inline]
@@ -232,27 +253,8 @@ unsafe impl AllocRef for Global {
         layout: Layout,
         new_size: usize,
     ) -> Result<NonNull<[u8]>, AllocErr> {
-        debug_assert!(
+        // SAFETY: all conditions must be upheld by the caller
-            new_size >= layout.size(),
+        unsafe { self.grow_impl(ptr, layout, new_size, true) }
-            "`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.
-        // If `new_size` is zero, then `old_size` has to be zero as well.
-        // Other conditions must be upheld by the caller
-        unsafe {
-            match layout.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(NonNull::slice_from_raw_parts(ptr, new_size))
-                }
-            }
-        }
     }
 
     #[inline]
@@ -262,30 +264,28 @@ unsafe impl AllocRef for Global {
         layout: Layout,
         new_size: usize,
     ) -> Result<NonNull<[u8]>, AllocErr> {
-        let old_size = layout.size();
         debug_assert!(
-            new_size <= old_size,
+            new_size <= layout.size(),
             "`new_size` must be smaller than or equal to `layout.size()`"
         );
 
-        let ptr = if new_size == 0 {
+        match new_size {
             // SAFETY: conditions must be upheld by the caller
-            unsafe {
+            0 => unsafe {
                 self.dealloc(ptr, layout);
-            }
+                Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0))
-            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(NonNull::slice_from_raw_parts(ptr, new_size))
+            // SAFETY: `new_size` is non-zero. Other conditions must be upheld by the caller
+            new_size => unsafe {
+                // `realloc` probably checks for `new_size <= size` or something similar.
+                intrinsics::assume(new_size <= layout.size());
+
+                let raw_ptr = realloc(ptr.as_ptr(), layout, new_size);
+                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                Ok(NonNull::slice_from_raw_parts(ptr, new_size))
+            },
+        }
     }
 }
 

library/core/src/alloc/mod.rs

@@ -151,6 +151,11 @@ pub unsafe trait AllocRef {
     /// alignment and a size given by `new_size`. To accomplish this, the allocator may extend the
     /// allocation referenced by `ptr` to fit the new layout.
     ///
+    /// If this returns `Ok`, then ownership of the memory block referenced by `ptr` has been
+    /// transferred to this allocator. The memory may or may not have been freed, and should be
+    /// considered unusable unless it was transferred back to the caller again via the return value
+    /// of this method.
+    ///
     /// If this method returns `Err`, then ownership of the memory block has not been transferred to
     /// this allocator, and the contents of the memory block are unaltered.
     ///
@@ -192,12 +197,9 @@ pub unsafe trait AllocRef {
             "`new_size` must be greater than or equal to `layout.size()`"
         );
 
-        let new_layout =
+        // SAFETY: the caller must ensure that the `new_size` does not overflow.
-            // SAFETY: the caller must ensure that the `new_size` does not overflow.
+        // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid for a Layout.
-            // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid for a Layout.
+        let new_layout = unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
-            // The caller must ensure that `new_size` is greater than or equal to zero. If it's equal
-            // to zero, it's catched beforehand.
-            unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
         let new_ptr = self.alloc(new_layout)?;
 
         // SAFETY: because `new_size` must be greater than or equal to `size`, both the old and new
@@ -206,10 +208,11 @@ pub unsafe trait AllocRef {
         // `copy_nonoverlapping` is safe.
         // The safety contract for `dealloc` must be upheld by the caller.
         unsafe {
-            ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_non_null_ptr().as_ptr(), size);
+            ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_mut_ptr(), size);
             self.dealloc(ptr, layout);
-            Ok(new_ptr)
         }
+
+        Ok(new_ptr)
     }
 
     /// Behaves like `grow`, but also ensures that the new contents are set to zero before being
@@ -218,12 +221,12 @@ pub unsafe trait AllocRef {
     /// The memory block will contain the following contents after a successful call to
     /// `grow_zeroed`:
     ///   * Bytes `0..layout.size()` are preserved from the original allocation.
-    ///   * Bytes `layout.size()..old_size` will either be preserved or zeroed,
+    ///   * Bytes `layout.size()..old_size` will either be preserved or zeroed, depending on the
-    ///     depending on the allocator implementation. `old_size` refers to the size of
+    ///     allocator implementation. `old_size` refers to the size of the memory block prior to
-    ///     the `MemoryBlock` prior to the `grow_zeroed` call, which may be larger than the size
+    ///     the `grow_zeroed` call, which may be larger than the size that was originally requested
-    ///     that was originally requested when it was allocated.
+    ///     when it was allocated.
-    ///   * Bytes `old_size..new_size` are zeroed. `new_size` refers to
+    ///   * Bytes `old_size..new_size` are zeroed. `new_size` refers to the size of the memory
-    ///     the size of the `MemoryBlock` returned by the `grow` call.
+    ///     block returned by the `grow` call.
     ///
     /// # Safety
     ///
@@ -261,12 +264,9 @@ pub unsafe trait AllocRef {
             "`new_size` must be greater than or equal to `layout.size()`"
         );
 
-        let new_layout =
+        // SAFETY: the caller must ensure that the `new_size` does not overflow.
-            // SAFETY: the caller must ensure that the `new_size` does not overflow.
+        // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid for a Layout.
-            // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid for a Layout.
+        let new_layout = unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
-            // The caller must ensure that `new_size` is greater than or equal to zero. If it's equal
-            // to zero, it's caught beforehand.
-            unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
         let new_ptr = self.alloc_zeroed(new_layout)?;
 
         // SAFETY: because `new_size` must be greater than or equal to `size`, both the old and new
@@ -275,10 +275,11 @@ pub unsafe trait AllocRef {
         // `copy_nonoverlapping` is safe.
         // The safety contract for `dealloc` must be upheld by the caller.
         unsafe {
-            ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_non_null_ptr().as_ptr(), size);
+            ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_mut_ptr(), size);
             self.dealloc(ptr, layout);
-            Ok(new_ptr)
         }
+
+        Ok(new_ptr)
     }
 
     /// Attempts to shrink the memory block.
@@ -290,8 +291,8 @@ pub unsafe trait AllocRef {
     ///
     /// If this returns `Ok`, then ownership of the memory block referenced by `ptr` has been
     /// transferred to this allocator. The memory may or may not have been freed, and should be
-    /// considered unusable unless it was transferred back to the caller again via the
+    /// considered unusable unless it was transferred back to the caller again via the return value
-    /// return value of this method.
+    /// of this method.
     ///
     /// If this method returns `Err`, then ownership of the memory block has not been transferred to
     /// this allocator, and the contents of the memory block are unaltered.
@@ -332,11 +333,9 @@ pub unsafe trait AllocRef {
             "`new_size` must be smaller than or equal to `layout.size()`"
         );
 
-        let new_layout =
         // SAFETY: the caller must ensure that the `new_size` does not overflow.
         // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid for a Layout.
-        // The caller must ensure that `new_size` is greater than zero.
+        let new_layout = unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
-            unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
         let new_ptr = self.alloc(new_layout)?;
 
         // SAFETY: because `new_size` must be lower than or equal to `size`, both the old and new
@@ -345,10 +344,11 @@ pub unsafe trait AllocRef {
         // `copy_nonoverlapping` is safe.
         // The safety contract for `dealloc` must be upheld by the caller.
         unsafe {
-            ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_non_null_ptr().as_ptr(), size);
+            ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_mut_ptr(), size);
             self.dealloc(ptr, layout);
-            Ok(new_ptr)
         }
+
+        Ok(new_ptr)
     }
 
     /// Creates a "by reference" adaptor for this instance of `AllocRef`.

library/std/src/alloc.rs

@@ -131,32 +131,75 @@ pub use alloc_crate::alloc::*;
 #[derive(Debug, Default, Copy, Clone)]
 pub struct System;
 
+impl System {
+    #[inline]
+    fn alloc_impl(&mut self, layout: Layout, zeroed: bool) -> Result<NonNull<[u8]>, AllocErr> {
+        match layout.size() {
+            0 => Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0)),
+            // SAFETY: `layout` is non-zero in size,
+            size => unsafe {
+                let raw_ptr = if zeroed {
+                    GlobalAlloc::alloc_zeroed(self, layout)
+                } else {
+                    GlobalAlloc::alloc(self, layout)
+                };
+                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                Ok(NonNull::slice_from_raw_parts(ptr, size))
+            },
+        }
+    }
+
+    // Safety: Same as `AllocRef::grow`
+    #[inline]
+    unsafe fn grow_impl(
+        &mut self,
+        ptr: NonNull<u8>,
+        layout: Layout,
+        new_size: usize,
+        zeroed: bool,
+    ) -> Result<NonNull<[u8]>, AllocErr> {
+        debug_assert!(
+            new_size >= layout.size(),
+            "`new_size` must be greater than or equal to `layout.size()`"
+        );
+
+        match layout.size() {
+            // SAFETY: the caller must ensure that the `new_size` does not overflow.
+            // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid for a Layout.
+            0 => unsafe {
+                let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
+                self.alloc_impl(new_layout, zeroed)
+            },
+
+            // SAFETY: `new_size` is non-zero as `old_size` is greater than or equal to `new_size`
+            // as required by safety conditions. Other conditions must be upheld by the caller
+            old_size => unsafe {
+                // `realloc` probably checks for `new_size >= size` or something similar.
+                intrinsics::assume(new_size >= layout.size());
+
+                let raw_ptr = GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size);
+                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                if zeroed {
+                    raw_ptr.add(old_size).write_bytes(0, new_size - old_size);
+                }
+                Ok(NonNull::slice_from_raw_parts(ptr, new_size))
+            },
+        }
+    }
+}
+
 // The AllocRef impl checks the layout size to be non-zero and forwards to the GlobalAlloc impl,
 // which is in `std::sys::*::alloc`.
 #[unstable(feature = "allocator_api", issue = "32838")]
 unsafe impl AllocRef for System {
     #[inline]
     fn alloc(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
-        let size = layout.size();
+        self.alloc_impl(layout, false)
-        let ptr = if size == 0 {
-            layout.dangling()
-        } else {
-            // SAFETY: `layout` is non-zero in size,
-            unsafe { NonNull::new(GlobalAlloc::alloc(&System, layout)).ok_or(AllocErr)? }
-        };
-        Ok(NonNull::slice_from_raw_parts(ptr, size))
     }
 
     #[inline]
     fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
-        let size = layout.size();
+        self.alloc_impl(layout, true)
-        let ptr = if size == 0 {
-            layout.dangling()
-        } else {
-            // SAFETY: `layout` is non-zero in size,
-            unsafe { NonNull::new(GlobalAlloc::alloc_zeroed(&System, layout)).ok_or(AllocErr)? }
-        };
-        Ok(NonNull::slice_from_raw_parts(ptr, size))
     }
 
     #[inline]
@@ -164,7 +207,7 @@ unsafe impl AllocRef for System {
         if layout.size() != 0 {
             // SAFETY: `layout` is non-zero in size,
             // other conditions must be upheld by the caller
-            unsafe { GlobalAlloc::dealloc(&System, ptr.as_ptr(), layout) }
+            unsafe { GlobalAlloc::dealloc(self, ptr.as_ptr(), layout) }
         }
     }
 
@@ -175,26 +218,8 @@ unsafe impl AllocRef for System {
         layout: Layout,
         new_size: usize,
     ) -> Result<NonNull<[u8]>, AllocErr> {
-        debug_assert!(
+        // SAFETY: all conditions must be upheld by the caller
-            new_size >= layout.size(),
+        unsafe { self.grow_impl(ptr, layout, new_size, false) }
-            "`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.
-        // If `new_size` is zero, then `old_size` has to be zero as well.
-        // Other conditions must be upheld by the caller
-        unsafe {
-            match layout.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 = GlobalAlloc::realloc(&System, ptr.as_ptr(), layout, new_size);
-                    let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
-                    Ok(NonNull::slice_from_raw_parts(ptr, new_size))
-                }
-            }
-        }
     }
 
     #[inline]
@@ -204,27 +229,8 @@ unsafe impl AllocRef for System {
         layout: Layout,
         new_size: usize,
     ) -> Result<NonNull<[u8]>, AllocErr> {
-        debug_assert!(
+        // SAFETY: all conditions must be upheld by the caller
-            new_size >= layout.size(),
+        unsafe { self.grow_impl(ptr, layout, new_size, true) }
-            "`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.
-        // If `new_size` is zero, then `old_size` has to be zero as well.
-        // Other conditions must be upheld by the caller
-        unsafe {
-            match layout.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 = GlobalAlloc::realloc(&System, 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(NonNull::slice_from_raw_parts(ptr, new_size))
-                }
-            }
-        }
     }
 
     #[inline]
@@ -234,32 +240,31 @@ unsafe impl AllocRef for System {
         layout: Layout,
         new_size: usize,
     ) -> Result<NonNull<[u8]>, AllocErr> {
-        let old_size = layout.size();
         debug_assert!(
-            new_size <= old_size,
+            new_size <= layout.size(),
             "`new_size` must be smaller than or equal to `layout.size()`"
         );
 
-        let ptr = if new_size == 0 {
+        match new_size {
             // SAFETY: conditions must be upheld by the caller
-            unsafe {
+            0 => unsafe {
                 self.dealloc(ptr, layout);
-            }
+                Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0))
-            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);
-                GlobalAlloc::realloc(&System, ptr.as_ptr(), layout, new_size)
-            };
-            NonNull::new(raw_ptr).ok_or(AllocErr)?
-        };
 
-        Ok(NonNull::slice_from_raw_parts(ptr, new_size))
+            // SAFETY: `new_size` is non-zero. Other conditions must be upheld by the caller
+            new_size => unsafe {
+                // `realloc` probably checks for `new_size <= size` or something similar.
+                intrinsics::assume(new_size <= layout.size());
+
+                let raw_ptr = GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size);
+                let ptr = NonNull::new(raw_ptr).ok_or(AllocErr)?;
+                Ok(NonNull::slice_from_raw_parts(ptr, new_size))
+            },
+        }
     }
 }
+
 static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 
 /// Registers a custom allocation error hook, replacing any that was previously registered.