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,85 @@
+use std::{alloc::Layout, ptr::NonNull};
+
+use scoped_tls::scoped_thread_local;
+
+use crate::Allocator;
+
+scoped_thread_local!(pub(crate) static ALLOC: Allocator);
+
+#[derive(Clone, Copy)]
+pub struct SwcAlloc {
+    pub(crate) is_custom: bool,
+}
+
+impl Default for SwcAlloc {
+    fn default() -> Self {
+        Self {
+            is_custom: ALLOC.is_set(),
+        }
+    }
+}
+
+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) -> T) -> T {
+        if self.is_custom {
+            ALLOC.with(|a| {
+                //
+                f(&&**a as &dyn allocator_api2::alloc::Allocator)
+            })
+        } else {
+            f(&allocator_api2::alloc::Global)
+        }
+    }
+}
+
+unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        self.with_allocator(|a| a.allocate(layout))
+    }
+
+    fn allocate_zeroed(
+        &self,
+        layout: Layout,
+    ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        self.with_allocator(|a| a.allocate_zeroed(layout))
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        self.with_allocator(|a| a.deallocate(ptr, layout))
+    }
+
+    unsafe fn grow(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        self.with_allocator(|a| a.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> {
+        self.with_allocator(|a| a.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> {
+        self.with_allocator(|a| a.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,69 @@
+use std::ops::{Deref, DerefMut};
+
+mod rkyv;
+mod serde;
+
+use crate::alloc::SwcAlloc;
+
+/// 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> 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::default(),
+        ))
+    }
+
+    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> 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
+    }
+}

crates/swc_allocator/src/boxed/rkyv.rs

@@ -0,0 +1,59 @@
+use std::{alloc, cmp};
+
+use rkyv::{
+    boxed::{ArchivedBox, BoxResolver},
+    Archive, ArchivePointee, ArchiveUnsized, Deserialize, DeserializeUnsized, Fallible, Serialize,
+    SerializeUnsized,
+};
+
+use super::Box;
+
+impl<T: ArchiveUnsized + ?Sized> Archive for Box<T> {
+    type Archived = ArchivedBox<T::Archived>;
+    type Resolver = BoxResolver<T::MetadataResolver>;
+
+    #[inline]
+    unsafe fn resolve(&self, pos: usize, resolver: Self::Resolver, out: *mut Self::Archived) {
+        ArchivedBox::resolve_from_ref(self.as_ref(), pos, resolver, out);
+    }
+}
+
+impl<T: SerializeUnsized<S> + ?Sized, S: Fallible + ?Sized> Serialize<S> for Box<T> {
+    #[inline]
+    fn serialize(&self, serializer: &mut S) -> Result<Self::Resolver, S::Error> {
+        ArchivedBox::serialize_from_ref(self.as_ref(), serializer)
+    }
+}
+
+impl<T, D> Deserialize<Box<T>, D> for ArchivedBox<T::Archived>
+where
+    T: ArchiveUnsized + ?Sized,
+    T::Archived: DeserializeUnsized<T, D>,
+    D: Fallible + ?Sized,
+{
+    #[inline]
+    fn deserialize(&self, deserializer: &mut D) -> Result<Box<T>, D::Error> {
+        unsafe {
+            let data_address = self
+                .get()
+                .deserialize_unsized(deserializer, |layout| alloc::alloc(layout))?;
+            let metadata = self.get().deserialize_metadata(deserializer)?;
+            let ptr = ptr_meta::from_raw_parts_mut(data_address, metadata);
+            Ok(Box::from_raw(ptr))
+        }
+    }
+}
+
+impl<T: ArchivePointee + PartialEq<U> + ?Sized, U: ?Sized> PartialEq<Box<U>> for ArchivedBox<T> {
+    #[inline]
+    fn eq(&self, other: &Box<U>) -> bool {
+        self.get().eq(other.as_ref())
+    }
+}
+
+impl<T: ArchivePointee + PartialOrd<U> + ?Sized, U: ?Sized> PartialOrd<Box<U>> for ArchivedBox<T> {
+    #[inline]
+    fn partial_cmp(&self, other: &Box<U>) -> Option<cmp::Ordering> {
+        self.get().partial_cmp(other.as_ref())
+    }
+}

crates/swc_allocator/src/boxed/serde.rs

@@ -0,0 +1,27 @@
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+use super::Box;
+
+impl<T> Serialize for Box<T>
+where
+    T: Serialize,
+{
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        self.0.serialize(serializer)
+    }
+}
+
+impl<'de, T> Deserialize<'de> for Box<T>
+where
+    T: Deserialize<'de>,
+{
+    fn deserialize<D>(deserializer: D) -> Result<Box<T>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        Ok(Box(allocator_api2::boxed::Box::deserialize(deserializer)?))
+    }
+}