[WinEH] Make setjmp work correctly with EH

[majnemer]

Summary

32-bit X86 EH on Windows utilizes a stack of registration nodes
allocated and deallocated on entry/exit. A registration node contains a
bunch of EH personality specific information like which try-state we are
currently in.

Because a setjmp target allows control flow from arbitrary program
points, there is no way to ensure that the try-state we are in is
correctly updated once we transfer control.

MSVC compatible compilers, like MSVC and ICC, utilize runtime helpers to
reinitialize the try-state when a longjmp occurs. This is implemented
by adding additional arguments to _setjmp3: the desired try-state and
a helper routine to update the try-state.

[majnemer]

• Simplify SetJmp3CallSites

[rnk]

lgtm

lib/Target/X86/X86WinEHState.cpp

520 ↗

(On Diff #49381)

Anything wrong with the old spelling?

llvm/trunk/lib/Target/X86/X86WinEHState.cpp

private:
   void unlinkExceptionRegistration(IRBuilder<> &Builder);
   void addStateStores(Function &F, WinEHFuncInfo &FuncInfo);
   void insertStateNumberStore(Instruction *IP, int State);
 
   Value *emitEHLSDA(IRBuilder<> &Builder, Function *F);
 
   Function *generateLSDAInEAXThunk(Function *ParentFunc);
 
+  bool isStateStoreNeeded(EHPersonality Personality, CallSite CS);
+  void rewriteSetJmpCallSite(IRBuilder<> &Builder, Function &F, CallSite CS,
+                             Value *State);
+  int getBaseStateForBB(DenseMap<BasicBlock *, ColorVector> &BlockColors,
+                        WinEHFuncInfo &FuncInfo, BasicBlock *BB);
+  int getStateForCallSite(DenseMap<BasicBlock *, ColorVector> &BlockColors,
+                          WinEHFuncInfo &FuncInfo, CallSite CS);
+
   // Module-level type getters.
   Type *getEHLinkRegistrationType();
   Type *getSEHRegistrationType();
   Type *getCXXEHRegistrationType();
 
   // Per-module data.
   Module *TheModule = nullptr;
   StructType *EHLinkRegistrationTy = nullptr;
   StructType *CXXEHRegistrationTy = nullptr;
   StructType *SEHRegistrationTy = nullptr;
-  Function *FrameRecover = nullptr;
-  Function *FrameAddress = nullptr;
-  Function *FrameEscape = nullptr;
+  Constant *SetJmp3 = nullptr;
+  Constant *CxxLongjmpUnwind = nullptr;
 
   // Per-function state
   EHPersonality Personality = EHPersonality::Unknown;
   Function *PersonalityFn = nullptr;
   bool UseStackGuard = false;
   int ParentBaseState;
+  Constant *SehLongjmpUnwind = nullptr;
+  Constant *Cookie = nullptr;
 
   /// The stack allocation containing all EH data, including the link in the
   /// fs:00 chain and the current state.
   AllocaInst *RegNode = nullptr;
 
   /// The index of the state field of RegNode.
   int StateFieldIndex = ~0U;
 

 }
 
 bool WinEHStatePass::doFinalization(Module &M) {
   assert(TheModule == &M);
   TheModule = nullptr;
   EHLinkRegistrationTy = nullptr;
   CXXEHRegistrationTy = nullptr;
   SEHRegistrationTy = nullptr;
-  FrameEscape = nullptr;
-  FrameRecover = nullptr;
-  FrameAddress = nullptr;
+  SetJmp3 = nullptr;
+  CxxLongjmpUnwind = nullptr;
+  SehLongjmpUnwind = nullptr;
+  Cookie = nullptr;
   return false;
 }
 
 void WinEHStatePass::getAnalysisUsage(AnalysisUsage &AU) const {
   // This pass should only insert a stack allocation, memory accesses, and
   // localrecovers.
   AU.setPreservesCFG();
 }
for (BasicBlock &BB : F) {
     if (BB.isEHPad()) {
       HasPads = true;
       break;
     }
   }
   if (!HasPads)
     return false;
 
-  FrameEscape = Intrinsic::getDeclaration(TheModule, Intrinsic::localescape);
-  FrameRecover = Intrinsic::getDeclaration(TheModule, Intrinsic::localrecover);
-  FrameAddress = Intrinsic::getDeclaration(TheModule, Intrinsic::frameaddress);
+  Type *Int8PtrType = Type::getInt8PtrTy(TheModule->getContext());
+  SetJmp3 = TheModule->getOrInsertFunction(
+      "_setjmp3", FunctionType::get(
+                      Type::getInt32Ty(TheModule->getContext()),
+                      {Int8PtrType, Type::getInt32Ty(TheModule->getContext())},
+                      /*isVarArg=*/true));
 
   // Disable frame pointer elimination in this function.
   // FIXME: Do the nested handlers need to keep the parent ebp in ebp, or can we
   // use an arbitrary register?
   F.addFnAttr("no-frame-pointer-elim", "true");
 
   emitExceptionRegistrationRecord(&F);
 
if (Personality == EHPersonality::MSVC_CXX) {
     // TryLevel = -1
     StateFieldIndex = 2;
     ParentBaseState = -1;
     insertStateNumberStore(&*Builder.GetInsertPoint(), ParentBaseState);
     // Handler = __ehhandler$F
     Function *Trampoline = generateLSDAInEAXThunk(F);
     Link = Builder.CreateStructGEP(RegNodeTy, RegNode, 1);
     linkExceptionRegistration(Builder, Trampoline);
+
+    CxxLongjmpUnwind = TheModule->getOrInsertFunction(
+        "__CxxLongjmpUnwind",
+        FunctionType::get(Type::getVoidTy(TheModule->getContext()), Int8PtrType,
+                          /*isVarArg=*/false));
+    cast<Function>(CxxLongjmpUnwind->stripPointerCasts())
+        ->setCallingConv(CallingConv::X86_StdCall);
   } else if (Personality == EHPersonality::MSVC_X86SEH) {
     // If _except_handler4 is in use, some additional guard checks and prologue
     // stuff is required.
     RegNodeTy = getSEHRegistrationType();
     RegNode = Builder.CreateAlloca(RegNodeTy);
     // SavedESP = llvm.stacksave()
     Value *SP = Builder.CreateCall(
         Intrinsic::getDeclaration(TheModule, Intrinsic::stacksave), {});
     Builder.CreateStore(SP, Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
     // TryLevel = -2 / -1
     StateFieldIndex = 4;
     StringRef PersonalityName = PersonalityFn->getName();
     UseStackGuard = (PersonalityName == "_except_handler4");
     ParentBaseState = UseStackGuard ? -2 : -1;
     insertStateNumberStore(&*Builder.GetInsertPoint(), ParentBaseState);
     // ScopeTable = llvm.x86.seh.lsda(F)
-    Value *FI8 = Builder.CreateBitCast(F, Int8PtrType);
-    Value *LSDA = Builder.CreateCall(
-        Intrinsic::getDeclaration(TheModule, Intrinsic::x86_seh_lsda), FI8);
+    Value *LSDA = emitEHLSDA(Builder, F);
     Type *Int32Ty = Type::getInt32Ty(TheModule->getContext());
     LSDA = Builder.CreatePtrToInt(LSDA, Int32Ty);
     // If using _except_handler4, xor the address of the table with
     // __security_cookie.
     if (UseStackGuard) {
-      Value *Cookie =
-          TheModule->getOrInsertGlobal("__security_cookie", Int32Ty);
+      Cookie = TheModule->getOrInsertGlobal("__security_cookie", Int32Ty);
       Value *Val = Builder.CreateLoad(Int32Ty, Cookie);
       LSDA = Builder.CreateXor(LSDA, Val);
     }
     Builder.CreateStore(LSDA, Builder.CreateStructGEP(RegNodeTy, RegNode, 3));
     Link = Builder.CreateStructGEP(RegNodeTy, RegNode, 2);
     linkExceptionRegistration(Builder, PersonalityFn);
+
+    SehLongjmpUnwind = TheModule->getOrInsertFunction(
+        UseStackGuard ? "_seh_longjmp_unwind4" : "_seh_longjmp_unwind",
+        FunctionType::get(Type::getVoidTy(TheModule->getContext()), Int8PtrType,
+                          /*isVarArg=*/false));
+    cast<Function>(SehLongjmpUnwind->stripPointerCasts())
+        ->setCallingConv(CallingConv::X86_StdCall);
   } else {
     llvm_unreachable("unexpected personality function");
   }
 
   // Insert an unlink before all returns.
   for (BasicBlock &BB : *F) {
     TerminatorInst *T = BB.getTerminator();
     if (!isa<ReturnInst>(T))
void WinEHStatePass::unlinkExceptionRegistration(IRBuilder<> &Builder) {
   // [fs:00] = Link->Next
   Value *Next =
       Builder.CreateLoad(Builder.CreateStructGEP(LinkTy, Link, 0));
   Constant *FSZero =
       Constant::getNullValue(LinkTy->getPointerTo()->getPointerTo(257));
   Builder.CreateStore(Next, FSZero);
 }
 
+// Calls to setjmp(p) are lowered to _setjmp3(p, 0) by the frontend.
+// The idea behind _setjmp3 is that it takes an optional number of personality
+// specific parameters to indicate how to restore the personality-specific frame
+// state when longjmp is initiated.  Typically, the current TryLevel is saved.
+void WinEHStatePass::rewriteSetJmpCallSite(IRBuilder<> &Builder, Function &F,
+                                           CallSite CS, Value *State) {
+  // Don't rewrite calls with a weird number of arguments.
+  if (CS.getNumArgOperands() != 2)
+    return;
+
+  Instruction *Inst = CS.getInstruction();
+
+  SmallVector<OperandBundleDef, 1> OpBundles;
+  CS.getOperandBundlesAsDefs(OpBundles);
+
+  SmallVector<Value *, 3> OptionalArgs;
+  if (Personality == EHPersonality::MSVC_CXX) {
+    OptionalArgs.push_back(CxxLongjmpUnwind);
+    OptionalArgs.push_back(State);
+    OptionalArgs.push_back(emitEHLSDA(Builder, &F));
+  } else if (Personality == EHPersonality::MSVC_X86SEH) {
+    OptionalArgs.push_back(SehLongjmpUnwind);
+    OptionalArgs.push_back(State);
+    if (UseStackGuard)
+      OptionalArgs.push_back(Cookie);
+  } else {
+    llvm_unreachable("unhandled personality!");
+  }
+
+  SmallVector<Value *, 5> Args;
+  Args.push_back(
+      Builder.CreateBitCast(CS.getArgOperand(0), Builder.getInt8PtrTy()));
+  Args.push_back(Builder.getInt32(OptionalArgs.size()));
+  Args.append(OptionalArgs.begin(), OptionalArgs.end());
+
+  CallSite NewCS;
+  if (CS.isCall()) {
+    auto *CI = cast<CallInst>(Inst);
+    CallInst *NewCI = Builder.CreateCall(SetJmp3, Args, OpBundles);
+    NewCI->setTailCallKind(CI->getTailCallKind());
+    NewCS = NewCI;
+  } else {
+    auto *II = cast<InvokeInst>(Inst);
+    NewCS = Builder.CreateInvoke(
+        SetJmp3, II->getNormalDest(), II->getUnwindDest(), Args, OpBundles);
+  }
+  NewCS.setCallingConv(CS.getCallingConv());
+  NewCS.setAttributes(CS.getAttributes());
+  NewCS->setDebugLoc(CS->getDebugLoc());
+
+  Instruction *NewInst = NewCS.getInstruction();
+  NewInst->takeName(Inst);
+  Inst->replaceAllUsesWith(NewInst);
+  Inst->eraseFromParent();
+}
+
 // Figure out what state we should assign calls in this block.
-static int getBaseStateForBB(DenseMap<BasicBlock *, ColorVector> &BlockColors,
-                             WinEHFuncInfo &FuncInfo, BasicBlock *BB) {
-  int BaseState = -1;
+int WinEHStatePass::getBaseStateForBB(
+    DenseMap<BasicBlock *, ColorVector> &BlockColors, WinEHFuncInfo &FuncInfo,
+    BasicBlock *BB) {
+  int BaseState = ParentBaseState;
   auto &BBColors = BlockColors[BB];
 
   assert(BBColors.size() == 1 && "multi-color BB not removed by preparation");
   BasicBlock *FuncletEntryBB = BBColors.front();
   if (auto *FuncletPad =
           dyn_cast<FuncletPadInst>(FuncletEntryBB->getFirstNonPHI())) {
     auto BaseStateI = FuncInfo.FuncletBaseStateMap.find(FuncletPad);
     if (BaseStateI != FuncInfo.FuncletBaseStateMap.end())
       BaseState = BaseStateI->second;
   }
 
   return BaseState;
 }
 
 // Calculate the state a call-site is in.
-static int getStateForCallSite(DenseMap<BasicBlock *, ColorVector> &BlockColors,
-                               WinEHFuncInfo &FuncInfo, CallSite CS) {
+int WinEHStatePass::getStateForCallSite(
+    DenseMap<BasicBlock *, ColorVector> &BlockColors, WinEHFuncInfo &FuncInfo,
+    CallSite CS) {
   if (auto *II = dyn_cast<InvokeInst>(CS.getInstruction())) {
     // Look up the state number of the EH pad this unwinds to.
     assert(FuncInfo.InvokeStateMap.count(II) && "invoke has no state!");
     return FuncInfo.InvokeStateMap[II];
   }
   // Possibly throwing call instructions have no actions to take after
   // an unwind. Ensure they are in the -1 state.
   return getBaseStateForBB(BlockColors, FuncInfo, CS.getParent());
for (BasicBlock *SuccBB : successors(BB)) {
     // conservatively report this basic block as overdefined.
     if (CommonState != SuccState)
       return OverdefinedState;
   }
 
   return CommonState;
 }
 
-static bool isStateStoreNeeded(EHPersonality Personality, CallSite CS) {
+bool WinEHStatePass::isStateStoreNeeded(EHPersonality Personality,
+                                        CallSite CS) {
   if (!CS)
     return false;
 
+  // If the function touches memory, it needs a state store.
   if (isAsynchronousEHPersonality(Personality))
     return !CS.doesNotAccessMemory();
 
+  // If the function throws, it needs a state store.
   return !CS.doesNotThrow();
 }
 
 void WinEHStatePass::addStateStores(Function &F, WinEHFuncInfo &FuncInfo) {
   // Mark the registration node. The backend needs to know which alloca it is so
   // that it can recover the original frame pointer.
   IRBuilder<> Builder(RegNode->getParent(), std::next(RegNode->getIterator()));
   Value *RegNodeI8 = Builder.CreateBitCast(RegNode, Builder.getInt8PtrTy());
for (BasicBlock *BB : RPOT) {
     }
 
     // We might have hoisted a state store into this block, emit it now.
     auto EndState = FinalStates.find(BB);
     if (EndState != FinalStates.end())
       if (EndState->second != PrevState)
         insertStateNumberStore(BB->getTerminator(), EndState->second);
   }
+
+  SmallVector<CallSite, 1> SetJmp3CallSites;
+  for (BasicBlock *BB : RPOT) {
+    for (Instruction &I : *BB) {
+      CallSite CS(&I);
+      if (!CS)
+        continue;
+      if (CS.getCalledValue()->stripPointerCasts() !=
+          SetJmp3->stripPointerCasts())
+        continue;
+
+      SetJmp3CallSites.push_back(CS);
+    }
+  }
+
+  for (CallSite CS : SetJmp3CallSites) {
+    auto &BBColors = BlockColors[CS->getParent()];
+    BasicBlock *FuncletEntryBB = BBColors.front();
+    bool InCleanup = isa<CleanupPadInst>(FuncletEntryBB->getFirstNonPHI());
+
+    IRBuilder<> Builder(CS.getInstruction());
+    Value *State;
+    if (InCleanup) {
+      Value *StateField =
+          Builder.CreateStructGEP(nullptr, RegNode, StateFieldIndex);
+      State = Builder.CreateLoad(StateField);
+    } else {
+      State = Builder.getInt32(getStateForCallSite(BlockColors, FuncInfo, CS));
+    }
+    rewriteSetJmpCallSite(Builder, F, CS, State);
+  }
 }
 
 void WinEHStatePass::insertStateNumberStore(Instruction *IP, int State) {
   IRBuilder<> Builder(IP);
   Value *StateField =
       Builder.CreateStructGEP(nullptr, RegNode, StateFieldIndex);
   Builder.CreateStore(Builder.getInt32(State), StateField);
 }

llvm/trunk/test/CodeGen/WinEH/wineh-setjmp.ll

+; RUN: opt -mtriple=i686-pc-windows-msvc -S -x86-winehstate  < %s | FileCheck %s
+target datalayout = "e-m:x-p:32:32-i64:64-f80:32-n8:16:32-a:0:32-S32"
+target triple = "i686-pc-windows-msvc"
+
+@jb = external global i8
+
+define i32 @test1() personality i32 (...)* @__CxxFrameHandler3 {
+entry:
+; CHECK-LABEL: define i32 @test1(
+; CHECK: %[[eh_reg:.*]] = alloca
+; CHECK: %[[gep:.*]] = getelementptr inbounds {{.*}}, {{.*}} %[[eh_reg]], i32 0, i32 2
+; CHECK: store i32 -1, i32* %[[gep]]
+; CHECK: %[[gep:.*]] = getelementptr inbounds {{.*}}, {{.*}} %[[eh_reg]], i32 0, i32 2
+; CHECK: store i32 0, i32* %[[gep]]
+; CHECK: %[[lsda:.*]] = call i8* @llvm.x86.seh.lsda(i8* bitcast (i32 ()* @test1 to i8*))
+; CHECK: invoke i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 3, void (i8*)* @__CxxLongjmpUnwind, i32 0, i8* %[[lsda]])
+  %inv = invoke i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 0) #2
+          to label %invoke.cont unwind label %ehcleanup
+
+invoke.cont:
+; CHECK: %[[gep:.*]] = getelementptr inbounds {{.*}}, {{.*}} %[[eh_reg]], i32 0, i32 2
+; CHECK: store i32 -1, i32* %[[gep]]
+; CHECK: %[[lsda:.*]] = call i8* @llvm.x86.seh.lsda(i8* bitcast (i32 ()* @test1 to i8*))
+; CHECK: call i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 3, void (i8*)* @__CxxLongjmpUnwind, i32 -1, i8* %[[lsda]])
+  call i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 0)
+  call void @cleanup()
+  ret i32 %inv
+
+ehcleanup:
+  %cp = cleanuppad within none []
+; CHECK: %[[gep:.*]] = getelementptr inbounds {{.*}}, {{.*}} %[[eh_reg]], i32 0, i32 2
+; CHECK: %[[load:.*]] = load i32, i32* %[[gep]]
+; CHECK: %[[lsda:.*]] = call i8* @llvm.x86.seh.lsda(i8* bitcast (i32 ()* @test1 to i8*))
+; CHECK: call i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 3, void (i8*)* @__CxxLongjmpUnwind, i32 %[[load]], i8* %[[lsda]]) [ "funclet"(token %cp) ]
+  %cal = call i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 0) [ "funclet"(token %cp) ]
+  call void @cleanup() [ "funclet"(token %cp) ]
+  cleanupret from %cp unwind to caller
+}
+
+define i32 @test2() personality i32 (...)* @_except_handler3 {
+entry:
+; CHECK-LABEL: define i32 @test2(
+; CHECK: %[[eh_reg:.*]] = alloca
+; CHECK: %[[gep:.*]] = getelementptr inbounds {{.*}}, {{.*}} %[[eh_reg]], i32 0, i32 4
+; CHECK: store i32 -1, i32* %[[gep]]
+; CHECK: %[[gep:.*]] = getelementptr inbounds {{.*}}, {{.*}} %[[eh_reg]], i32 0, i32 4
+; CHECK: store i32 0, i32* %[[gep]]
+; CHECK: invoke i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 2, void (i8*)* @_seh_longjmp_unwind, i32 0)
+  %inv = invoke i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 0) #2
+          to label %invoke.cont unwind label %ehcleanup
+
+invoke.cont:
+; CHECK: %[[gep:.*]] = getelementptr inbounds {{.*}}, {{.*}} %[[eh_reg]], i32 0, i32 4
+; CHECK: store i32 -1, i32* %[[gep]]
+; CHECK: call i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 2, void (i8*)* @_seh_longjmp_unwind, i32 -1)
+  call i32 (i8*, i32, ...) @_setjmp3(i8* @jb, i32 0)
+  call void @cleanup()
+  ret i32 %inv
+
+ehcleanup:
+  %cp = cleanuppad within none []
+  call void @cleanup() [ "funclet"(token %cp) ]
+  cleanupret from %cp unwind to caller
+}
+
+; Function Attrs: returns_twice
+declare i32 @_setjmp3(i8*, i32, ...) #2
+
+declare i32 @__CxxFrameHandler3(...)
+
+declare i32 @_except_handler3(...)
+
+declare void @cleanup()
+
+attributes #2 = { returns_twice }