refactor(allocator): Use RAII guard instead of scope (#9254)

**Related issue:**

- https://github.com/swc-project/swc/issues/9253

crates/swc_allocator/benches/bench.rs

@@ -39,30 +39,30 @@ fn bench_alloc(c: &mut Criterion) {
     fn direct_alloc_scoped(b: &mut Bencher, times: usize) {
         b.iter(|| {
             let allocator = Allocator::default();
+            let _guard = unsafe { allocator.guard() };
 
-            allocator.scope(|| {
-                let mut vec = SwcVec::new();
+            let mut vec = SwcVec::new();
 
-                for i in 0..times {
-                    let item: SwcBox<usize> = black_box(SwcBox::new(black_box(i)));
-                    vec.push(item);
-                }
-            });
+            for i in 0..times {
+                let item: SwcBox<usize> = black_box(SwcBox::new(black_box(i)));
+                vec.push(item);
+            }
         })
     }
 
     fn fast_alloc_scoped(b: &mut Bencher, times: usize) {
         b.iter(|| {
-            Allocator::default().scope(|| {
-                let alloc = FastAlloc::default();
+            let alloc = Allocator::default();
+            let _guard = unsafe { alloc.guard() };
 
-                let mut vec = SwcVec::new_in(alloc);
+            let alloc = FastAlloc::default();
 
-                for i in 0..times {
-                    let item: SwcBox<usize> = black_box(SwcBox::new_in(black_box(i), alloc));
-                    vec.push(item);
-                }
-            });
+            let mut vec = SwcVec::new_in(alloc);
+
+            for i in 0..times {
+                let item: SwcBox<usize> = black_box(SwcBox::new_in(black_box(i), alloc));
+                vec.push(item);
+            }
         })
     }
 

crates/swc_allocator/src/alloc.rs

@@ -21,27 +21,33 @@ pub struct Allocator {
     alloc: Bump,
 }
 
+pub struct AllocGuard {
+    orig: Option<&'static Allocator>,
+}
+
+impl Drop for AllocGuard {
+    fn drop(&mut self) {
+        ALLOC.set(self.orig.take());
+    }
+}
+
 impl Allocator {
-    /// Invokes `f` in a scope where the allocations are done in this allocator.
+    /// Creates a RAII guard that enables optimized allocation.
     ///
     /// # Safety
     ///
-    /// [Allocator] must be dropped after dropping all [crate::boxed::Box] and
-    /// [crate::vec::Vec] created in the scope.
-    #[inline(always)]
-    pub fn scope<'a, F, R>(&'a self, f: F) -> R
-    where
-        F: FnOnce() -> R,
-    {
+    /// [Allocator] must outlive [crate::boxed::Box] and [crate::vec::Vec]
+    /// created while the guard is active.
+    pub unsafe fn guard(&self) -> AllocGuard {
+        let orig = ALLOC.get();
+
         let s = unsafe {
             // Safery: We are using a scoped API
-            transmute::<&'a Allocator, &'static Allocator>(self)
+            transmute::<&Allocator, &'static Allocator>(self)
         };
 
         ALLOC.set(Some(s));
-        let ret = f();
-        ALLOC.set(None);
-        ret
+        AllocGuard { orig }
     }
 }
 

crates/swc_allocator/src/vec/mod.rs

@@ -1,6 +1,9 @@
 //! Faster vec type.
 
-use std::ops::{Deref, DerefMut};
+use std::{
+    fmt, io,
+    ops::{Deref, DerefMut},
+};
 
 #[cfg(feature = "rkyv")]
 mod rkyv;
@@ -349,3 +352,32 @@ impl<T> Extend<T> for Vec<T> {
         self.0.extend(iter)
     }
 }
+
+impl io::Write for Vec<u8> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        io::Write::write(&mut self.0, buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        io::Write::flush(&mut self.0)
+    }
+
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        io::Write::write_all(&mut self.0, buf)
+    }
+
+    fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
+        io::Write::write_vectored(&mut self.0, bufs)
+    }
+
+    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
+        io::Write::write_fmt(&mut self.0, fmt)
+    }
+
+    fn by_ref(&mut self) -> &mut Self
+    where
+        Self: Sized,
+    {
+        self
+    }
+}

crates/swc_allocator/tests/apis.rs

@@ -23,14 +23,13 @@ fn direct_alloc_no_scope() {
 #[test]
 fn direct_alloc_in_scope() {
     let allocator = Allocator::default();
+    let _guard = unsafe { allocator.guard() };
 
-    allocator.scope(|| {
-        let mut vec = swc_allocator::vec::Vec::new();
+    let mut vec = swc_allocator::vec::Vec::new();
 
-        for i in 0..1000 {
-            let item: swc_allocator::boxed::Box<usize> =
-                black_box(swc_allocator::boxed::Box::new(black_box(i)));
-            vec.push(item);
-        }
-    });
+    for i in 0..1000 {
+        let item: swc_allocator::boxed::Box<usize> =
+            black_box(swc_allocator::boxed::Box::new(black_box(i)));
+        vec.push(item);
+    }
 }

crates/swc_allocator/tests/escape.rs

@@ -4,7 +4,10 @@ use swc_allocator::{boxed::Box, Allocator, FastAlloc};
 fn escape() {
     let allocator = Allocator::default();
 
-    let obj = allocator.scope(|| Box::new(1234));
+    let obj = {
+        let _guard = unsafe { allocator.guard() };
+        Box::new(1234)
+    };
 
     assert_eq!(*obj, 1234);
     // It should not segfault, because the allocator is still alive.
@@ -14,8 +17,10 @@ fn escape() {
 #[test]
 fn global_allocator() {
     let allocator = Allocator::default();
-
-    let obj = allocator.scope(|| Box::new_in(1234, FastAlloc::global()));
+    let obj = {
+        let _guard = unsafe { allocator.guard() };
+        Box::new_in(1234, FastAlloc::global())
+    };
 
     assert_eq!(*obj, 1234);
     drop(allocator);