From be992ca672bab90cc1c58de9e3ad76a670a2248f Mon Sep 17 00:00:00 2001
From: Tom Dohrmann <erbse.13@gmx.de>
Date: Sun, 7 May 2023 11:01:42 +0200
Subject: [PATCH] fix stack update for x86_intrcc with error code

Previously the x86_intrcc calling convention built two
STACKALLOC_W_PROBING machine instructions if the function took an error
code. This was caused by an additional call to emitSPUpdate in
llvm/lib/Target/X86/X86FrameLowering.cpp:1650. The code for inlining the
stack probes only handles the first STACKALLOC_W_PROBING machine
instruction and ignores any futher `STACKALLOC_W_PROBING` instructions.
This lead to miscompilations where the stack pointer wasn't properly
updated (see https://github.com/rust-lang/rust/issues/109918).
This patch prevents a second STACKALLOC_W_PROBING instruction from being
built by hard-coding the stack update to `push rax`.

To be honest I don't quite understand why this didn't lead to more
noticeable miscompilations previously.
---
 llvm/lib/Target/X86/X86FrameLowering.cpp | 14 +++++++++++++-
 llvm/test/CodeGen/X86/x86-64-intrcc.ll   | 17 +++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/X86/X86FrameLowering.cpp b/llvm/lib/Target/X86/X86FrameLowering.cpp
index 10c57e43cf9506..f512de1ac3ac6b 100644
--- a/llvm/lib/Target/X86/X86FrameLowering.cpp
+++ b/llvm/lib/Target/X86/X86FrameLowering.cpp
@@ -1644,7 +1644,19 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
       Fn.arg_size() == 2) {
     StackSize += 8;
     MFI.setStackSize(StackSize);
-    emitSPUpdate(MBB, MBBI, DL, -8, /*InEpilogue=*/false);
+
+    // Update the stack pointer by pushing a register. This is the instruction
+    // emitted that would be end up being emitted by a call to `emitSPUpdate`.
+    // Hard-coding the update to a push avoids emitting a second
+    // `STACKALLOC_W_PROBING` instruction in the save block: We know that stack
+    // probing isn't needed anyways for an 8-byte update.
+    // Pushing a register leaves us in a similar situation to a regular
+    // function call where we know that the address at (rsp-8) is writeable.
+    // That way we avoid any off-by-ones with stack probing for additional
+    // stack pointer updates later on.
+    BuildMI(MBB, MBBI, DL, TII.get(X86::PUSH64r))
+        .addReg(X86::RAX, RegState::Undef)
+        .setMIFlag(MachineInstr::FrameSetup);
   }
 
   // If this is x86-64 and the Red Zone is not disabled, if we are a leaf
diff --git a/llvm/test/CodeGen/X86/x86-64-intrcc.ll b/llvm/test/CodeGen/X86/x86-64-intrcc.ll
index 1e1e01214c90fa..e11fc08c4fc480 100644
--- a/llvm/test/CodeGen/X86/x86-64-intrcc.ll
+++ b/llvm/test/CodeGen/X86/x86-64-intrcc.ll
@@ -174,5 +174,22 @@ entry:
   ret void
 }
 
+define x86_intrcc void @test_stack_allocation(ptr byval(%struct.interrupt_frame) %frame, i64 %err) #1 {
+  ; CHECK-LABEL: test_stack_allocation:
+  ; CHECK: # %bb.0: # %entry
+  
+  ;; Ensure that STACKALLOC_W_PROBING isn't emitted.
+  ; CHECK-NOT: # fixed size alloca with probing
+  ;; Ensure that stack space is allocated. 
+  ; CHECK: subq    $280, %rsp
+entry:
+  %some_allocation = alloca i64
+  ;; Call a un-inlineable function to ensure the allocation isn't put in the red zone.
+  call void @external_function(ptr %some_allocation)
+  ret void
+}
+
+declare void @external_function(ptr)
 
 attributes #0 = { nounwind "frame-pointer"="all" }
+attributes #1 = { nounwind "probe-stack"="inline-asm" }