fix(allocator): Fix allocator & add benchmark (#9234)

**Related issue:**

 - This PR is a part of https://github.com/swc-project/swc/pull/9230.

.github/workflows/CI.yml

@@ -330,6 +330,7 @@ jobs:
           tool: cargo-hack@0.5.29
 
       - name: Check compilation
+        if: matrix.settings.os == 'ubuntu-latest'
         run: |
           ./scripts/github/run-cargo-hack.sh ${{ matrix.settings.crate }}
 

Cargo.lock

@@ -3702,11 +3702,14 @@ version = "0.1.2"
 dependencies = [
  "allocator-api2",
  "bumpalo",
+ "codspeed-criterion-compat",
+ "criterion",
  "ptr_meta",
  "rkyv",
  "scoped-tls",
  "serde",
  "serde_derive",
+ "swc_malloc",
 ]
 
 [[package]]

crates/swc_allocator/Cargo.toml

@@ -24,3 +24,15 @@ rkyv = { workspace = true, optional = true }
 scoped-tls = { workspace = true }
 serde = { workspace = true, optional = true }
 serde_derive = { workspace = true, optional = true }
+
+
+[dev-dependencies]
+criterion = { workspace = true }
+
+codspeed-criterion-compat = { workspace = true }
+swc_malloc                = { version = "0.5.10", path = "../swc_malloc" }
+
+
+[[bench]]
+harness = false
+name    = "bench"

crates/swc_allocator/benches/bench.rs

@@ -0,0 +1,92 @@
+extern crate swc_malloc;
+
+use codspeed_criterion_compat::{black_box, criterion_group, criterion_main, Bencher, Criterion};
+use swc_allocator::{FastAlloc, MemorySpace};
+
+fn bench_alloc(c: &mut Criterion) {
+    fn direct_alloc_std(b: &mut Bencher, times: usize) {
+        b.iter(|| {
+            let mut buf = std::vec::Vec::new();
+            for i in 0..times {
+                let item: std::boxed::Box<usize> = black_box(std::boxed::Box::new(black_box(i)));
+                buf.push(item);
+            }
+        })
+    }
+
+    fn direct_alloc_no_scope(b: &mut Bencher, times: usize) {
+        b.iter(|| {
+            let mut vec = swc_allocator::vec::Vec::new();
+            for i in 0..times {
+                let item: swc_allocator::boxed::Box<usize> =
+                    black_box(swc_allocator::boxed::Box::new(black_box(i)));
+                vec.push(item);
+            }
+        })
+    }
+
+    fn fast_alloc_no_scope(b: &mut Bencher, times: usize) {
+        b.iter(|| {
+            let allocator = FastAlloc::default();
+
+            let mut vec = allocator.vec();
+            for i in 0..times {
+                let item: swc_allocator::boxed::Box<usize> =
+                    black_box(allocator.alloc(black_box(i)));
+                vec.push(item);
+            }
+        })
+    }
+
+    fn direct_alloc_scoped(b: &mut Bencher, times: usize) {
+        b.iter(|| {
+            let allocator = MemorySpace::default();
+
+            allocator.scope(|| {
+                let mut vec = swc_allocator::vec::Vec::new();
+
+                for i in 0..times {
+                    let item: swc_allocator::boxed::Box<usize> =
+                        black_box(swc_allocator::boxed::Box::new(black_box(i)));
+                    vec.push(item);
+                }
+            });
+        })
+    }
+
+    fn fast_alloc_scoped(b: &mut Bencher, times: usize) {
+        b.iter(|| {
+            MemorySpace::default().scope(|| {
+                let allocator = FastAlloc::default();
+
+                let mut vec = allocator.vec();
+
+                for i in 0..times {
+                    let item: swc_allocator::boxed::Box<usize> =
+                        black_box(allocator.alloc(black_box(i)));
+                    vec.push(item);
+                }
+            });
+        })
+    }
+
+    c.bench_function("common/allocator/alloc/std/1000000", |b| {
+        direct_alloc_std(b, 1000000)
+    });
+    c.bench_function("common/allocator/alloc/no-scope/1000000", |b| {
+        direct_alloc_no_scope(b, 1000000)
+    });
+    c.bench_function("common/allocator/alloc/scoped/1000000", |b| {
+        direct_alloc_scoped(b, 1000000)
+    });
+
+    c.bench_function("common/allocator/alloc/cached-no-scope/1000000", |b| {
+        fast_alloc_no_scope(b, 1000000)
+    });
+    c.bench_function("common/allocator/alloc/cached-scoped/1000000", |b| {
+        fast_alloc_scoped(b, 1000000)
+    });
+}
+
+criterion_group!(benches, bench_alloc);
+criterion_main!(benches);

crates/swc_allocator/src/alloc.rs

@@ -3,12 +3,30 @@ use std::{alloc::Layout, ptr::NonNull};
 use allocator_api2::alloc::Global;
 use scoped_tls::scoped_thread_local;
 
-use crate::Allocator;
+use crate::{FastAlloc, MemorySpace};
 
-scoped_thread_local!(pub(crate) static ALLOC: Allocator);
+scoped_thread_local!(pub(crate) static ALLOC: MemorySpace);
 
-#[derive(Debug, Clone, Copy, Default)]
-pub struct SwcAlloc;
+#[derive(Debug, Clone, Copy)]
+pub struct SwcAlloc {
+    pub(crate) is_arena_mode: bool,
+}
+
+impl Default for FastAlloc {
+    fn default() -> Self {
+        Self {
+            is_arena_mode: ALLOC.is_set(),
+        }
+    }
+}
+
+impl Default for SwcAlloc {
+    fn default() -> Self {
+        SwcAlloc {
+            is_arena_mode: ALLOC.is_set(),
+        }
+    }
+}
 
 impl SwcAlloc {
     /// `true` is passed to `f` if the box is allocated with a custom allocator.
@@ -16,7 +34,7 @@ impl SwcAlloc {
         &self,
         f: impl FnOnce(&dyn allocator_api2::alloc::Allocator, bool) -> T,
     ) -> T {
-        if ALLOC.is_set() {
+        if self.is_arena_mode {
             ALLOC.with(|a| {
                 //
                 f(&&**a as &dyn allocator_api2::alloc::Allocator, true)
@@ -27,21 +45,8 @@ impl SwcAlloc {
     }
 }
 
-/// Set the last bit to 1
 fn mark_ptr_as_arena_mode(ptr: NonNull<[u8]>) -> NonNull<[u8]> {
-    let (mut raw_ptr, metadata) = ptr_meta::PtrExt::to_raw_parts(ptr.as_ptr());
-
-    raw_ptr = (raw_ptr as usize | 1) as *mut ();
-
-    unsafe {
-        // Safety:
-        NonNull::new_unchecked(ptr_meta::from_raw_parts_mut(raw_ptr, metadata))
-    }
-}
-
-fn is_ptr_in_arena_mode(ptr: NonNull<u8>) -> bool {
-    let ptr = ptr.as_ptr() as usize;
-    ptr & 1 == 1
+    ptr
 }
 
 unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
@@ -73,7 +78,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
     }
 
     unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
-        if is_ptr_in_arena_mode(ptr) {
+        if self.is_arena_mode {
             debug_assert!(
                 ALLOC.is_set(),
                 "Deallocating a pointer allocated with arena mode with a non-arena mode allocator"
@@ -96,7 +101,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
         old_layout: Layout,
         new_layout: Layout,
     ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
-        if is_ptr_in_arena_mode(ptr) {
+        if self.is_arena_mode {
             debug_assert!(
                 ALLOC.is_set(),
                 "Growing a pointer allocated with arena mode with a non-arena mode allocator"
@@ -114,7 +119,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
         old_layout: Layout,
         new_layout: Layout,
     ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
-        if is_ptr_in_arena_mode(ptr) {
+        if self.is_arena_mode {
             debug_assert!(
                 ALLOC.is_set(),
                 "Growing a pointer allocated with arena mode with a non-arena mode allocator"
@@ -132,7 +137,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
         old_layout: Layout,
         new_layout: Layout,
     ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
-        if is_ptr_in_arena_mode(ptr) {
+        if self.is_arena_mode {
             debug_assert!(
                 ALLOC.is_set(),
                 "Shrinking a pointer allocated with arena mode with a non-arena mode allocator"

crates/swc_allocator/src/boxed/mod.rs

@@ -7,7 +7,7 @@ use std::{
     pin::Pin,
 };
 
-use crate::alloc::SwcAlloc;
+use crate::{alloc::SwcAlloc, FastAlloc};
 
 #[cfg(feature = "rkyv")]
 mod rkyv;
@@ -54,7 +54,10 @@ impl<T> Box<T> {
     /// See [`std::boxed::Box::new`].
     #[inline(always)]
     pub fn new(value: T) -> Self {
-        Self(allocator_api2::boxed::Box::new_in(value, SwcAlloc))
+        Self(allocator_api2::boxed::Box::new_in(
+            value,
+            SwcAlloc::default(),
+        ))
     }
 
     /// Moves the value out of the box.
@@ -106,7 +109,10 @@ impl<T: ?Sized> Box<T> {
     /// [memory layout]: self#memory-layout
     /// [`Layout`]: crate::Layout
     pub unsafe fn from_raw(raw: *mut T) -> Self {
-        Self(allocator_api2::boxed::Box::from_raw_in(raw, SwcAlloc))
+        Self(allocator_api2::boxed::Box::from_raw_in(
+            raw,
+            SwcAlloc::default(),
+        ))
     }
 
     /// Consumes the `Box`, returning a wrapped raw pointer.
@@ -621,3 +627,9 @@ where
         self.0.len()
     }
 }
+
+impl FastAlloc {
+    pub fn alloc<T>(self, t: T) -> Box<T> {
+        Box(allocator_api2::boxed::Box::new_in(t, self.swc_alloc()))
+    }
+}

crates/swc_allocator/src/lib.rs

@@ -4,6 +4,7 @@
 
 #![allow(clippy::needless_doctest_main)]
 
+use alloc::SwcAlloc;
 use std::ops::{Deref, DerefMut};
 
 use bumpalo::Bump;
@@ -14,12 +15,25 @@ mod alloc;
 pub mod boxed;
 pub mod vec;
 
+#[derive(Debug, Clone, Copy)]
+pub struct FastAlloc {
+    is_arena_mode: bool,
+}
+
+impl FastAlloc {
+    fn swc_alloc(self) -> SwcAlloc {
+        SwcAlloc {
+            is_arena_mode: self.is_arena_mode,
+        }
+    }
+}
+
 #[derive(Default)]
-pub struct Allocator {
+pub struct MemorySpace {
     alloc: Bump,
 }
 
-impl Allocator {
+impl MemorySpace {
     /// Invokes `f` in a scope where the allocations are done in this allocator.
     #[inline(always)]
     pub fn scope<F, R>(&self, f: F) -> R
@@ -30,13 +44,13 @@ impl Allocator {
     }
 }
 
-impl From<Bump> for Allocator {
+impl From<Bump> for MemorySpace {
     fn from(alloc: Bump) -> Self {
         Self { alloc }
     }
 }
 
-impl Deref for Allocator {
+impl Deref for MemorySpace {
     type Target = Bump;
 
     fn deref(&self) -> &Bump {
@@ -44,7 +58,7 @@ impl Deref for Allocator {
     }
 }
 
-impl DerefMut for Allocator {
+impl DerefMut for MemorySpace {
     fn deref_mut(&mut self) -> &mut Bump {
         &mut self.alloc
     }

crates/swc_allocator/src/vec/mod.rs

@@ -3,7 +3,7 @@ use std::ops::{Deref, DerefMut};
 #[cfg(feature = "rkyv")]
 mod rkyv;
 
-use crate::{alloc::SwcAlloc, boxed::Box};
+use crate::{alloc::SwcAlloc, boxed::Box, FastAlloc};
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[repr(transparent)]
@@ -20,7 +20,8 @@ impl<T> Vec<T> {
 
     pub fn with_capacity(capacity: usize) -> Self {
         Self(allocator_api2::vec::Vec::with_capacity_in(
-            capacity, SwcAlloc,
+            capacity,
+            SwcAlloc::default(),
         ))
     }
 
@@ -163,7 +164,10 @@ impl<T> Vec<T> {
     /// ```
     pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
         Self(allocator_api2::vec::Vec::from_raw_parts_in(
-            ptr, length, capacity, SwcAlloc,
+            ptr,
+            length,
+            capacity,
+            SwcAlloc::default(),
         ))
     }
 }
@@ -184,7 +188,7 @@ impl<T> DerefMut for Vec<T> {
 
 impl<T> Default for Vec<T> {
     fn default() -> Self {
-        Self(allocator_api2::vec::Vec::new_in(SwcAlloc))
+        Self(allocator_api2::vec::Vec::new_in(SwcAlloc::default()))
     }
 }
 
@@ -239,3 +243,16 @@ impl<T> Extend<T> for Vec<T> {
         self.0.extend(iter)
     }
 }
+
+impl FastAlloc {
+    pub fn vec<T>(self) -> Vec<T> {
+        Vec(allocator_api2::vec::Vec::new_in(self.swc_alloc()))
+    }
+
+    pub fn vec_with_capacity<T>(self, capacity: usize) -> Vec<T> {
+        Vec(allocator_api2::vec::Vec::with_capacity_in(
+            capacity,
+            self.swc_alloc(),
+        ))
+    }
+}

crates/swc_allocator/tests/apis.rs

@@ -1 +1,36 @@
+use criterion::black_box;
+use swc_allocator::MemorySpace;
 
+#[test]
+fn direct_alloc_std() {
+    let mut buf = std::vec::Vec::new();
+    for i in 0..1000 {
+        let item: std::boxed::Box<usize> = black_box(std::boxed::Box::new(black_box(i)));
+        buf.push(item);
+    }
+}
+
+#[test]
+fn direct_alloc_no_scope() {
+    let mut vec = swc_allocator::vec::Vec::new();
+    for i in 0..1000 {
+        let item: swc_allocator::boxed::Box<usize> =
+            black_box(swc_allocator::boxed::Box::new(black_box(i)));
+        vec.push(item);
+    }
+}
+
+#[test]
+fn direct_alloc_in_scope() {
+    let allocator = MemorySpace::default();
+
+    allocator.scope(|| {
+        let mut vec = swc_allocator::vec::Vec::new();
+
+        for i in 0..1000 {
+            let item: swc_allocator::boxed::Box<usize> =
+                black_box(swc_allocator::boxed::Box::new(black_box(i)));
+            vec.push(item);
+        }
+    });
+}