Reduce duplicate code across implementations of TurboMalloc

The two implementations were identical, aside from the allocator used. Move the compile-time branch into an inlined function.
vercel · Jun 12, 2024 · bfcbb2b · bfcbb2b
1 parent 98b7065
commit bfcbb2b
Showing 1 changed file with 17 additions and 42 deletions.
diff --git a/crates/turbo-tasks-malloc/src/lib.rs b/crates/turbo-tasks-malloc/src/lib.rs
@@ -67,70 +67,45 @@ impl TurboMalloc {
     }
 }
 
-#[cfg(all(
-    feature = "custom_allocator",
-    not(any(target_family = "wasm", target_env = "musl"))
-))]
-unsafe impl GlobalAlloc for TurboMalloc {
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let ret = mimalloc::MiMalloc.alloc(layout);
-        if !ret.is_null() {
-            add(layout.size());
-        }
-        ret
-    }
-
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        mimalloc::MiMalloc.dealloc(ptr, layout);
-        remove(layout.size());
-    }
-
-    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
-        let ret = mimalloc::MiMalloc.alloc_zeroed(layout);
-        if !ret.is_null() {
-            add(layout.size());
-        }
-        ret
-    }
-
-    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
-        let ret = mimalloc::MiMalloc.realloc(ptr, layout, new_size);
-        if !ret.is_null() {
-            let old_size = layout.size();
-            update(old_size, new_size);
-        }
-        ret
-    }
+/// Get the allocator for this platform that we should wrap with TurboMalloc.
+#[inline]
+fn base_alloc() -> &'static impl GlobalAlloc {
+    #[cfg(all(
+        feature = "custom_allocator",
+        not(any(target_family = "wasm", target_env = "musl"))
+    ))]
+    return &mimalloc::MiMalloc;
+    #[cfg(any(
+        not(feature = "custom_allocator"),
+        any(target_family = "wasm", target_env = "musl")
+    ))]
+    return &std::alloc::System;
 }
 
-#[cfg(any(
-    not(feature = "custom_allocator"),
-    any(target_family = "wasm", target_env = "musl")
-))]
 unsafe impl GlobalAlloc for TurboMalloc {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let ret = std::alloc::System.alloc(layout);
+        let ret = base_alloc().alloc(layout);
         if !ret.is_null() {
             add(layout.size());
         }
         ret
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        std::alloc::System.dealloc(ptr, layout);
+        base_alloc().dealloc(ptr, layout);
         remove(layout.size());
     }
 
     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
-        let ret = std::alloc::System.alloc_zeroed(layout);
+        let ret = base_alloc().alloc_zeroed(layout);
         if !ret.is_null() {
             add(layout.size());
         }
         ret
     }
 
     unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
-        let ret = std::alloc::System.realloc(ptr, layout, new_size);
+        let ret = base_alloc().realloc(ptr, layout, new_size);
         if !ret.is_null() {
             let old_size = layout.size();
             update(old_size, new_size);