Skip to content

Commit

Permalink
Refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
@wenyuzhao
wenyuzhao committed May 23, 2024
1 parent be7dc54 commit ad9897b
Show file tree
Hide file tree
Showing 4 changed files with 273 additions and 252 deletions.
48 changes: 33 additions & 15 deletions mallockit/macros/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,24 +9,42 @@ pub fn plan(_attr: TokenStream, item: TokenStream) -> TokenStream {
#input

mod __mallockit_plan {
pub(super) static PLAN: ::mallockit::util::Lazy<super::#name> = ::mallockit::util::Lazy::new(|| <super::#name as ::mallockit::Plan>::new());
type Plan = super::#name;

::mallockit::export_malloc_api!(PLAN, super::super::#name);
}
pub(super) static PLAN: ::mallockit::util::Lazy<Plan> = ::mallockit::util::Lazy::new(|| <Plan as ::mallockit::Plan>::new());

::mallockit::export_rust_global_alloc_api!(PLAN, #name);
#[cfg(any(feature = "malloc", feature = "mallockit/malloc"))]
#[::mallockit::ctor]
unsafe fn ctor() {
::mallockit::util::sys::hooks::process_start(&*PLAN);
}

#[cfg(target_os = "macos")]
#[no_mangle]
pub extern "C" fn mallockit_initialize_macos_tls() -> *mut u8 {
<Plan as ::mallockit::Plan>::Mutator::current() as *mut ::mallockit::Mutator as _
}

impl ::mallockit::plan::Singleton for #name {
fn singleton() -> &'static Self {
unsafe { &__mallockit_plan::PLAN }
impl ::mallockit::plan::Singleton for super::#name {
fn singleton() -> &'static Self {
unsafe { &PLAN }
}
}

::mallockit::export_malloc_api!(PLAN, super::super::#name);
::mallockit::export_rust_global_alloc_api!(super::super::#name);
}

include!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../target/generated_tests.rs"
));
::mallockit::rust_allocator_tests!(Global);
pub use __mallockit_plan::__mallockit_rust_api::Global;

#[cfg(test)]
mod tests {
include!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../target/generated_tests.rs"
));
::mallockit::rust_allocator_tests!(crate::Global);
}
};
result.into()
}
Expand All @@ -42,12 +60,12 @@ pub fn mutator(_attr: TokenStream, item: TokenStream) -> TokenStream {
#[cfg(not(target_os = "macos"))]
mod __mallockit_mutator {
#[thread_local]
pub(super) static mut MUTATOR: mallockit::util::Lazy<super::#name, mallockit::util::Local> = mallockit::util::Lazy::new(|| <super::#name as mallockit::Mutator>::new());
pub(super) static mut MUTATOR: ::mallockit::util::Lazy<super::#name, ::mallockit::util::Local> = ::mallockit::util::Lazy::new(|| <super::#name as ::mallockit::Mutator>::new());
}

impl mallockit::mutator::TLS for #name {
impl ::mallockit::mutator::TLS for #name {
fn new() -> Self {
<Self as mallockit::Mutator>::new()
<Self as ::mallockit::Mutator>::new()
}

#[cfg(not(target_os = "macos"))]
Expand Down
187 changes: 84 additions & 103 deletions mallockit/src/util/malloc/malloc_api.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -224,128 +224,109 @@ impl<P: Plan> MallocAPI<P> {
#[doc(hidden)]
macro_rules! export_malloc_api {
($plan: expr, $plan_ty: ty) => {
pub mod __mallockit {
#[cfg(any(feature = "malloc", feature = "mallockit/malloc"))]
pub mod __mallockit_malloc_api {
use super::*;
use $crate::Plan;
type ConcretePlan = $plan_ty;
type Malloc = $crate::util::malloc::MallocAPI<ConcretePlan>;
static MALLOC_IMPL: Malloc =
$crate::util::malloc::MallocAPI::<ConcretePlan>::new(&$plan);

#[cfg(any(feature = "malloc", feature = "mallockit/malloc"))]
#[$crate::ctor]
unsafe fn ctor() {
$crate::util::sys::hooks::process_start(&*$plan);
$crate::libc::atexit($crate::util::sys::hooks::process_exit);
type Malloc = $crate::util::malloc::MallocAPI<$plan_ty>;
static MALLOC_IMPL: Malloc = $crate::util::malloc::MallocAPI::<$plan_ty>::new(&$plan);

#[$crate::interpose]
pub unsafe extern "C" fn malloc(size: usize) -> *mut u8 {
MALLOC_IMPL.alloc_or_enomem(size, Malloc::MIN_ALIGNMENT)
}

#[cfg(target_os = "macos")]
#[no_mangle]
pub extern "C" fn mallockit_initialize_macos_tls() -> *mut u8 {
MALLOC_IMPL.mutator() as *mut _ as _
#[$crate::interpose]
pub unsafe extern "C" fn malloc_size(ptr: *mut u8) -> usize {
MALLOC_IMPL.malloc_size(ptr.into())
}

#[cfg(any(feature = "malloc", feature = "mallockit/malloc"))]
pub mod __malloc_api {
use super::{Malloc, MALLOC_IMPL};

#[$crate::interpose]
pub unsafe extern "C" fn malloc(size: usize) -> *mut u8 {
MALLOC_IMPL.alloc_or_enomem(size, Malloc::MIN_ALIGNMENT)
}

#[cfg(target_os = "macos")]
#[$crate::interpose]
pub unsafe extern "C" fn malloc_size(ptr: *mut u8) -> usize {
MALLOC_IMPL.malloc_size(ptr.into())
}

// #[cfg(target_os = "macos")]
// #[$crate::interpose]
// pub unsafe fn malloc_good_size(ptr: *mut u8) -> usize {
// MALLOC_IMPL.malloc_size(ptr.into())
// }
// #[cfg(target_os = "macos")]
// #[$crate::interpose]
// pub unsafe fn malloc_good_size(ptr: *mut u8) -> usize {
// MALLOC_IMPL.malloc_size(ptr.into())
// }

#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn malloc_usable_size(ptr: *mut u8) -> usize {
MALLOC_IMPL.malloc_size(ptr.into())
}
#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn malloc_usable_size(ptr: *mut u8) -> usize {
MALLOC_IMPL.malloc_size(ptr.into())
}

#[$crate::interpose]
pub unsafe extern "C" fn free(ptr: *mut u8) {
MALLOC_IMPL.free(ptr)
}
#[$crate::interpose]
pub unsafe extern "C" fn free(ptr: *mut u8) {
MALLOC_IMPL.free(ptr)
}

#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn cfree(ptr: *mut u8) {
MALLOC_IMPL.free(ptr)
}
#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn cfree(ptr: *mut u8) {
MALLOC_IMPL.free(ptr)
}

#[$crate::interpose]
pub unsafe extern "C" fn calloc(count: usize, size: usize) -> *mut u8 {
let size = count * size;
let ptr = MALLOC_IMPL.alloc_or_enomem(size, Malloc::MIN_ALIGNMENT);
std::ptr::write_bytes(ptr, 0, size);
ptr
}
#[$crate::interpose]
pub unsafe extern "C" fn calloc(count: usize, size: usize) -> *mut u8 {
let size = count * size;
let ptr = MALLOC_IMPL.alloc_or_enomem(size, Malloc::MIN_ALIGNMENT);
std::ptr::write_bytes(ptr, 0, size);
ptr
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
#[$crate::interpose]
pub unsafe extern "C" fn valloc(size: usize) -> *mut u8 {
MALLOC_IMPL.alloc_or_enomem(size, Malloc::PAGE_SIZE)
}
#[cfg(any(target_os = "linux", target_os = "macos"))]
#[$crate::interpose]
pub unsafe extern "C" fn valloc(size: usize) -> *mut u8 {
MALLOC_IMPL.alloc_or_enomem(size, Malloc::PAGE_SIZE)
}

#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn pvalloc(size: usize) -> *mut u8 {
MALLOC_IMPL.alloc_or_enomem(size, Malloc::PAGE_SIZE)
}
#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn pvalloc(size: usize) -> *mut u8 {
MALLOC_IMPL.alloc_or_enomem(size, Malloc::PAGE_SIZE)
}

#[$crate::interpose]
pub unsafe extern "C" fn realloc(ptr: *mut u8, size: usize) -> *mut u8 {
MALLOC_IMPL.reallocate_or_enomem(
ptr,
size,
cfg!(any(target_os = "linux", target_os = "windows")),
false,
)
}
#[$crate::interpose]
pub unsafe extern "C" fn realloc(ptr: *mut u8, size: usize) -> *mut u8 {
MALLOC_IMPL.reallocate_or_enomem(
ptr,
size,
cfg!(any(target_os = "linux", target_os = "windows")),
false,
)
}

#[cfg(target_os = "macos")]
#[$crate::interpose]
pub unsafe extern "C" fn reallocf(ptr: *mut u8, size: usize) -> *mut u8 {
MALLOC_IMPL.reallocate_or_enomem(ptr, size, false, true)
}
#[cfg(target_os = "macos")]
#[$crate::interpose]
pub unsafe extern "C" fn reallocf(ptr: *mut u8, size: usize) -> *mut u8 {
MALLOC_IMPL.reallocate_or_enomem(ptr, size, false, true)
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
#[$crate::interpose]
pub unsafe extern "C" fn posix_memalign(
ptr: *mut *mut u8,
alignment: usize,
size: usize,
) -> i32 {
MALLOC_IMPL.posix_memalign(ptr, alignment, size)
}
#[cfg(any(target_os = "linux", target_os = "macos"))]
#[$crate::interpose]
pub unsafe extern "C" fn posix_memalign(
ptr: *mut *mut u8,
alignment: usize,
size: usize,
) -> i32 {
MALLOC_IMPL.posix_memalign(ptr, alignment, size)
}

#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn memalign(alignment: usize, size: usize) -> *mut u8 {
MALLOC_IMPL.memalign(alignment, size)
}
#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn memalign(alignment: usize, size: usize) -> *mut u8 {
MALLOC_IMPL.memalign(alignment, size)
}

#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn aligned_alloc(alignment: usize, size: usize) -> *mut u8 {
MALLOC_IMPL.aligned_alloc(size, alignment, true, false)
}
#[cfg(target_os = "linux")]
#[$crate::interpose]
pub unsafe extern "C" fn aligned_alloc(alignment: usize, size: usize) -> *mut u8 {
MALLOC_IMPL.aligned_alloc(size, alignment, true, false)
}

#[cfg(target_os = "windows")]
#[$crate::interpose]
pub unsafe extern "C" fn _aligned_malloc(size: usize, alignment: usize) -> *mut u8 {
MALLOC_IMPL.aligned_alloc(size, alignment, false, true)
}
#[cfg(target_os = "windows")]
#[$crate::interpose]
pub unsafe extern "C" fn _aligned_malloc(size: usize, alignment: usize) -> *mut u8 {
MALLOC_IMPL.aligned_alloc(size, alignment, false, true)
}
}
};
Expand Down
Loading

0 comments on commit ad9897b

Please sign in to comment.