Allow ZSTs in `AllocRef`

5 files changed, 96 insertions, 81 deletions

diff --git a/src/liballoc/alloc.rs b/src/liballoc/alloc.rs
index 73e8121868a..9f82b2c6fa6 100644
--- a/src/liballoc/alloc.rs
+++ b/src/liballoc/alloc.rs
@@ -165,13 +165,19 @@ pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
 #[unstable(feature = "allocator_api", issue = "32838")]
 unsafe impl AllocRef for Global {
     #[inline]
-    unsafe fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
-        NonNull::new(alloc(layout)).ok_or(AllocErr).map(|p| (p, layout.size()))
+    fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
+        if layout.size() == 0 {
+            Ok((layout.dangling(), 0))
+        } else {
+            unsafe { NonNull::new(alloc(layout)).ok_or(AllocErr).map(|p| (p, layout.size())) }
+        }
     }
 
     #[inline]
     unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
-        dealloc(ptr.as_ptr(), layout)
+        if layout.size() != 0 {
+            dealloc(ptr.as_ptr(), layout)
+        }
     }
 
     #[inline]
@@ -181,12 +187,28 @@ unsafe impl AllocRef for Global {
         layout: Layout,
         new_size: usize,
     ) -> Result<(NonNull<u8>, usize), AllocErr> {
-        NonNull::new(realloc(ptr.as_ptr(), layout, new_size)).ok_or(AllocErr).map(|p| (p, new_size))
+        match (layout.size(), new_size) {
+            (0, 0) => Ok((layout.dangling(), 0)),
+            (0, _) => self.alloc(Layout::from_size_align_unchecked(new_size, layout.align())),
+            (_, 0) => {
+                self.dealloc(ptr, layout);
+                Ok((layout.dangling(), 0))
+            }
+            (_, _) => NonNull::new(realloc(ptr.as_ptr(), layout, new_size))
+                .ok_or(AllocErr)
+                .map(|p| (p, new_size)),
+        }
     }
 
     #[inline]
-    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
-        NonNull::new(alloc_zeroed(layout)).ok_or(AllocErr).map(|p| (p, layout.size()))
+    fn alloc_zeroed(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
+        if layout.size() == 0 {
+            Ok((layout.dangling(), 0))
+        } else {
+            unsafe {
+                NonNull::new(alloc_zeroed(layout)).ok_or(AllocErr).map(|p| (p, layout.size()))
+            }
+        }
     }
 }
 
diff --git a/src/liballoc/raw_vec.rs b/src/liballoc/raw_vec.rs
index 345834d7daa..b31fec7f037 100644
--- a/src/liballoc/raw_vec.rs
+++ b/src/liballoc/raw_vec.rs
@@ -73,30 +73,28 @@ impl<T, A: AllocRef> RawVec<T, A> {
     }
 
     fn allocate_in(mut capacity: usize, zeroed: bool, mut a: A) -> Self {
-        unsafe {
-            let elem_size = mem::size_of::<T>();
+        let elem_size = mem::size_of::<T>();
 
-            let alloc_size = capacity.checked_mul(elem_size).unwrap_or_else(|| capacity_overflow());
-            alloc_guard(alloc_size).unwrap_or_else(|_| capacity_overflow());
+        let alloc_size = capacity.checked_mul(elem_size).unwrap_or_else(|| capacity_overflow());
+        alloc_guard(alloc_size).unwrap_or_else(|_| capacity_overflow());
 
-            // Handles ZSTs and `capacity == 0` alike.
-            let ptr = if alloc_size == 0 {
-                NonNull::<T>::dangling()
-            } else {
-                let align = mem::align_of::<T>();
-                let layout = Layout::from_size_align(alloc_size, align).unwrap();
-                let result = if zeroed { a.alloc_zeroed(layout) } else { a.alloc(layout) };
-                match result {
-                    Ok((ptr, size)) => {
-                        capacity = size / elem_size;
-                        ptr.cast()
-                    }
-                    Err(_) => handle_alloc_error(layout),
+        // Handles ZSTs and `capacity == 0` alike.
+        let ptr = if alloc_size == 0 {
+            NonNull::<T>::dangling()
+        } else {
+            let align = mem::align_of::<T>();
+            let layout = Layout::from_size_align(alloc_size, align).unwrap();
+            let result = if zeroed { a.alloc_zeroed(layout) } else { a.alloc(layout) };
+            match result {
+                Ok((ptr, size)) => {
+                    capacity = size / elem_size;
+                    ptr.cast()
                 }
-            };
+                Err(_) => handle_alloc_error(layout),
+            }
+        };
 
-            RawVec { ptr: ptr.into(), cap: capacity, a }
-        }
+        RawVec { ptr: ptr.into(), cap: capacity, a }
     }
 }
 
diff --git a/src/liballoc/raw_vec/tests.rs b/src/liballoc/raw_vec/tests.rs
index 860058debe1..21a8a76d0a7 100644
--- a/src/liballoc/raw_vec/tests.rs
+++ b/src/liballoc/raw_vec/tests.rs
@@ -20,7 +20,7 @@ fn allocator_param() {
         fuel: usize,
     }
     unsafe impl AllocRef for BoundedAlloc {
-        unsafe fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
+        fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
             let size = layout.size();
             if size > self.fuel {
                 return Err(AllocErr);
diff --git a/src/libcore/alloc.rs b/src/libcore/alloc.rs
index 0a7a8ab266a..d2a513451cc 100644
--- a/src/libcore/alloc.rs
+++ b/src/libcore/alloc.rs
@@ -606,20 +606,11 @@ pub unsafe trait GlobalAlloc {
 ///   method (`dealloc`) or by being passed to a reallocation method
 ///  (see above) that returns `Ok`.
 ///
-/// A note regarding zero-sized types and zero-sized layouts: many
-/// methods in the `AllocRef` trait state that allocation requests
-/// must be non-zero size, or else undefined behavior can result.
-///
-/// * If an `AllocRef` implementation chooses to return `Ok` in this
-///   case (i.e., the pointer denotes a zero-sized inaccessible block)
-///   then that returned pointer must be considered "currently
-///   allocated". On such an allocator, *all* methods that take
-///   currently-allocated pointers as inputs must accept these
-///   zero-sized pointers, *without* causing undefined behavior.
-///
-/// * In other words, if a zero-sized pointer can flow out of an
-///   allocator, then that allocator must likewise accept that pointer
-///   flowing back into its deallocation and reallocation methods.
+/// Unlike [`GlobalAlloc`], zero-sized allocations are allowed in
+/// `AllocRef`. If an underlying allocator does not support this (like
+/// jemalloc) or return a null pointer (such as `libc::malloc`), this case
+/// must be caught. In this case [`Layout::dangling()`] can be used to
+/// create a dangling, but aligned `NonNull<u8>`.
 ///
 /// Some of the methods require that a layout *fit* a memory block.
 /// What it means for a layout to "fit" a memory block means (or
@@ -649,6 +640,9 @@ pub unsafe trait GlobalAlloc {
 ///  * if an allocator does not support overallocating, it is fine to
 ///    simply return `layout.size()` as the allocated size.
 ///
+/// [`GlobalAlloc`]: self::GlobalAlloc
+/// [`Layout::dangling()`]: self::Layout::dangling
+///
 /// # Safety
 ///
 /// The `AllocRef` trait is an `unsafe` trait for a number of reasons, and
@@ -669,14 +663,6 @@ pub unsafe trait GlobalAlloc {
 /// the future.
 #[unstable(feature = "allocator_api", issue = "32838")]
 pub unsafe trait AllocRef {
-    // (Note: some existing allocators have unspecified but well-defined
-    // behavior in response to a zero size allocation request ;
-    // e.g., in C, `malloc` of 0 will either return a null pointer or a
-    // unique pointer, but will not have arbitrary undefined
-    // behavior.
-    // However in jemalloc for example,
-    // `mallocx(0)` is documented as undefined behavior.)
-
     /// On success, returns a pointer meeting the size and alignment
     /// guarantees of `layout` and the actual size of the allocated block,
     /// which must be greater than or equal to `layout.size()`.
@@ -690,15 +676,6 @@ pub unsafe trait AllocRef {
     /// behavior, e.g., to ensure initialization to particular sets of
     /// bit patterns.)
     ///
-    /// # Safety
-    ///
-    /// This function is unsafe because undefined behavior can result
-    /// if the caller does not ensure that `layout` has non-zero size.
-    ///
-    /// (Extension subtraits might provide more specific bounds on
-    /// behavior, e.g., guarantee a sentinel address or a null pointer
-    /// in response to a zero-size allocation request.)
-    ///
     /// # Errors
     ///
     /// Returning `Err` indicates that either memory is exhausted or
@@ -716,7 +693,7 @@ pub unsafe trait AllocRef {
     /// rather than directly invoking `panic!` or similar.
     ///
     /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
-    unsafe fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr>;
+    fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr>;
 
     /// Deallocate the memory referenced by `ptr`.
     ///
@@ -738,10 +715,6 @@ pub unsafe trait AllocRef {
     /// Behaves like `alloc`, but also ensures that the contents
     /// are set to zero before being returned.
     ///
-    /// # Safety
-    ///
-    /// This function is unsafe for the same reasons that `alloc` is.
-    ///
     /// # Errors
     ///
     /// Returning `Err` indicates that either memory is exhausted or
@@ -753,17 +726,17 @@ pub unsafe trait AllocRef {
     /// rather than directly invoking `panic!` or similar.
     ///
     /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
-    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
+    fn alloc_zeroed(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
         let size = layout.size();
         let result = self.alloc(layout);
         if let Ok((p, _)) = result {
-            ptr::write_bytes(p.as_ptr(), 0, size);
+            unsafe { ptr::write_bytes(p.as_ptr(), 0, size) }
         }
         result
     }
 
     // == METHODS FOR MEMORY REUSE ==
-    // realloc. alloc_excess, realloc_excess
+    // realloc, realloc_zeroed, grow_in_place, grow_in_place_zeroed, shrink_in_place
 
     /// Returns a pointer suitable for holding data described by
     /// a new layout with `layout`’s alignment and a size given
@@ -793,8 +766,6 @@ pub unsafe trait AllocRef {
     /// * `layout` must *fit* the `ptr` (see above). (The `new_size`
     ///   argument need not fit it.)
     ///
-    /// * `new_size` must be greater than zero.
-    ///
     /// * `new_size`, when rounded up to the nearest multiple of `layout.align()`,
     ///   must not overflow (i.e., the rounded value must be less than `usize::MAX`).
     ///
@@ -1009,8 +980,7 @@ pub unsafe trait AllocRef {
     /// * `layout` must *fit* the `ptr` (see above); note the
     ///   `new_size` argument need not fit it,
     ///
-    /// * `new_size` must not be greater than `layout.size()`
-    ///   (and must be greater than zero),
+    /// * `new_size` must not be greater than `layout.size()`,
     ///
     /// # Errors
     ///
diff --git a/src/libstd/alloc.rs b/src/libstd/alloc.rs
index 2da18e06d99..25f3ddcbeba 100644
--- a/src/libstd/alloc.rs
+++ b/src/libstd/alloc.rs
@@ -133,24 +133,41 @@ pub use alloc_crate::alloc::*;
 #[derive(Debug, Default, Copy, Clone)]
 pub struct System;
 
-// The AllocRef impl just forwards to the GlobalAlloc impl, which is in `std::sys::*::alloc`.
+// 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]
-    unsafe fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
-        NonNull::new(GlobalAlloc::alloc(self, layout)).ok_or(AllocErr).map(|p| (p, layout.size()))
+    fn alloc(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
+        if layout.size() == 0 {
+            Ok((layout.dangling(), 0))
+        } else {
+            unsafe {
+                NonNull::new(GlobalAlloc::alloc(self, layout))
+                    .ok_or(AllocErr)
+                    .map(|p| (p, layout.size()))
+            }
+        }
     }
 
     #[inline]
-    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
-        NonNull::new(GlobalAlloc::alloc_zeroed(self, layout))
-            .ok_or(AllocErr)
-            .map(|p| (p, layout.size()))
+    fn alloc_zeroed(&mut self, layout: Layout) -> Result<(NonNull<u8>, usize), AllocErr> {
+        if layout.size() == 0 {
+            Ok((layout.dangling(), 0))
+        } else {
+            unsafe {
+                NonNull::new(GlobalAlloc::alloc_zeroed(self, layout))
+                    .ok_or(AllocErr)
+                    .map(|p| (p, layout.size()))
+            }
+        }
     }
 
     #[inline]
     unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
-        GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
+        if layout.size() != 0 {
+            GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
+        }
     }
 
     #[inline]
@@ -160,9 +177,17 @@ unsafe impl AllocRef for System {
         layout: Layout,
         new_size: usize,
     ) -> Result<(NonNull<u8>, usize), AllocErr> {
-        NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size))
-            .ok_or(AllocErr)
-            .map(|p| (p, new_size))
+        match (layout.size(), new_size) {
+            (0, 0) => Ok((layout.dangling(), 0)),
+            (0, _) => self.alloc(Layout::from_size_align_unchecked(new_size, layout.align())),
+            (_, 0) => {
+                self.dealloc(ptr, layout);
+                Ok((layout.dangling(), 0))
+            }
+            (_, _) => NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size))
+                .ok_or(AllocErr)
+                .map(|p| (p, new_size)),
+        }
     }
 }