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
+ }
+}