From 8ca85febe08f854a30164393f07793d6a212aa2c Mon Sep 17 00:00:00 2001
From: Nicholas Mosier <nmosier@stanford.edu>
Date: Sat, 20 Jan 2024 17:28:41 +0000
Subject: [PATCH] fix bug

---
 llvm/lib/CodeGen/SafeStack.cpp                | 50 +++++++++++--------
 .../SafeStack/X86/builtin-setjmp.ll           | 26 ++++++++++
 2 files changed, 55 insertions(+), 21 deletions(-)
 create mode 100644 llvm/test/Transforms/SafeStack/X86/builtin-setjmp.ll

diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index 0a26247a4d16..485a3b10580f 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -400,6 +400,12 @@ void SafeStack::findInsts(Function &F,
         Returns.push_back(CI);
       else
         Returns.push_back(RI);
+    } else if (auto II = dyn_cast<IntrinsicInst>(&I)) {
+      if (II->getIntrinsicID() == Intrinsic::gcroot)
+        report_fatal_error(
+            "gcroot intrinsic not compatible with safestack attribute");
+      if (II->getIntrinsicID() == Intrinsic::eh_sjlj_setjmp)
+        StackRestorePoints.push_back(II);
     } else if (auto CI = dyn_cast<CallInst>(&I)) {
       // setjmps require stack restore.
       if (CI->getCalledFunction() && CI->canReturnTwice())
@@ -407,10 +413,6 @@ void SafeStack::findInsts(Function &F,
     } else if (auto LP = dyn_cast<LandingPadInst>(&I)) {
       // Exception landing pads require stack restore.
       StackRestorePoints.push_back(LP);
-    } else if (auto II = dyn_cast<IntrinsicInst>(&I)) {
-      if (II->getIntrinsicID() == Intrinsic::gcroot)
-        report_fatal_error(
-            "gcroot intrinsic not compatible with safestack attribute");
     }
   }
   for (Argument &Arg : F.args()) {
@@ -697,24 +699,30 @@ void SafeStack::moveDynamicAllocasToUnsafeStack(
 
   if (!DynamicAllocas.empty()) {
     // Now go through the instructions again, replacing stacksave/stackrestore.
-    for (Instruction &I : llvm::make_early_inc_range(instructions(&F))) {
-      auto *II = dyn_cast<IntrinsicInst>(&I);
-      if (!II)
-        continue;
-
-      if (II->getIntrinsicID() == Intrinsic::stacksave) {
-        IRBuilder<> IRB(II);
-        Instruction *LI = IRB.CreateLoad(StackPtrTy, UnsafeStackPtr);
-        LI->takeName(II);
-        II->replaceAllUsesWith(LI);
-        II->eraseFromParent();
-      } else if (II->getIntrinsicID() == Intrinsic::stackrestore) {
-        IRBuilder<> IRB(II);
-        Instruction *SI = IRB.CreateStore(II->getArgOperand(0), UnsafeStackPtr);
-        SI->takeName(II);
-        assert(II->use_empty());
-        II->eraseFromParent();
+    SmallVector<IntrinsicInst *> StackSaves;
+    for (Instruction &I : instructions(F))
+      if (auto *II = dyn_cast<IntrinsicInst>(&I);
+          II && II->getIntrinsicID() == Intrinsic::stacksave)
+        StackSaves.push_back(II);
+    for (IntrinsicInst *StackSave : StackSaves) {
+      Instruction *LI = nullptr;
+      for (Use &U : llvm::make_early_inc_range(StackSave->uses())) {
+        auto *UserII = dyn_cast<IntrinsicInst>(U.getUser());
+        if (!UserII)
+          continue;
+        if (UserII->getIntrinsicID() == Intrinsic::stackrestore) {
+          if (!LI)
+            LI = IRBuilder<>(StackSave).CreateLoad(StackPtrTy, UnsafeStackPtr);
+          IRBuilder<>(UserII).CreateStore(LI, UnsafeStackPtr);
+          UserII->eraseFromParent();
+        } else {
+          report_fatal_error("SafeStack does not support passing the result of "
+                             "llvm.stacksave to incompatible intrinsic " +
+                             Intrinsic::getName(UserII->getIntrinsicID()));
+        }
       }
+      if (StackSave->use_empty())
+        StackSave->eraseFromParent();
     }
   }
 }
diff --git a/llvm/test/Transforms/SafeStack/X86/builtin-setjmp.ll b/llvm/test/Transforms/SafeStack/X86/builtin-setjmp.ll
new file mode 100644
index 000000000000..219abd9f37d1
--- /dev/null
+++ b/llvm/test/Transforms/SafeStack/X86/builtin-setjmp.ll
@@ -0,0 +1,26 @@
+; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -passes=safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -passes=safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+
+define void @main() safestack {
+  ; CHECK: %[[SP:.*]] = load ptr, ptr @__safestack_unsafe_stack_ptr
+  ; CHECK: %[[STATICTOP:.*]] = getelementptr i8, ptr %[[SP]]
+  %env = alloca [5 x ptr], align 16
+  ; CHECK: %sp = call ptr @llvm.stacksave.p0()
+  %sp = call ptr @llvm.stacksave.p0()
+  %env_sp_slot = getelementptr inbounds ptr, ptr %env, i64 2
+  store ptr %sp, ptr %env_sp_slot, align 16
+  ; CHECK: %setjmp_rv = call i32 @llvm.eh.sjlj.setjmp(
+  %setjmp_rv = call i32 @llvm.eh.sjlj.setjmp(ptr %env)
+  ; CHECK: store ptr %[[STATICTOP]], ptr @__safestack_unsafe_stack_ptr
+  %cmp = icmp eq i32 %setjmp_rv, 0
+  br i1 %cmp, label %longjmp, label %end
+
+longjmp:
+  call void @llvm.eh.sjlj.longjmp(ptr %env)
+  unreachable
+
+end:
+  ret void
+}
-- 
2.34.1