From 8547e02456bd335f7ac4276e4195b034c95b96bf Mon Sep 17 00:00:00 2001
From: Boshen <1430279+Boshen@users.noreply.github.com>
Date: Fri, 20 Dec 2024 13:26:48 +0000
Subject: [PATCH] feat(ast): implement `allocator_api2` for `Allocator` (#8043)

---
 Cargo.lock                                 |  1 -
 crates/oxc_allocator/src/allocator_api2.rs | 47 ++++++++++++++++++++++
 crates/oxc_allocator/src/lib.rs            |  1 +
 crates/oxc_semantic/Cargo.toml             |  1 -
 crates/oxc_semantic/src/scope.rs           |  8 ++--
 5 files changed, 52 insertions(+), 6 deletions(-)
 create mode 100644 crates/oxc_allocator/src/allocator_api2.rs

diff --git a/Cargo.lock b/Cargo.lock
index f9707090841d0..b0446737852f1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1926,7 +1926,6 @@ name = "oxc_semantic"
 version = "0.42.0"
 dependencies = [
  "assert-unchecked",
- "bumpalo",
  "hashbrown 0.15.2",
  "insta",
  "itertools",
diff --git a/crates/oxc_allocator/src/allocator_api2.rs b/crates/oxc_allocator/src/allocator_api2.rs
new file mode 100644
index 0000000000000..cff3271ba0884
--- /dev/null
+++ b/crates/oxc_allocator/src/allocator_api2.rs
@@ -0,0 +1,47 @@
+use std::{alloc::Layout, ptr::NonNull};
+
+use allocator_api2::alloc::{AllocError, Allocator};
+
+/// SAFETY:
+/// <https://github.com/fitzgen/bumpalo/blob/4eeab8847c85d5cde135ca21ae14a54e56b05224/src/lib.rs#L1938>
+unsafe impl Allocator for &crate::Allocator {
+    #[inline]
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        (&self.bump).allocate(layout)
+    }
+
+    #[inline]
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        (&self.bump).deallocate(ptr, layout);
+    }
+
+    #[inline]
+    unsafe fn shrink(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        (&self.bump).shrink(ptr, old_layout, new_layout)
+    }
+
+    #[inline]
+    unsafe fn grow(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        (&self.bump).grow(ptr, old_layout, new_layout)
+    }
+
+    #[inline]
+    unsafe fn grow_zeroed(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        (&self.bump).grow_zeroed(ptr, old_layout, new_layout)
+    }
+}
diff --git a/crates/oxc_allocator/src/lib.rs b/crates/oxc_allocator/src/lib.rs
index fbfc15a2e8f2b..a6f5646f74139 100644
--- a/crates/oxc_allocator/src/lib.rs
+++ b/crates/oxc_allocator/src/lib.rs
@@ -48,6 +48,7 @@ pub use bumpalo::collections::String;
 use bumpalo::Bump;
 
 mod address;
+mod allocator_api2;
 mod boxed;
 mod clone_in;
 mod convert;
diff --git a/crates/oxc_semantic/Cargo.toml b/crates/oxc_semantic/Cargo.toml
index 98f06e30f2035..9158d130690e8 100644
--- a/crates/oxc_semantic/Cargo.toml
+++ b/crates/oxc_semantic/Cargo.toml
@@ -30,7 +30,6 @@ oxc_span = { workspace = true }
 oxc_syntax = { workspace = true }
 
 assert-unchecked = { workspace = true }
-bumpalo = { workspace = true, features = ["allocator-api2"] }
 hashbrown = { workspace = true, features = ["allocator-api2"] }
 itertools = { workspace = true }
 phf = { workspace = true, features = ["macros"] }
diff --git a/crates/oxc_semantic/src/scope.rs b/crates/oxc_semantic/src/scope.rs
index cb89f89280ace..a62c8e6e6b4be 100644
--- a/crates/oxc_semantic/src/scope.rs
+++ b/crates/oxc_semantic/src/scope.rs
@@ -2,7 +2,7 @@ use std::mem;
 
 use rustc_hash::{FxBuildHasher, FxHashMap};
 
-use bumpalo::Bump;
+use oxc_allocator::Allocator;
 use oxc_index::IndexVec;
 use oxc_span::CompactStr;
 use oxc_syntax::{
@@ -12,7 +12,7 @@ use oxc_syntax::{
     symbol::SymbolId,
 };
 
-pub(crate) type Bindings<'a> = hashbrown::HashMap<&'a str, SymbolId, FxBuildHasher, &'a Bump>;
+pub(crate) type Bindings<'a> = hashbrown::HashMap<&'a str, SymbolId, FxBuildHasher, &'a Allocator>;
 pub type UnresolvedReferences = FxHashMap<CompactStr, Vec<ReferenceId>>;
 
 /// Scope Tree
@@ -54,7 +54,7 @@ impl Default for ScopeTree {
             node_ids: IndexVec::new(),
             flags: IndexVec::new(),
             root_unresolved_references: UnresolvedReferences::default(),
-            cell: ScopeTreeCell::new(Bump::new(), |_bump| ScopeTreeInner {
+            cell: ScopeTreeCell::new(Allocator::default(), |_bump| ScopeTreeInner {
                 bindings: IndexVec::new(),
             }),
         }
@@ -63,7 +63,7 @@ impl Default for ScopeTree {
 
 self_cell::self_cell!(
     pub(crate) struct ScopeTreeCell {
-        owner: Bump,
+        owner: Allocator,
         #[covariant]
         dependent: ScopeTreeInner,
     }