Wrap HeapAlloc and never inline

This commit is contained in:
王宇逸 2024-01-24 02:43:06 +08:00
parent d9d89fd53d
commit 27a6e6e67b

View File

@ -95,7 +95,7 @@ fn HeapReAlloc(
#[inline] #[inline]
fn init_or_get_process_heap() -> c::HANDLE { fn init_or_get_process_heap() -> c::HANDLE {
let heap = HEAP.load(Ordering::Relaxed); let heap = HEAP.load(Ordering::Relaxed);
if heap.is_null() { if core::intrinsics::unlikely(heap.is_null()) {
// `HEAP` has not yet been successfully initialized // `HEAP` has not yet been successfully initialized
let heap = unsafe { GetProcessHeap() }; let heap = unsafe { GetProcessHeap() };
if !heap.is_null() { if !heap.is_null() {
@ -115,6 +115,16 @@ fn init_or_get_process_heap() -> c::HANDLE {
} }
} }
#[inline(never)]
fn process_heap_alloc(flags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID {
let heap = init_or_get_process_heap();
if core::intrinsics::unlikely(heap.is_null()) {
return ptr::null_mut();
}
// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
unsafe { HeapAlloc(heap, 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]
@ -133,25 +143,17 @@ unsafe fn get_process_heap() -> c::HANDLE {
// initialized. // initialized.
#[inline] #[inline]
unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 { unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
let heap = init_or_get_process_heap();
if heap.is_null() {
// Allocation has failed, could not get the current process heap.
return ptr::null_mut();
}
// Allocated memory will be either zeroed or uninitialized. // Allocated memory will be either zeroed or uninitialized.
let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 }; let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };
if layout.align() <= MIN_ALIGN { if layout.align() <= MIN_ALIGN {
// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
// The returned pointer points to the start of an allocated block. // The returned pointer points to the start of an allocated block.
unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 } process_heap_alloc(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();
// SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`. let ptr = process_heap_alloc(flags, total) as *mut u8;
let ptr = unsafe { HeapAlloc(heap, 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();