[WinEH] Call llvm.stackrestore in __except blocks

We have to do this manually, the runtime only sets up ebp. Fixes a crash
when returning after catching an exception.

llvm-svn: 239451

2 files changed, 23 insertions, 1 deletions

diff --git a/llvm/lib/Target/X86/X86WinEHState.cpp b/llvm/lib/Target/X86/X86WinEHState.cpp
index 16a94daa819..afad3f930da 100644
--- a/llvm/lib/Target/X86/X86WinEHState.cpp
+++ b/llvm/lib/Target/X86/X86WinEHState.cpp
@@ -263,7 +263,6 @@ void WinEHStatePass::emitExceptionRegistrationRecord(Function *F) {
   if (Personality == EHPersonality::MSVC_CXX) {
     RegNodeTy = getCXXEHRegistrationType();
     RegNode = Builder.CreateAlloca(RegNodeTy);
-    // FIXME: We can skip this in -GS- mode, when we figure that out.
     // SavedESP = llvm.stacksave()
     Value *SP = Builder.CreateCall(
         Intrinsic::getDeclaration(TheModule, Intrinsic::stacksave), {});
@@ -490,6 +489,7 @@ void WinEHStatePass::addSEHStateStores(Function &F, MachineModuleInfo &MMI) {
 
   // Iterate all the instructions and emit state number stores.
   int CurState = 0;
+  SmallPtrSet<BasicBlock *, 4> ExceptBlocks;
   for (BasicBlock &BB : F) {
     for (auto I = BB.begin(), E = BB.end(); I != E; ++I) {
       if (auto *CI = dyn_cast<CallInst>(I)) {
@@ -517,11 +517,29 @@ void WinEHStatePass::addSEHStateStores(Function &F, MachineModuleInfo &MMI) {
           assert(!ActionList.empty());
           CurState += ActionList.size();
           State += ActionList.size() - 1;
+
+          // Remember all the __except block targets.
+          for (auto &Handler : ActionList) {
+            if (auto *CH = dyn_cast<CatchHandler>(Handler.get())) {
+              auto *BA = cast<BlockAddress>(CH->getHandlerBlockOrFunc());
+              ExceptBlocks.insert(BA->getBasicBlock());
+            }
+          }
         }
         insertStateNumberStore(RegNode, II, State);
       }
     }
   }
+
+  // Insert llvm.stackrestore into each __except block.
+  Function *StackRestore =
+      Intrinsic::getDeclaration(TheModule, Intrinsic::stackrestore);
+  for (BasicBlock *ExceptBB : ExceptBlocks) {
+    IRBuilder<> Builder(ExceptBB->begin());
+    Value *SP =
+        Builder.CreateLoad(Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
+    Builder.CreateCall(StackRestore, {SP});
+  }
 }
 
 /// Rewrite llvm.eh.exceptioncode and llvm.eh.exceptioninfo to memory loads in
diff --git a/llvm/test/CodeGen/X86/seh-safe-div-win32.ll b/llvm/test/CodeGen/X86/seh-safe-div-win32.ll
index 66967797076..e5cbc587bb2 100644
--- a/llvm/test/CodeGen/X86/seh-safe-div-win32.ll
+++ b/llvm/test/CodeGen/X86/seh-safe-div-win32.ll
@@ -77,11 +77,15 @@ __try.cont:
 
 ; CHECK: [[handler0:Ltmp[0-9]+]]: # Block address taken
 ; CHECK: # %handler0
+; 	Restore SP
+; CHECK: movl {{.*}}(%ebp), %esp
 ; CHECK: calll _puts
 ; CHECK: jmp [[cont_bb]]
 
 ; CHECK: [[handler1:Ltmp[0-9]+]]: # Block address taken
 ; CHECK: # %handler1
+; 	Restore SP
+; CHECK: movl {{.*}}(%ebp), %esp
 ; CHECK: calll _puts
 ; CHECK: jmp [[cont_bb]]