feat(allocator): Implement SwcAlloc (#9232)

**Related issue:**

 - This is a part of #9230

Cargo.toml

@@ -40,6 +40,7 @@ resolver = "2"
 
   Inflector                 = "0.11.4"
   ahash                     = "0.8.8"
+  allocator-api2            = "0.2.18"
   ansi_term                 = "0.12.1"
   anyhow                    = "1.0.81"
   arbitrary                 = "1"
@@ -95,6 +96,7 @@ resolver = "2"
   phf                       = "0.11.2"
   pretty_assertions         = "1.3"
   proc-macro2               = "1.0.24"
+  ptr_meta                  = "0.1.4"
   quote                     = "1.0.7"
   rayon                     = "1.7.0"
   regex                     = "1.5.4"

crates/swc_allocator/Cargo.toml

@@ -8,5 +8,19 @@ name          = "swc_allocator"
 repository    = { workspace = true }
 version       = "0.1.1"
 
+[features]
+rkyv  = ["dep:rkyv"]
+serde = ["dep:serde", "dep:serde_derive"]
+
 [dependencies]
-bumpalo = { workspace = true, features = ["boxed", "collections"] }
+allocator-api2 = { workspace = true, features = ["serde"] }
+bumpalo = { workspace = true, features = [
+    "allocator-api2",
+    "boxed",
+    "collections",
+] }
+ptr_meta = { workspace = true }
+rkyv = { workspace = true, optional = true }
+scoped-tls = { workspace = true }
+serde = { workspace = true, optional = true }
+serde_derive = { workspace = true, optional = true }

crates/swc_allocator/src/alloc.rs

@@ -0,0 +1,153 @@
+use std::{alloc::Layout, ptr::NonNull};
+
+use allocator_api2::alloc::Global;
+use scoped_tls::scoped_thread_local;
+
+use crate::Allocator;
+
+scoped_thread_local!(pub(crate) static ALLOC: Allocator);
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct SwcAlloc;
+
+impl SwcAlloc {
+    /// `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 ALLOC.is_set() {
+            ALLOC.with(|a| {
+                //
+                f(&&**a as &dyn allocator_api2::alloc::Allocator, true)
+            })
+        } else {
+            f(&allocator_api2::alloc::Global, false)
+        }
+    }
+}
+
+/// 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
+}
+
+unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        self.with_allocator(|a, is_arena_mode| {
+            let ptr = a.allocate(layout)?;
+
+            if is_arena_mode {
+                Ok(mark_ptr_as_arena_mode(ptr))
+            } else {
+                Ok(ptr)
+            }
+        })
+    }
+
+    fn allocate_zeroed(
+        &self,
+        layout: Layout,
+    ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        self.with_allocator(|a, is_arena_mode| {
+            let ptr = a.allocate_zeroed(layout)?;
+
+            if is_arena_mode {
+                Ok(mark_ptr_as_arena_mode(ptr))
+            } else {
+                Ok(ptr)
+            }
+        })
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        if is_ptr_in_arena_mode(ptr) {
+            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)
+                }
+            })
+        } else {
+            Global.deallocate(ptr, layout)
+        }
+    }
+
+    unsafe fn grow(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        if is_ptr_in_arena_mode(ptr) {
+            debug_assert!(
+                ALLOC.is_set(),
+                "Growing a pointer allocated with arena mode with a non-arena mode allocator"
+            );
+
+            ALLOC.with(|alloc| (&**alloc).grow(ptr, old_layout, new_layout))
+        } else {
+            Global.grow(ptr, old_layout, new_layout)
+        }
+    }
+
+    unsafe fn grow_zeroed(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        if is_ptr_in_arena_mode(ptr) {
+            debug_assert!(
+                ALLOC.is_set(),
+                "Growing a pointer allocated with arena mode with a non-arena mode allocator"
+            );
+
+            ALLOC.with(|alloc| (&**alloc).grow_zeroed(ptr, old_layout, new_layout))
+        } else {
+            Global.grow_zeroed(ptr, old_layout, new_layout)
+        }
+    }
+
+    unsafe fn shrink(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        if is_ptr_in_arena_mode(ptr) {
+            debug_assert!(
+                ALLOC.is_set(),
+                "Shrinking a pointer allocated with arena mode with a non-arena mode allocator"
+            );
+
+            ALLOC.with(|alloc| (&**alloc).shrink(ptr, old_layout, new_layout))
+        } else {
+            Global.shrink(ptr, old_layout, new_layout)
+        }
+    }
+
+    fn by_ref(&self) -> &Self
+    where
+        Self: Sized,
+    {
+        self
+    }
+}