diff --git a/crates/swc_allocator/benches/bench.rs b/crates/swc_allocator/benches/bench.rs index 355e6f939240..8b1b10328466 100644 --- a/crates/swc_allocator/benches/bench.rs +++ b/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 = black_box(SwcBox::new(black_box(i))); - vec.push(item); - } - }); + for i in 0..times { + let item: SwcBox = 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 alloc = FastAlloc::default(); - let mut vec = SwcVec::new_in(alloc); + let mut vec = SwcVec::new_in(alloc); - for i in 0..times { - let item: SwcBox = black_box(SwcBox::new_in(black_box(i), alloc)); - vec.push(item); - } - }); + for i in 0..times { + let item: SwcBox = black_box(SwcBox::new_in(black_box(i), alloc)); + vec.push(item); + } }) } diff --git a/crates/swc_allocator/src/alloc.rs b/crates/swc_allocator/src/alloc.rs index caf934f661f1..94197fcc5fcd 100644 --- a/crates/swc_allocator/src/alloc.rs +++ b/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 } } } diff --git a/crates/swc_allocator/src/vec/mod.rs b/crates/swc_allocator/src/vec/mod.rs index dfcd2bdf9b98..7569af262f9a 100644 --- a/crates/swc_allocator/src/vec/mod.rs +++ b/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 Extend for Vec { self.0.extend(iter) } } + +impl io::Write for Vec { + fn write(&mut self, buf: &[u8]) -> io::Result { + 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 { + 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 + } +} diff --git a/crates/swc_allocator/tests/apis.rs b/crates/swc_allocator/tests/apis.rs index 377251cd50f2..bf35cea53d35 100644 --- a/crates/swc_allocator/tests/apis.rs +++ b/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 = - black_box(swc_allocator::boxed::Box::new(black_box(i))); - vec.push(item); - } - }); + for i in 0..1000 { + let item: swc_allocator::boxed::Box = + black_box(swc_allocator::boxed::Box::new(black_box(i))); + vec.push(item); + } } diff --git a/crates/swc_allocator/tests/escape.rs b/crates/swc_allocator/tests/escape.rs index a215d29e6bb2..f12ab5520b0e 100644 --- a/crates/swc_allocator/tests/escape.rs +++ b/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);