refactor(allocator): Use & instead of a thread-local (#9235)

**Description:**

This is a part of #9230. I
profiled the performance, and `thread_local` took too long to get the
address of the thread-local variable. So, I inlined the reference into
the allocator.

# Benchmark result

```

Gnuplot not found, using plotters backend
common/allocator/alloc/std/1000000
                        time:   [4.9478 ms 4.9653 ms 4.9922 ms]
Found 17 outliers among 100 measurements (17.00%)
  4 (4.00%) high mild
  13 (13.00%) high severe

common/allocator/alloc/no-scope/1000000
                        time:   [5.4821 ms 5.4938 ms 5.5068 ms]
Found 17 outliers among 100 measurements (17.00%)
  2 (2.00%) high mild
  15 (15.00%) high severe

common/allocator/alloc/scoped/1000000
                        time:   [3.1401 ms 3.1456 ms 3.1518 ms]
Found 12 outliers among 100 measurements (12.00%)
  3 (3.00%) high mild
  9 (9.00%) high severe

common/allocator/alloc/cached-no-scope/1000000
                        time:   [5.0992 ms 5.1090 ms 5.1198 ms]
Found 11 outliers among 100 measurements (11.00%)
  2 (2.00%) high mild
  9 (9.00%) high severe

common/allocator/alloc/cached-scoped/1000000
                        time:   [3.0191 ms 3.0230 ms 3.0273 ms]
Found 11 outliers among 100 measurements (11.00%)
  2 (2.00%) low mild
  1 (1.00%) high mild
  8 (8.00%) high severe
```

crates/swc_allocator/Cargo.toml

@@ -24,6 +24,7 @@ rkyv = { workspace = true, optional = true }
 scoped-tls = { workspace = true }
 serde = { workspace = true, optional = true }
 serde_derive = { workspace = true, optional = true }
+triomphe = "0.1.13"
 
 
 [dev-dependencies]

crates/swc_allocator/benches/bench.rs

@@ -1,7 +1,7 @@
 extern crate swc_malloc;
 
 use codspeed_criterion_compat::{black_box, criterion_group, criterion_main, Bencher, Criterion};
-use swc_allocator::{FastAlloc, MemorySpace};
+use swc_allocator::{FastAlloc, SwcAllocator};
 
 fn bench_alloc(c: &mut Criterion) {
  fn direct_alloc_std(b: &mut Bencher, times: usize) {
@@ -40,7 +40,7 @@ fn bench_alloc(c: &mut Criterion) {
 
  fn direct_alloc_scoped(b: &mut Bencher, times: usize) {
  b.iter(|| {
- let allocator = MemorySpace::default();
+ let allocator = SwcAllocator::default();
 
  allocator.scope(|| {
  let mut vec = swc_allocator::vec::Vec::new();
@@ -56,7 +56,7 @@ fn bench_alloc(c: &mut Criterion) {
 
  fn fast_alloc_scoped(b: &mut Bencher, times: usize) {
  b.iter(|| {
- MemorySpace::default().scope(|| {
+ SwcAllocator::default().scope(|| {
  let allocator = FastAlloc::default();
 
  let mut vec = allocator.vec();

crates/swc_allocator/src/alloc.rs

@@ -1,44 +1,51 @@
-use std::{alloc::Layout, ptr::NonNull};
+use std::{alloc::Layout, mem::transmute, ptr::NonNull};
 
 use allocator_api2::alloc::Global;
 use scoped_tls::scoped_thread_local;
 
 use crate::{FastAlloc, MemorySpace};
 
-scoped_thread_local!(pub(crate) static ALLOC: MemorySpace);
+scoped_thread_local!(pub(crate) static ALLOC: &'static SwcAllocator);
 
-#[derive(Debug, Clone, Copy)]
-pub struct SwcAlloc {
- pub(crate) is_arena_mode: bool,
-}
+#[derive(Default)]
+pub struct SwcAllocator(MemorySpace);
 
-impl Default for FastAlloc {
- fn default() -> Self {
- Self {
- is_arena_mode: ALLOC.is_set(),
- }
+impl SwcAllocator {
+ /// Invokes `f` in a scope where the allocations are done in this allocator.
+ #[inline(always)]
+ pub fn scope<'a, F, R>(&'a self, f: F) -> R
+ where
+ F: FnOnce() -> R,
+ {
+ let s = unsafe {
+ // Safery: We are using a scoped API
+ transmute::<&'a SwcAllocator, &'static SwcAllocator>(self)
+ };
+
+ ALLOC.set(&s, f)
  }
 }
 
-impl Default for SwcAlloc {
+impl Default for FastAlloc {
  fn default() -> Self {
- SwcAlloc {
- is_arena_mode: ALLOC.is_set(),
+ Self {
+ alloc: if ALLOC.is_set() {
+ Some(ALLOC.with(|v| *v))
+ } else {
+ None
+ },
  }
  }
 }
 
-impl SwcAlloc {
+impl FastAlloc {
  /// `true` is passed to `f` if the box is allocated with a custom allocator.
  fn with_allocator<T>(
  &self,
  f: impl FnOnce(&dyn allocator_api2::alloc::Allocator, bool) -> T,
  ) -> T {
- if self.is_arena_mode {
- ALLOC.with(|a| {
- //
- f(&&**a as &dyn allocator_api2::alloc::Allocator, true)
- })
+ if let Some(arena) = &self.alloc {
+ f((&&*arena.0) as &dyn allocator_api2::alloc::Allocator, true)
  } else {
  f(&allocator_api2::alloc::Global, false)
  }
@@ -49,7 +56,7 @@ fn mark_ptr_as_arena_mode(ptr: NonNull<[u8]>) -> NonNull<[u8]> {
  ptr
 }
 
-unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
+unsafe impl allocator_api2::alloc::Allocator for FastAlloc {
  fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
  self.with_allocator(|a, is_arena_mode| {
  let ptr = a.allocate(layout)?;
@@ -78,18 +85,13 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
  }
 
  unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
- if self.is_arena_mode {
+ if self.alloc.is_some() {
  debug_assert!(
  ALLOC.is_set(),
  "Deallocating a pointer allocated with arena mode with a non-arena mode allocator"
  );
 
- ALLOC.with(|alloc| {
- unsafe {
- // Safety: We are in unsafe fn
- (&**alloc).deallocate(ptr, layout)
- }
- })
+ self.with_allocator(|alloc, _| alloc.deallocate(ptr, layout))
  } else {
  Global.deallocate(ptr, layout)
  }
@@ -101,16 +103,15 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
  old_layout: Layout,
  new_layout: Layout,
  ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
- if self.is_arena_mode {
- debug_assert!(
- ALLOC.is_set(),
- "Growing a pointer allocated with arena mode with a non-arena mode allocator"
- );
+ self.with_allocator(|alloc, is_arena_mode| {
+ let ptr = alloc.grow(ptr, old_layout, new_layout)?;
 
- ALLOC.with(|alloc| (&**alloc).grow(ptr, old_layout, new_layout))
- } else {
- Global.grow(ptr, old_layout, new_layout)
- }
+ if is_arena_mode {
+ Ok(mark_ptr_as_arena_mode(ptr))
+ } else {
+ Ok(ptr)
+ }
+ })
  }
 
  unsafe fn grow_zeroed(
@@ -119,16 +120,15 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
  old_layout: Layout,
  new_layout: Layout,
  ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
- if self.is_arena_mode {
- debug_assert!(
- ALLOC.is_set(),
- "Growing a pointer allocated with arena mode with a non-arena mode allocator"
- );
+ self.with_allocator(|alloc, is_arena_mode| {
+ let ptr = alloc.grow_zeroed(ptr, old_layout, new_layout)?;
 
- ALLOC.with(|alloc| (&**alloc).grow_zeroed(ptr, old_layout, new_layout))
- } else {
- Global.grow_zeroed(ptr, old_layout, new_layout)
- }
+ if is_arena_mode {
+ Ok(mark_ptr_as_arena_mode(ptr))
+ } else {
+ Ok(ptr)
+ }
+ })
  }
 
  unsafe fn shrink(
@@ -137,16 +137,15 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
  old_layout: Layout,
  new_layout: Layout,
  ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
- if self.is_arena_mode {
- debug_assert!(
- ALLOC.is_set(),
- "Shrinking a pointer allocated with arena mode with a non-arena mode allocator"
- );
+ self.with_allocator(|alloc, is_arena_mode| {
+ let ptr = alloc.shrink(ptr, old_layout, new_layout)?;
 
- ALLOC.with(|alloc| (&**alloc).shrink(ptr, old_layout, new_layout))
- } else {
- Global.shrink(ptr, old_layout, new_layout)
- }
+ if is_arena_mode {
+ Ok(mark_ptr_as_arena_mode(ptr))
+ } else {
+ Ok(ptr)
+ }
+ })
  }
 
  fn by_ref(&self) -> &Self

crates/swc_allocator/src/boxed/mod.rs

@@ -7,7 +7,7 @@ use std::{
  pin::Pin,
 };
 
-use crate::{alloc::SwcAlloc, FastAlloc};
+use crate::FastAlloc;
 
 #[cfg(feature = "rkyv")]
 mod rkyv;
@@ -23,7 +23,7 @@ mod serde;
 /// The last bit is 1 if the box is allocated with a custom allocator.
 #[repr(transparent)]
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub struct Box<T: ?Sized>(pub(crate) allocator_api2::boxed::Box<T, SwcAlloc>);
+pub struct Box<T: ?Sized>(pub(crate) allocator_api2::boxed::Box<T, FastAlloc>);
 
 impl<T> From<T> for Box<T> {
  #[inline(always)]
@@ -32,9 +32,9 @@ impl<T> From<T> for Box<T> {
  }
 }
 
-impl<T: ?Sized> From<allocator_api2::boxed::Box<T, SwcAlloc>> for Box<T> {
+impl<T: ?Sized> From<allocator_api2::boxed::Box<T, FastAlloc>> for Box<T> {
  #[inline(always)]
- fn from(v: allocator_api2::boxed::Box<T, SwcAlloc>) -> Self {
+ fn from(v: allocator_api2::boxed::Box<T, FastAlloc>) -> Self {
  Box(v)
  }
 }
@@ -56,7 +56,7 @@ impl<T> Box<T> {
  pub fn new(value: T) -> Self {
  Self(allocator_api2::boxed::Box::new_in(
  value,
- SwcAlloc::default(),
+ FastAlloc::default(),
  ))
  }
 
@@ -111,7 +111,7 @@ impl<T: ?Sized> Box<T> {
  pub unsafe fn from_raw(raw: *mut T) -> Self {
  Self(allocator_api2::boxed::Box::from_raw_in(
  raw,
- SwcAlloc::default(),
+ FastAlloc::default(),
  ))
  }
 
@@ -629,7 +629,7 @@ where
 }
 
 impl FastAlloc {
- pub fn alloc<T>(self, t: T) -> Box<T> {
- Box(allocator_api2::boxed::Box::new_in(t, self.swc_alloc()))
+ pub fn alloc<T>(&self, t: T) -> Box<T> {
+ Box(allocator_api2::boxed::Box::new_in(t, self.clone()))
  }
 }

crates/swc_allocator/src/lib.rs

@@ -4,46 +4,26 @@
 
 #![allow(clippy::needless_doctest_main)]
 
-use alloc::SwcAlloc;
 use std::ops::{Deref, DerefMut};
 
 use bumpalo::Bump;
 
-use crate::alloc::ALLOC;
+pub use crate::alloc::SwcAllocator;
 
 mod alloc;
 pub mod boxed;
 pub mod vec;
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Clone)]
 pub struct FastAlloc {
- is_arena_mode: bool,
-}
-
-impl FastAlloc {
- fn swc_alloc(self) -> SwcAlloc {
- SwcAlloc {
- is_arena_mode: self.is_arena_mode,
- }
- }
+ alloc: Option<&'static SwcAllocator>,
 }
 
 #[derive(Default)]
-pub struct MemorySpace {
+struct MemorySpace {
  alloc: Bump,
 }
 
-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
- where
- F: FnOnce() -> R,
- {
- ALLOC.set(self, f)
- }
-}
-
 impl From<Bump> for MemorySpace {
  fn from(alloc: Bump) -> Self {
  Self { alloc }