Optimize process_heap_alloc

This commit is contained in:
John Kåre Alsaker 2024-03-11 05:43:12 +01:00
parent 3b1717c052
commit 50760aa2b5

View File

@ -6,6 +6,7 @@
use crate::sync::atomic::{AtomicPtr, Ordering}; use crate::sync::atomic::{AtomicPtr, Ordering};
use crate::sys::c; use crate::sys::c;
use crate::sys::common::alloc::{realloc_fallback, MIN_ALIGN}; use crate::sys::common::alloc::{realloc_fallback, MIN_ALIGN};
use core::mem::MaybeUninit;
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
@ -94,29 +95,30 @@ fn HeapReAlloc(
// a non-null handle returned by `GetProcessHeap`. // a non-null handle returned by `GetProcessHeap`.
#[inline] #[inline]
fn init_or_get_process_heap() -> c::HANDLE { fn init_or_get_process_heap() -> c::HANDLE {
let heap = HEAP.load(Ordering::Relaxed); // `HEAP` has not yet been successfully initialized
if core::intrinsics::unlikely(heap.is_null()) { let heap = unsafe { GetProcessHeap() };
// `HEAP` has not yet been successfully initialized if !heap.is_null() {
let heap = unsafe { GetProcessHeap() }; // SAFETY: No locking is needed because within the same process,
if !heap.is_null() { // successful calls to `GetProcessHeap` will always return the same value, even on different threads.
// SAFETY: No locking is needed because within the same process, HEAP.store(heap, Ordering::Release);
// successful calls to `GetProcessHeap` will always return the same value, even on different threads.
HEAP.store(heap, Ordering::Release);
// SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
heap
} else {
// Could not get the current process heap.
ptr::null_mut()
}
} else {
// SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap` // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
heap heap
} else {
// Could not get the current process heap.
ptr::null_mut()
} }
} }
/// This is outlined from `process_heap_alloc` so that `process_heap_alloc`
/// does not need any stack allocations.
#[inline(never)] #[inline(never)]
fn process_heap_alloc(flags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID { #[cold]
extern "C" fn process_heap_init_and_alloc(
_heap: MaybeUninit<c::HANDLE>, // We pass this argument to match the ABI of `HeapAlloc`
flags: c::DWORD,
dwBytes: c::SIZE_T,
) -> c::LPVOID {
let heap = init_or_get_process_heap(); let heap = init_or_get_process_heap();
if core::intrinsics::unlikely(heap.is_null()) { if core::intrinsics::unlikely(heap.is_null()) {
return ptr::null_mut(); return ptr::null_mut();
@ -125,6 +127,21 @@ fn process_heap_alloc(flags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID {
unsafe { HeapAlloc(heap, flags, dwBytes) } unsafe { HeapAlloc(heap, flags, dwBytes) }
} }
#[inline(never)]
fn process_heap_alloc(
_heap: MaybeUninit<c::HANDLE>, // We pass this argument to match the ABI of `HeapAlloc`,
flags: c::DWORD,
dwBytes: c::SIZE_T,
) -> c::LPVOID {
let heap = HEAP.load(Ordering::Relaxed);
if core::intrinsics::likely(!heap.is_null()) {
// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
unsafe { HeapAlloc(heap, flags, dwBytes) }
} else {
process_heap_init_and_alloc(MaybeUninit::uninit(), flags, dwBytes)
}
}
// Get a non-null handle to the default heap of the current process. // Get a non-null handle to the default heap of the current process.
// SAFETY: `HEAP` must have been successfully initialized. // SAFETY: `HEAP` must have been successfully initialized.
#[inline] #[inline]
@ -148,12 +165,12 @@ unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
if layout.align() <= MIN_ALIGN { if layout.align() <= MIN_ALIGN {
// The returned pointer points to the start of an allocated block. // The returned pointer points to the start of an allocated block.
process_heap_alloc(flags, layout.size()) as *mut u8 process_heap_alloc(MaybeUninit::uninit(), flags, layout.size()) as *mut u8
} else { } else {
// Allocate extra padding in order to be able to satisfy the alignment. // Allocate extra padding in order to be able to satisfy the alignment.
let total = layout.align() + layout.size(); let total = layout.align() + layout.size();
let ptr = process_heap_alloc(flags, total) as *mut u8; let ptr = process_heap_alloc(MaybeUninit::uninit(), flags, total) as *mut u8;
if ptr.is_null() { if ptr.is_null() {
// Allocation has failed. // Allocation has failed.
return ptr::null_mut(); return ptr::null_mut();