Implement the GlobalAlloc trait, upgrade to nightly-2018-04-14

CC rust-lang/rust#49669
SimonSapin · Apr 13, 2018 · a1c7d92 · a1c7d92
1 parent e9c2722
commit a1c7d92
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 131 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "jemallocator"
-version = "0.1.4"
+version = "0.1.5"
 authors = ["Alex Crichton <alex@alexcrichton.com>"]
 license = "MIT/Apache-2.0"
 readme = "README.md"

diff --git a/src/lib.rs b/src/lib.rs
@@ -15,15 +15,15 @@
 //! `Alloc` trait and is suitable both as a memory allocator and as a
 //! global allocator.
 
-#![feature(allocator_api)]
+#![feature(allocator_api, nonnull_cast)]
 #![deny(missing_docs)]
 
 extern crate jemalloc_sys;
 extern crate libc;
 
 use std::mem;
-use std::ptr;
-use std::heap::{Alloc, Layout, Excess, CannotReallocInPlace, AllocErr, System};
+use std::ptr::{self, NonNull};
+use std::heap::{GlobalAlloc, Alloc, Layout, Opaque, Excess, CannotReallocInPlace, AllocErr, System};
 
 use libc::{c_int, c_void};
 
@@ -45,16 +45,16 @@ const MIN_ALIGN: usize = 8;
               target_arch = "sparc64")))]
 const MIN_ALIGN: usize = 16;
 
-fn layout_to_flags(layout: &Layout) -> c_int {
+fn layout_to_flags(align: usize, size: usize) -> c_int {
     // If our alignment is less than the minimum alignment they we may not
     // have to pass special flags asking for a higher alignment. If the
     // alignment is greater than the size, however, then this hits a sort of odd
     // case where we still need to ask for a custom alignment. See #25 for more
     // info.
-    if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
+    if align <= MIN_ALIGN && align <= size {
         0
     } else {
-        ffi::MALLOCX_ALIGN(layout.align())
+        ffi::MALLOCX_ALIGN(align)
     }
 }
 
@@ -66,160 +66,106 @@ fn layout_to_flags(layout: &Layout) -> c_int {
 /// allocator.
 pub struct Jemalloc;
 
-unsafe impl Alloc for Jemalloc {
-    #[inline]
-    unsafe fn alloc(&mut self, layout: Layout) -> Result<*mut u8, AllocErr> {
-        (&*self).alloc(layout)
-    }
-
-    #[inline]
-    unsafe fn alloc_zeroed(&mut self, layout: Layout)
-        -> Result<*mut u8, AllocErr>
-    {
-        (&*self).alloc_zeroed(layout)
-    }
-
+unsafe impl GlobalAlloc for Jemalloc {
     #[inline]
-    unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
-        (&*self).dealloc(ptr, layout)
+    unsafe fn alloc(&self, layout: Layout) -> *mut Opaque {
+        let flags = layout_to_flags(layout.align(), layout.size());
+        let ptr = ffi::mallocx(layout.size(), flags);
+        ptr as *mut Opaque
     }
 
     #[inline]
-    unsafe fn realloc(&mut self,
-                      ptr: *mut u8,
-                      old_layout: Layout,
-                      new_layout: Layout) -> Result<*mut u8, AllocErr> {
-        (&*self).realloc(ptr, old_layout, new_layout)
-    }
-
-    fn oom(&mut self, err: AllocErr) -> ! {
-        (&*self).oom(err)
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut Opaque {
+        let ptr = if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
+            ffi::calloc(1, layout.size())
+        } else {
+            let flags = layout_to_flags(layout.align(), layout.size()) | ffi::MALLOCX_ZERO;
+            ffi::mallocx(layout.size(), flags)
+        };
+        ptr as *mut Opaque
     }
 
     #[inline]
-    fn usable_size(&self, layout: &Layout) -> (usize, usize) {
-        (&self).usable_size(layout)
+    unsafe fn dealloc(&self, ptr: *mut Opaque, layout: Layout) {
+        let flags = layout_to_flags(layout.align(), layout.size());
+        ffi::sdallocx(ptr as *mut c_void, layout.size(), flags)
     }
 
     #[inline]
-    unsafe fn alloc_excess(&mut self, layout: Layout) -> Result<Excess, AllocErr> {
-        (&*self).alloc_excess(layout)
+    unsafe fn realloc(&self,
+                      ptr: *mut Opaque,
+                      layout: Layout,
+                      new_size: usize) -> *mut Opaque {
+        let flags = layout_to_flags(layout.align(), new_size);
+        let ptr = ffi::rallocx(ptr as *mut c_void, new_size, flags);
+        ptr as *mut Opaque
     }
 
     #[inline]
-    unsafe fn realloc_excess(&mut self,
-                             ptr: *mut u8,
-                             layout: Layout,
-                             new_layout: Layout) -> Result<Excess, AllocErr> {
-        (&*self).realloc_excess(ptr, layout, new_layout)
+    fn oom(&self) -> ! {
+        System.oom()
     }
+}
 
+unsafe impl Alloc for Jemalloc {
     #[inline]
-    unsafe fn grow_in_place(&mut self,
-                            ptr: *mut u8,
-                            layout: Layout,
-                            new_layout: Layout) -> Result<(), CannotReallocInPlace> {
-        (&*self).grow_in_place(ptr, layout, new_layout)
+    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<Opaque>, AllocErr> {
+        NonNull::new(GlobalAlloc::alloc(self, layout)).ok_or(AllocErr)
     }
 
     #[inline]
-    unsafe fn shrink_in_place(&mut self,
-                              ptr: *mut u8,
-                              layout: Layout,
-                              new_layout: Layout) -> Result<(), CannotReallocInPlace> {
-        (&*self).shrink_in_place(ptr, layout, new_layout)
+    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<Opaque>, AllocErr> {
+        NonNull::new(GlobalAlloc::alloc_zeroed(self, layout)).ok_or(AllocErr)
     }
-}
 
-unsafe impl<'a> Alloc for &'a Jemalloc {
     #[inline]
-    unsafe fn alloc(&mut self, layout: Layout) -> Result<*mut u8, AllocErr> {
-        let flags = layout_to_flags(&layout);
-        let ptr = ffi::mallocx(layout.size(), flags);
-        if ptr.is_null() {
-            Err(AllocErr::Exhausted { request: layout })
-        } else {
-            Ok(ptr as *mut u8)
-        }
+    unsafe fn dealloc(&mut self, ptr: NonNull<Opaque>, layout: Layout) {
+        GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
     }
 
     #[inline]
-    unsafe fn alloc_zeroed(&mut self, layout: Layout)
-        -> Result<*mut u8, AllocErr>
-    {
-        let ptr = if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
-            ffi::calloc(1, layout.size())
-        } else {
-            let flags = layout_to_flags(&layout) | ffi::MALLOCX_ZERO;
-            ffi::mallocx(layout.size(), flags)
-        };
-        if ptr.is_null() {
-            Err(AllocErr::Exhausted { request: layout })
-        } else {
-            Ok(ptr as *mut u8)
-        }
+    unsafe fn realloc(&mut self,
+                      ptr: NonNull<Opaque>,
+                      layout: Layout,
+                      new_size: usize) -> Result<NonNull<Opaque>, AllocErr> {
+        NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size)).ok_or(AllocErr)
     }
 
     #[inline]
     unsafe fn alloc_excess(&mut self, layout: Layout) -> Result<Excess, AllocErr> {
-        let flags = layout_to_flags(&layout);
+        let flags = layout_to_flags(layout.align(), layout.size());
         let ptr = ffi::mallocx(layout.size(), flags);
-        if ptr.is_null() {
-            Err(AllocErr::Exhausted { request: layout })
-        } else {
+        if let Some(nonnull) = NonNull::new(ptr as *mut Opaque) {
             let excess = ffi::nallocx(layout.size(), flags);
-            Ok(Excess(ptr as *mut u8, excess))
-        }
-    }
-
-    #[inline]
-    unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
-        let flags = layout_to_flags(&layout);
-        ffi::sdallocx(ptr as *mut c_void, layout.size(), flags)
-    }
-
-    #[inline]
-    unsafe fn realloc(&mut self,
-                      ptr: *mut u8,
-                      old_layout: Layout,
-                      new_layout: Layout) -> Result<*mut u8, AllocErr> {
-        if old_layout.align() != new_layout.align() {
-            return Err(AllocErr::Unsupported { details: "cannot change align" })
-        }
-        let flags = layout_to_flags(&new_layout);
-        let ptr = ffi::rallocx(ptr as *mut c_void, new_layout.size(), flags);
-        if ptr.is_null() {
-            Err(AllocErr::Exhausted { request: new_layout })
+            Ok(Excess(nonnull, excess))
         } else {
-            Ok(ptr as *mut u8)
+            Err(AllocErr)
         }
     }
 
     #[inline]
     unsafe fn realloc_excess(&mut self,
-                      ptr: *mut u8,
-                      old_layout: Layout,
-                      new_layout: Layout) -> Result<Excess, AllocErr> {
-        if old_layout.align() != new_layout.align() {
-            return Err(AllocErr::Unsupported { details: "cannot change align" })
-        }
-        let flags = layout_to_flags(&new_layout);
-        let ptr = ffi::rallocx(ptr as *mut c_void, new_layout.size(), flags);
-        if ptr.is_null() {
-            Err(AllocErr::Exhausted { request: new_layout })
+                      ptr: NonNull<Opaque>,
+                      layout: Layout,
+                      new_size: usize) -> Result<Excess, AllocErr> {
+        let flags = layout_to_flags(layout.align(), new_size);
+        let ptr = ffi::rallocx(ptr.cast().as_ptr(), new_size, flags);
+        if let Some(nonnull) = NonNull::new(ptr as *mut Opaque) {
+            let excess = ffi::nallocx(new_size, flags);
+            Ok(Excess(nonnull, excess))
         } else {
-            let excess = ffi::nallocx(new_layout.size(), flags);
-            Ok(Excess(ptr as *mut u8, excess))
+            Err(AllocErr)
         }
     }
 
-    fn oom(&mut self, err: AllocErr) -> ! {
-        System.oom(err)
+    #[inline]
+    fn oom(&mut self) -> ! {
+        System.oom()
     }
 
     #[inline]
     fn usable_size(&self, layout: &Layout) -> (usize, usize) {
-        let flags = layout_to_flags(&layout);
+        let flags = layout_to_flags(layout.align(), layout.size());
         unsafe {
             let max = ffi::nallocx(layout.size(), flags);
             (layout.size(), max)
@@ -228,23 +174,20 @@ unsafe impl<'a> Alloc for &'a Jemalloc {
 
     #[inline]
     unsafe fn grow_in_place(&mut self,
-                            ptr: *mut u8,
-                            old_layout: Layout,
-                            new_layout: Layout) -> Result<(), CannotReallocInPlace> {
-        self.shrink_in_place(ptr, old_layout, new_layout)
+                            ptr: NonNull<Opaque>,
+                            layout: Layout,
+                            new_size: usize) -> Result<(), CannotReallocInPlace> {
+        self.shrink_in_place(ptr, layout, new_size)
     }
 
     #[inline]
     unsafe fn shrink_in_place(&mut self,
-                              ptr: *mut u8,
-                              old_layout: Layout,
-                              new_layout: Layout) -> Result<(), CannotReallocInPlace> {
-        if old_layout.align() != new_layout.align() {
-            return Err(CannotReallocInPlace)
-        }
-        let flags = layout_to_flags(&new_layout);
-        let size = ffi::xallocx(ptr as *mut c_void, new_layout.size(), 0, flags);
-        if size >= new_layout.size() {
+                              ptr: NonNull<Opaque>,
+                              layout: Layout,
+                              new_size: usize) -> Result<(), CannotReallocInPlace> {
+        let flags = layout_to_flags(layout.align(), new_size);
+        let size = ffi::xallocx(ptr.cast().as_ptr(), new_size, 0, flags);
+        if size >= new_size {
             Err(CannotReallocInPlace)
         } else {
             Ok(())

diff --git a/tests/malloctl.rs b/tests/malloctl.rs
@@ -13,7 +13,7 @@ static A: Jemalloc = Jemalloc;
 fn smoke() {
     let layout = Layout::from_size_align(100, 8).unwrap();
     unsafe {
-        let ptr = Jemalloc.alloc(layout.clone()).unwrap_or_else(|e| Jemalloc.oom(e));
+        let ptr = Jemalloc.alloc(layout.clone()).unwrap_or_else(|_| Jemalloc.oom());
         Jemalloc.dealloc(ptr, layout);
     }
 }

diff --git a/tests/smoke.rs b/tests/smoke.rs
@@ -25,7 +25,8 @@ fn overaligned() {
             Jemalloc.alloc(Layout::from_size_align(size, align).unwrap()).unwrap()
         }).collect();
         for &ptr in &pointers {
-            assert_eq!((ptr as usize) % align, 0, "Got a pointer less aligned than requested")
+            assert_eq!((ptr.as_ptr() as usize) % align, 0,
+                       "Got a pointer less aligned than requested")
         }
 
         // Clean up