From 27a6e6e67b6585e0f9bebc00617c22709c0d4a7e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E5=AE=87=E9=80=B8?= <Strawberry_Str@hotmail.com>
Date: Wed, 24 Jan 2024 02:43:06 +0800
Subject: [PATCH] Wrap `HeapAlloc` and never inline

---
 library/std/src/sys/pal/windows/alloc.rs | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/library/std/src/sys/pal/windows/alloc.rs b/library/std/src/sys/pal/windows/alloc.rs
index d53ea16005f..270eca37b14 100644
--- a/library/std/src/sys/pal/windows/alloc.rs
+++ b/library/std/src/sys/pal/windows/alloc.rs
@@ -95,7 +95,7 @@ fn HeapReAlloc(
 #[inline]
 fn init_or_get_process_heap() -> c::HANDLE {
     let heap = HEAP.load(Ordering::Relaxed);
-    if heap.is_null() {
+    if core::intrinsics::unlikely(heap.is_null()) {
         // `HEAP` has not yet been successfully initialized
         let heap = unsafe { GetProcessHeap() };
         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.
 // SAFETY: `HEAP` must have been successfully initialized.
 #[inline]
@@ -133,25 +143,17 @@ unsafe fn get_process_heap() -> c::HANDLE {
 // initialized.
 #[inline]
 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.
     let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };
 
     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.
-        unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 }
+        process_heap_alloc(flags, layout.size()) as *mut u8
     } else {
         // Allocate extra padding in order to be able to satisfy the alignment.
         let total = layout.align() + layout.size();
 
-        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
-        let ptr = unsafe { HeapAlloc(heap, flags, total) as *mut u8 };
+        let ptr = process_heap_alloc(flags, total) as *mut u8;
         if ptr.is_null() {
             // Allocation has failed.
             return ptr::null_mut();