From 0a957ae73a50eeb135688d0574b064d6f5abf883 Mon Sep 17 00:00:00 2001
From: Xiaoyu Lu <xiaoyux.lu@intel.com>
Date: Fri, 14 Aug 2020 10:28:15 +0800
Subject: [PATCH] Add uefi arch x86 probestack support

1. In UEFI x86 arch, probestack need triple underscore.

2.  In UEFI, probestack function do things like _chkstk(in MSVC).
    MSVC x32's _chkstk and cygwin/mingw's _alloca adjust %esp themselves
    MSVC x64's __chkstk and cygwin/mingw's ___chkstk_ms do not adjust
    %rsp themselves.
    But current probestack doesn't adjust esp. And LLVM doesn't generate
    sub %eax, %esp after probestack. So we adjust esp in probestack like
    MSVC x32's _chkstk.
---
 library/compiler-builtins/src/probestack.rs | 68 ++++++++++++++++++++-
 1 file changed, 66 insertions(+), 2 deletions(-)

diff --git a/library/compiler-builtins/src/probestack.rs b/library/compiler-builtins/src/probestack.rs
index 9c78faa1d238..2f37a104ecee 100644
--- a/library/compiler-builtins/src/probestack.rs
+++ b/library/compiler-builtins/src/probestack.rs
@@ -76,7 +76,7 @@ macro_rules! define_rust_probestack {
     };
 }
 
-#[cfg(target_os = "uefi")]
+#[cfg(all(target_os = "uefi", target_arch = "x86_64"))]
 macro_rules! define_rust_probestack {
     ($body: expr) => {
         concat!(
@@ -104,6 +104,20 @@ macro_rules! define_rust_probestack {
     };
 }
 
+// In UEFI x86 arch, triple underscore is deliberate.
+#[cfg(all(target_os = "uefi", target_arch = "x86"))]
+macro_rules! define_rust_probestack {
+    ($body: expr) => {
+        concat!(
+            "
+            .globl ___rust_probestack
+        ___rust_probestack:
+            ",
+            $body
+        )
+    };
+}
+
 // Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
 // ensuring that if any pages are unmapped we'll make a page fault.
 //
@@ -231,7 +245,7 @@ global_asm!(define_rust_probestack!(
     "
 ));
 
-#[cfg(target_arch = "x86")]
+#[cfg(all(target_arch = "x86", not(target_os = "uefi")))]
 // This is the same as x86_64 above, only translated for 32-bit sizes. Note
 // that on Unix we're expected to restore everything as it was, this
 // function basically can't tamper with anything.
@@ -270,3 +284,53 @@ global_asm!(define_rust_probestack!(
     .cfi_endproc
     "
 ));
+
+#[cfg(all(target_arch = "x86", target_os = "uefi"))]
+// UEFI target is windows like target. LLVM will do _chkstk things like windows.
+// probestack function will also do things like _chkstk in MSVC.
+// So we need to sub %ax %sp in probestack when arch is x86.
+//
+// REF: Rust commit(74e80468347)
+// rust\src\llvm-project\llvm\lib\Target\X86\X86FrameLowering.cpp: 805
+// Comments in LLVM:
+//   MSVC x32's _chkstk and cygwin/mingw's _alloca adjust %esp themselves.
+//   MSVC x64's __chkstk and cygwin/mingw's ___chkstk_ms do not adjust %rsp
+//   themselves.
+global_asm!(define_rust_probestack!(
+    "
+    .cfi_startproc
+    push   %ebp
+    .cfi_adjust_cfa_offset 4
+    .cfi_offset %ebp, -8
+    mov    %esp, %ebp
+    .cfi_def_cfa_register %ebp
+    push   %ecx
+    push   %edx
+    mov    %eax,%ecx
+
+    cmp    $0x1000,%ecx
+    jna    3f
+2:
+    sub    $0x1000,%esp
+    test   %esp,8(%esp)
+    sub    $0x1000,%ecx
+    cmp    $0x1000,%ecx
+    ja     2b
+
+3:
+    sub    %ecx,%esp
+    test   %esp,8(%esp)
+    mov    4(%ebp),%edx
+    mov    %edx, 12(%esp)
+    add    %eax,%esp
+    pop    %edx
+    pop    %ecx
+    leave
+
+    sub   %eax, %esp
+    .cfi_def_cfa_register %esp
+    .cfi_adjust_cfa_offset -4
+    ret
+    .cfi_endproc
+    "
+));