feat(allocator): Introduce scoped allocator

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

@@ -9,4 +9,14 @@ repository    = { workspace = true }
 version       = "0.1.1"
 
 [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 }
+scoped-tls = { workspace = true }
+serde = { workspace = true }
+serde_derive = { workspace = 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 ptr = ptr.as_ptr() as *mut () as usize | 1;
+    unsafe { NonNull::new_unchecked(ptr as *mut [u8]) }
+}
+
+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| {
+                return (&**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| {
+                return (&**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| {
+                return (&**alloc).shrink(ptr, old_layout, new_layout);
+            })
+        } else {
+            Global.shrink(ptr, old_layout, new_layout)
+        }
+    }
+
+    fn by_ref(&self) -> &Self
+    where
+        Self: Sized,
+    {
+        self
+    }
+}

crates/swc_allocator/src/boxed/mod.rs

@@ -0,0 +1,93 @@
+use std::ops::{Deref, DerefMut};
+
+use crate::alloc::SwcAlloc;
+
+mod rkyv;
+mod serde;
+
+/// A special `Box` which has size of [`std::boxed::Box`] but **may** be
+/// allocated with a custom allocator.
+///
+///
+/// # Representation
+///
+/// The last bit is 1 if the box is allocated with a custom allocator.
+#[repr(transparent)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct Box<T: ?Sized>(pub(crate) allocator_api2::boxed::Box<T, SwcAlloc>);
+
+impl<T> From<T> for Box<T> {
+    #[inline(always)]
+    fn from(v: T) -> Self {
+        Box::new(v)
+    }
+}
+
+pub trait IntoBox<T>: Sized {
+    fn boxed(self) -> Box<T>;
+}
+
+impl<T, V> IntoBox<T> for V
+where
+    V: Into<T>,
+{
+    #[inline(always)]
+    fn boxed(self) -> Box<T> {
+        Box::new(self.into())
+    }
+}
+
+impl<T> Default for Box<T>
+where
+    T: Default,
+{
+    fn default() -> Self {
+        Box::new(Default::default())
+    }
+}
+
+impl<T> Box<T> {
+    #[inline(always)]
+    pub fn new(value: T) -> Self {
+        Self(allocator_api2::boxed::Box::new_in(value, SwcAlloc))
+    }
+
+    pub fn unbox(self) -> T {
+        allocator_api2::boxed::Box::into_inner(self.0)
+    }
+}
+
+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(),
+        ))
+    }
+}
+
+impl<T: ?Sized> AsRef<T> for Box<T> {
+    fn as_ref(&self) -> &T {
+        &self.0
+    }
+}
+
+impl<T: ?Sized> AsMut<T> for Box<T> {
+    fn as_mut(&mut self) -> &mut T {
+        &mut self.0
+    }
+}
+
+impl<T: ?Sized> Deref for Box<T> {
+    type Target = T;
+
+    fn deref(&self) -> &T {
+        &self.0
+    }
+}
+
+impl<T: ?Sized> DerefMut for Box<T> {
+    fn deref_mut(&mut self) -> &mut T {
+        &mut self.0
+    }
+}