Skip to content

Commit

Permalink
feat(allocator): Implement SwcAlloc (#9232)
Browse files Browse the repository at this point in the history
**Related issue:**

 - This is a part of #9230
  • Loading branch information
kdy1 authored Jul 14, 2024
1 parent c9ac23b commit e343eb6
Show file tree
Hide file tree
Showing 15 changed files with 1,259 additions and 113 deletions.
19 changes: 19 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"
Expand Down
16 changes: 15 additions & 1 deletion crates/swc_allocator/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
153 changes: 153 additions & 0 deletions crates/swc_allocator/src/alloc.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
Loading

0 comments on commit e343eb6

Please sign in to comment.