llvm/lib/Target/X86/X86WinEHState.cpp

   1 //===-- X86WinEHState - Insert EH state updates for win32 exceptions ------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // All functions using an MSVC EH personality use an explicitly updated state
  10 // number stored in an exception registration stack object. The registration
  11 // object is linked into a thread-local chain of registrations stored at fs:00.
  12 // This pass adds the registration object and EH state updates.
  13 //
  14 //===----------------------------------------------------------------------===//
  15
  16 #include "X86.h"
  17 #include "llvm/ADT/PostOrderIterator.h"
  18 #include "llvm/Analysis/CFG.h"
  19 #include "llvm/CodeGen/MachineModuleInfo.h"
  20 #include "llvm/CodeGen/WinEHFuncInfo.h"
  21 #include "llvm/IR/CFG.h"
  22 #include "llvm/IR/EHPersonalities.h"
  23 #include "llvm/IR/Function.h"
  24 #include "llvm/IR/IRBuilder.h"
  25 #include "llvm/IR/Instructions.h"
  26 #include "llvm/IR/IntrinsicInst.h"
  27 #include "llvm/IR/Intrinsics.h"
  28 #include "llvm/IR/IntrinsicsX86.h"
  29 #include "llvm/IR/Module.h"
  30 #include "llvm/Pass.h"
  31 #include "llvm/Support/Debug.h"
  32 #include <deque>
  33
  34 using namespace llvm;
  35
  36 #define DEBUG_TYPE "winehstate"
  37
  38 namespace {
  39 const int OverdefinedState = INT_MIN;
  40
  41 class WinEHStatePass : public FunctionPass {
  42 public:
  43   static char ID; // Pass identification, replacement for typeid.
  44
  45   WinEHStatePass() : FunctionPass(ID) {}
  46
  47   bool runOnFunction(Function &Fn) override;
  48
  49   bool doInitialization(Module &M) override;
  50
  51   bool doFinalization(Module &M) override;
  52
  53   void getAnalysisUsage(AnalysisUsage &AU) const override;
  54
  55   StringRef getPassName() const override {
  56     return "Windows 32-bit x86 EH state insertion";
  57   }
  58
  59 private:
  60   void emitExceptionRegistrationRecord(Function *F);
  61
  62   void linkExceptionRegistration(IRBuilder<> &Builder, Function *Handler);
  63   void unlinkExceptionRegistration(IRBuilder<> &Builder);
  64   void addStateStores(Function &F, WinEHFuncInfo &FuncInfo);
  65   void insertStateNumberStore(Instruction *IP, int State);
  66
  67   Value *emitEHLSDA(IRBuilder<> &Builder, Function *F);
  68
  69   Function *generateLSDAInEAXThunk(Function *ParentFunc);
  70
  71   bool isStateStoreNeeded(EHPersonality Personality, CallBase &Call);
  72   void rewriteSetJmpCall(IRBuilder<> &Builder, Function &F, CallBase &Call,
  73                          Value *State);
  74   int getBaseStateForBB(DenseMap<BasicBlock *, ColorVector> &BlockColors,
  75                         WinEHFuncInfo &FuncInfo, BasicBlock *BB);
  76   int getStateForCall(DenseMap<BasicBlock *, ColorVector> &BlockColors,
  77                       WinEHFuncInfo &FuncInfo, CallBase &Call);
  78
  79   void updateEspForInAllocas(Function &F);
  80
  81   // Module-level type getters.
  82   Type *getEHLinkRegistrationType();
  83   Type *getSEHRegistrationType();
  84   Type *getCXXEHRegistrationType();
  85
  86   // Per-module data.
  87   Module *TheModule = nullptr;
  88   StructType *EHLinkRegistrationTy = nullptr;
  89   StructType *CXXEHRegistrationTy = nullptr;
  90   StructType *SEHRegistrationTy = nullptr;
  91   FunctionCallee SetJmp3 = nullptr;
  92   FunctionCallee CxxLongjmpUnwind = nullptr;
  93
  94   // Per-function state
  95   EHPersonality Personality = EHPersonality::Unknown;
  96   Function *PersonalityFn = nullptr;
  97   bool UseStackGuard = false;
  98   int ParentBaseState = 0;
  99   FunctionCallee SehLongjmpUnwind = nullptr;
 100   Constant *Cookie = nullptr;
 101
 102   /// The stack allocation containing all EH data, including the link in the
 103   /// fs:00 chain and the current state.
 104   AllocaInst *RegNode = nullptr;
 105
 106   // Struct type of RegNode. Used for GEPing.
 107   Type *RegNodeTy = nullptr;
 108
 109   // The allocation containing the EH security guard.
 110   AllocaInst *EHGuardNode = nullptr;
 111
 112   /// The index of the state field of RegNode.
 113   int StateFieldIndex = ~0U;
 114
 115   /// The linked list node subobject inside of RegNode.
 116   Value *Link = nullptr;
 117 };
 118 } // namespace
 119
 120 FunctionPass *llvm::createX86WinEHStatePass() { return new WinEHStatePass(); }
 121
 122 char WinEHStatePass::ID = 0;
 123
 124 INITIALIZE_PASS(WinEHStatePass, "x86-winehstate",
 125                 "Insert stores for EH state numbers", false, false)
 126
 127 bool WinEHStatePass::doInitialization(Module &M) {
 128   TheModule = &M;
 129   return false;
 130 }
 131
 132 bool WinEHStatePass::doFinalization(Module &M) {
 133   assert(TheModule == &M);
 134   TheModule = nullptr;
 135   EHLinkRegistrationTy = nullptr;
 136   CXXEHRegistrationTy = nullptr;
 137   SEHRegistrationTy = nullptr;
 138   SetJmp3 = nullptr;
 139   CxxLongjmpUnwind = nullptr;
 140   SehLongjmpUnwind = nullptr;
 141   Cookie = nullptr;
 142   return false;
 143 }
 144
 145 void WinEHStatePass::getAnalysisUsage(AnalysisUsage &AU) const {
 146   // This pass should only insert a stack allocation, memory accesses, and
 147   // localrecovers.
 148   AU.setPreservesCFG();
 149 }
 150
 151 bool WinEHStatePass::runOnFunction(Function &F) {
 152   // Don't insert state stores or exception handler thunks for
 153   // available_externally functions. The handler needs to reference the LSDA,
 154   // which will not be emitted in this case.
 155   if (F.hasAvailableExternallyLinkage())
 156     return false;
 157
 158   // Check the personality. Do nothing if this personality doesn't use funclets.
 159   if (!F.hasPersonalityFn())
 160     return false;
 161   PersonalityFn = dyn_cast<Function>(F.getPersonalityFn()->stripPointerCasts());
 162   if (!PersonalityFn)
 163     return false;
 164   Personality = classifyEHPersonality(PersonalityFn);
 165   if (!isFuncletEHPersonality(Personality))
 166     return false;
 167
 168   // Skip this function if there are no EH pads and we aren't using IR-level
 169   // outlining.
 170   bool HasPads = false;
 171   for (BasicBlock &BB : F) {
 172     if (BB.isEHPad()) {
 173       HasPads = true;
 174       break;
 175     }
 176   }
 177   if (!HasPads)
 178     return false;
 179
 180   Type *Int8PtrType = PointerType::getUnqual(TheModule->getContext());
 181   SetJmp3 = TheModule->getOrInsertFunction(
 182       "_setjmp3", FunctionType::get(
 183                       Type::getInt32Ty(TheModule->getContext()),
 184                       {Int8PtrType, Type::getInt32Ty(TheModule->getContext())},
 185                       /*isVarArg=*/true));
 186
 187   emitExceptionRegistrationRecord(&F);
 188
 189   // The state numbers calculated here in IR must agree with what we calculate
 190   // later on for the MachineFunction. In particular, if an IR pass deletes an
 191   // unreachable EH pad after this point before machine CFG construction, we
 192   // will be in trouble. If this assumption is ever broken, we should turn the
 193   // numbers into an immutable analysis pass.
 194   WinEHFuncInfo FuncInfo;
 195   addStateStores(F, FuncInfo);
 196   updateEspForInAllocas(F);
 197
 198   // Reset per-function state.
 199   PersonalityFn = nullptr;
 200   Personality = EHPersonality::Unknown;
 201   UseStackGuard = false;
 202   RegNodeTy = nullptr;
 203   RegNode = nullptr;
 204   EHGuardNode = nullptr;
 205
 206   return true;
 207 }
 208
 209 /// Get the common EH registration subobject:
 210 ///   typedef _EXCEPTION_DISPOSITION (*PEXCEPTION_ROUTINE)(
 211 ///       _EXCEPTION_RECORD *, void *, _CONTEXT *, void *);
 212 ///   struct EHRegistrationNode {
 213 ///     EHRegistrationNode *Next;
 214 ///     PEXCEPTION_ROUTINE Handler;
 215 ///   };
 216 Type *WinEHStatePass::getEHLinkRegistrationType() {
 217   if (EHLinkRegistrationTy)
 218     return EHLinkRegistrationTy;
 219   LLVMContext &Context = TheModule->getContext();
 220   Type *FieldTys[] = {
 221       PointerType::getUnqual(Context), // EHRegistrationNode *Next
 222       PointerType::getUnqual(Context)  // EXCEPTION_DISPOSITION (*Handler)(...)
 223   };
 224   EHLinkRegistrationTy = StructType::create(FieldTys, "EHRegistrationNode");
 225   return EHLinkRegistrationTy;
 226 }
 227
 228 /// The __CxxFrameHandler3 registration node:
 229 ///   struct CXXExceptionRegistration {
 230 ///     void *SavedESP;
 231 ///     EHRegistrationNode SubRecord;
 232 ///     int32_t TryLevel;
 233 ///   };
 234 Type *WinEHStatePass::getCXXEHRegistrationType() {
 235   if (CXXEHRegistrationTy)
 236     return CXXEHRegistrationTy;
 237   LLVMContext &Context = TheModule->getContext();
 238   Type *FieldTys[] = {
 239       PointerType::getUnqual(Context), // void *SavedESP
 240       getEHLinkRegistrationType(),     // EHRegistrationNode SubRecord
 241       Type::getInt32Ty(Context)        // int32_t TryLevel
 242   };
 243   CXXEHRegistrationTy =
 244       StructType::create(FieldTys, "CXXExceptionRegistration");
 245   return CXXEHRegistrationTy;
 246 }
 247
 248 /// The _except_handler3/4 registration node:
 249 ///   struct EH4ExceptionRegistration {
 250 ///     void *SavedESP;
 251 ///     _EXCEPTION_POINTERS *ExceptionPointers;
 252 ///     EHRegistrationNode SubRecord;
 253 ///     int32_t EncodedScopeTable;
 254 ///     int32_t TryLevel;
 255 ///   };
 256 Type *WinEHStatePass::getSEHRegistrationType() {
 257   if (SEHRegistrationTy)
 258     return SEHRegistrationTy;
 259   LLVMContext &Context = TheModule->getContext();
 260   Type *FieldTys[] = {
 261       PointerType::getUnqual(Context), // void *SavedESP
 262       PointerType::getUnqual(Context), // void *ExceptionPointers
 263       getEHLinkRegistrationType(),     // EHRegistrationNode SubRecord
 264       Type::getInt32Ty(Context),       // int32_t EncodedScopeTable
 265       Type::getInt32Ty(Context)        // int32_t TryLevel
 266   };
 267   SEHRegistrationTy = StructType::create(FieldTys, "SEHExceptionRegistration");
 268   return SEHRegistrationTy;
 269 }
 270
 271 // Emit an exception registration record. These are stack allocations with the
 272 // common subobject of two pointers: the previous registration record (the old
 273 // fs:00) and the personality function for the current frame. The data before
 274 // and after that is personality function specific.
 275 void WinEHStatePass::emitExceptionRegistrationRecord(Function *F) {
 276   assert(Personality == EHPersonality::MSVC_CXX ||
 277          Personality == EHPersonality::MSVC_X86SEH);
 278
 279   IRBuilder<> Builder(&F->getEntryBlock(), F->getEntryBlock().begin());
 280   Type *Int8PtrType = Builder.getPtrTy();
 281   Type *Int32Ty = Builder.getInt32Ty();
 282   Type *VoidTy = Builder.getVoidTy();
 283
 284   if (Personality == EHPersonality::MSVC_CXX) {
 285     RegNodeTy = getCXXEHRegistrationType();
 286     RegNode = Builder.CreateAlloca(RegNodeTy);
 287     // SavedESP = llvm.stacksave()
 288     Value *SP = Builder.CreateStackSave();
 289     Builder.CreateStore(SP, Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
 290     // TryLevel = -1
 291     StateFieldIndex = 2;
 292     ParentBaseState = -1;
 293     insertStateNumberStore(&*Builder.GetInsertPoint(), ParentBaseState);
 294     // Handler = __ehhandler$F
 295     Function *Trampoline = generateLSDAInEAXThunk(F);
 296     Link = Builder.CreateStructGEP(RegNodeTy, RegNode, 1);
 297     linkExceptionRegistration(Builder, Trampoline);
 298
 299     CxxLongjmpUnwind = TheModule->getOrInsertFunction(
 300         "__CxxLongjmpUnwind",
 301         FunctionType::get(VoidTy, Int8PtrType, /*isVarArg=*/false));
 302     cast<Function>(CxxLongjmpUnwind.getCallee()->stripPointerCasts())
 303         ->setCallingConv(CallingConv::X86_StdCall);
 304   } else if (Personality == EHPersonality::MSVC_X86SEH) {
 305     // If _except_handler4 is in use, some additional guard checks and prologue
 306     // stuff is required.
 307     StringRef PersonalityName = PersonalityFn->getName();
 308     UseStackGuard = (PersonalityName == "_except_handler4");
 309
 310     // Allocate local structures.
 311     RegNodeTy = getSEHRegistrationType();
 312     RegNode = Builder.CreateAlloca(RegNodeTy);
 313     if (UseStackGuard)
 314       EHGuardNode = Builder.CreateAlloca(Int32Ty);
 315
 316     // SavedESP = llvm.stacksave()
 317     Value *SP = Builder.CreateStackSave();
 318     Builder.CreateStore(SP, Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
 319     // TryLevel = -2 / -1
 320     StateFieldIndex = 4;
 321     ParentBaseState = UseStackGuard ? -2 : -1;
 322     insertStateNumberStore(&*Builder.GetInsertPoint(), ParentBaseState);
 323     // ScopeTable = llvm.x86.seh.lsda(F)
 324     Value *LSDA = emitEHLSDA(Builder, F);
 325     LSDA = Builder.CreatePtrToInt(LSDA, Int32Ty);
 326     // If using _except_handler4, xor the address of the table with
 327     // __security_cookie.
 328     if (UseStackGuard) {
 329       Cookie = TheModule->getOrInsertGlobal("__security_cookie", Int32Ty);
 330       Value *Val = Builder.CreateLoad(Int32Ty, Cookie, "cookie");
 331       LSDA = Builder.CreateXor(LSDA, Val);
 332     }
 333     Builder.CreateStore(LSDA, Builder.CreateStructGEP(RegNodeTy, RegNode, 3));
 334
 335     // If using _except_handler4, the EHGuard contains: FramePtr xor Cookie.
 336     if (UseStackGuard) {
 337       Value *Val = Builder.CreateLoad(Int32Ty, Cookie);
 338       Value *FrameAddr = Builder.CreateIntrinsic(
 339           Intrinsic::frameaddress,
 340           Builder.getPtrTy(TheModule->getDataLayout().getAllocaAddrSpace()),
 341           Builder.getInt32(0), /*FMFSource=*/nullptr, "frameaddr");
 342       Value *FrameAddrI32 = Builder.CreatePtrToInt(FrameAddr, Int32Ty);
 343       FrameAddrI32 = Builder.CreateXor(FrameAddrI32, Val);
 344       Builder.CreateStore(FrameAddrI32, EHGuardNode);
 345     }
 346
 347     // Register the exception handler.
 348     Link = Builder.CreateStructGEP(RegNodeTy, RegNode, 2);
 349     linkExceptionRegistration(Builder, PersonalityFn);
 350
 351     SehLongjmpUnwind = TheModule->getOrInsertFunction(
 352         UseStackGuard ? "_seh_longjmp_unwind4" : "_seh_longjmp_unwind",
 353         FunctionType::get(Type::getVoidTy(TheModule->getContext()), Int8PtrType,
 354                           /*isVarArg=*/false));
 355     cast<Function>(SehLongjmpUnwind.getCallee()->stripPointerCasts())
 356         ->setCallingConv(CallingConv::X86_StdCall);
 357   } else {
 358     llvm_unreachable("unexpected personality function");
 359   }
 360
 361   // Insert an unlink before all returns.
 362   for (BasicBlock &BB : *F) {
 363     Instruction *T = BB.getTerminator();
 364     if (!isa<ReturnInst>(T))
 365       continue;
 366
 367     // If there is a musttail call, that's the de-facto terminator.
 368     if (CallInst *CI = BB.getTerminatingMustTailCall())
 369       T = CI;
 370
 371     Builder.SetInsertPoint(T);
 372     unlinkExceptionRegistration(Builder);
 373   }
 374 }
 375
 376 Value *WinEHStatePass::emitEHLSDA(IRBuilder<> &Builder, Function *F) {
 377   return Builder.CreateIntrinsic(Intrinsic::x86_seh_lsda, {}, F);
 378 }
 379
 380 /// Generate a thunk that puts the LSDA of ParentFunc in EAX and then calls
 381 /// PersonalityFn, forwarding the parameters passed to PEXCEPTION_ROUTINE:
 382 ///   typedef _EXCEPTION_DISPOSITION (*PEXCEPTION_ROUTINE)(
 383 ///       _EXCEPTION_RECORD *, void *, _CONTEXT *, void *);
 384 /// We essentially want this code:
 385 ///   movl $lsda, %eax
 386 ///   jmpl ___CxxFrameHandler3
 387 Function *WinEHStatePass::generateLSDAInEAXThunk(Function *ParentFunc) {
 388   LLVMContext &Context = ParentFunc->getContext();
 389   Type *Int32Ty = Type::getInt32Ty(Context);
 390   Type *Int8PtrType = PointerType::getUnqual(Context);
 391   Type *ArgTys[5] = {Int8PtrType, Int8PtrType, Int8PtrType, Int8PtrType,
 392                      Int8PtrType};
 393   FunctionType *TrampolineTy =
 394       FunctionType::get(Int32Ty, ArrayRef(&ArgTys[0], 4),
 395                         /*isVarArg=*/false);
 396   FunctionType *TargetFuncTy =
 397       FunctionType::get(Int32Ty, ArrayRef(&ArgTys[0], 5),
 398                         /*isVarArg=*/false);
 399   Function *Trampoline = Function::Create(
 400       TrampolineTy, GlobalValue::InternalLinkage,
 401       Twine("__ehhandler$") +
 402           GlobalValue::dropLLVMManglingEscape(ParentFunc->getName()),
 403       TheModule);
 404   if (auto *C = ParentFunc->getComdat())
 405     Trampoline->setComdat(C);
 406   BasicBlock *EntryBB = BasicBlock::Create(Context, "entry", Trampoline);
 407   IRBuilder<> Builder(EntryBB);
 408   Value *LSDA = emitEHLSDA(Builder, ParentFunc);
 409   auto AI = Trampoline->arg_begin();
 410   Value *Args[5] = {LSDA, &*AI++, &*AI++, &*AI++, &*AI++};
 411   CallInst *Call = Builder.CreateCall(TargetFuncTy, PersonalityFn, Args);
 412   // Can't use musttail due to prototype mismatch, but we can use tail.
 413   Call->setTailCall(true);
 414   // Set inreg so we pass it in EAX.
 415   Call->addParamAttr(0, Attribute::InReg);
 416   Builder.CreateRet(Call);
 417   return Trampoline;
 418 }
 419
 420 void WinEHStatePass::linkExceptionRegistration(IRBuilder<> &Builder,
 421                                                Function *Handler) {
 422   // Emit the .safeseh directive for this function.
 423   Handler->addFnAttr("safeseh");
 424
 425   LLVMContext &C = Builder.getContext();
 426   Type *LinkTy = getEHLinkRegistrationType();
 427   // Handler = Handler
 428   Builder.CreateStore(Handler, Builder.CreateStructGEP(LinkTy, Link, 1));
 429   // Next = [fs:00]
 430   Constant *FSZero = Constant::getNullValue(PointerType::get(C, X86AS::FS));
 431   Value *Next = Builder.CreateLoad(PointerType::getUnqual(C), FSZero);
 432   Builder.CreateStore(Next, Builder.CreateStructGEP(LinkTy, Link, 0));
 433   // [fs:00] = Link
 434   Builder.CreateStore(Link, FSZero);
 435 }
 436
 437 void WinEHStatePass::unlinkExceptionRegistration(IRBuilder<> &Builder) {
 438   // Clone Link into the current BB for better address mode folding.
 439   if (auto *GEP = dyn_cast<GetElementPtrInst>(Link)) {
 440     GEP = cast<GetElementPtrInst>(GEP->clone());
 441     Builder.Insert(GEP);
 442     Link = GEP;
 443   }
 444
 445   LLVMContext &C = Builder.getContext();
 446   Type *LinkTy = getEHLinkRegistrationType();
 447   // [fs:00] = Link->Next
 448   Value *Next = Builder.CreateLoad(PointerType::getUnqual(C),
 449                                    Builder.CreateStructGEP(LinkTy, Link, 0));
 450   Constant *FSZero = Constant::getNullValue(PointerType::get(C, X86AS::FS));
 451   Builder.CreateStore(Next, FSZero);
 452 }
 453
 454 // Calls to setjmp(p) are lowered to _setjmp3(p, 0) by the frontend.
 455 // The idea behind _setjmp3 is that it takes an optional number of personality
 456 // specific parameters to indicate how to restore the personality-specific frame
 457 // state when longjmp is initiated.  Typically, the current TryLevel is saved.
 458 void WinEHStatePass::rewriteSetJmpCall(IRBuilder<> &Builder, Function &F,
 459                                        CallBase &Call, Value *State) {
 460   // Don't rewrite calls with a weird number of arguments.
 461   if (Call.arg_size() != 2)
 462     return;
 463
 464   SmallVector<OperandBundleDef, 1> OpBundles;
 465   Call.getOperandBundlesAsDefs(OpBundles);
 466
 467   SmallVector<Value *, 3> OptionalArgs;
 468   if (Personality == EHPersonality::MSVC_CXX) {
 469     OptionalArgs.push_back(CxxLongjmpUnwind.getCallee());
 470     OptionalArgs.push_back(State);
 471     OptionalArgs.push_back(emitEHLSDA(Builder, &F));
 472   } else if (Personality == EHPersonality::MSVC_X86SEH) {
 473     OptionalArgs.push_back(SehLongjmpUnwind.getCallee());
 474     OptionalArgs.push_back(State);
 475     if (UseStackGuard)
 476       OptionalArgs.push_back(Cookie);
 477   } else {
 478     llvm_unreachable("unhandled personality!");
 479   }
 480
 481   SmallVector<Value *, 5> Args;
 482   Args.push_back(
 483       Builder.CreateBitCast(Call.getArgOperand(0), Builder.getPtrTy()));
 484   Args.push_back(Builder.getInt32(OptionalArgs.size()));
 485   Args.append(OptionalArgs.begin(), OptionalArgs.end());
 486
 487   CallBase *NewCall;
 488   if (auto *CI = dyn_cast<CallInst>(&Call)) {
 489     CallInst *NewCI = Builder.CreateCall(SetJmp3, Args, OpBundles);
 490     NewCI->setTailCallKind(CI->getTailCallKind());
 491     NewCall = NewCI;
 492   } else {
 493     auto *II = cast<InvokeInst>(&Call);
 494     NewCall = Builder.CreateInvoke(SetJmp3, II->getNormalDest(),
 495                                    II->getUnwindDest(), Args, OpBundles);
 496   }
 497   NewCall->setCallingConv(Call.getCallingConv());
 498   NewCall->setAttributes(Call.getAttributes());
 499   NewCall->setDebugLoc(Call.getDebugLoc());
 500
 501   NewCall->takeName(&Call);
 502   Call.replaceAllUsesWith(NewCall);
 503   Call.eraseFromParent();
 504 }
 505
 506 // Figure out what state we should assign calls in this block.
 507 int WinEHStatePass::getBaseStateForBB(
 508     DenseMap<BasicBlock *, ColorVector> &BlockColors, WinEHFuncInfo &FuncInfo,
 509     BasicBlock *BB) {
 510   int BaseState = ParentBaseState;
 511   auto &BBColors = BlockColors[BB];
 512
 513   assert(BBColors.size() == 1 && "multi-color BB not removed by preparation");
 514   BasicBlock *FuncletEntryBB = BBColors.front();
 515   if (auto *FuncletPad =
 516           dyn_cast<FuncletPadInst>(FuncletEntryBB->getFirstNonPHI())) {
 517     auto BaseStateI = FuncInfo.FuncletBaseStateMap.find(FuncletPad);
 518     if (BaseStateI != FuncInfo.FuncletBaseStateMap.end())
 519       BaseState = BaseStateI->second;
 520   }
 521
 522   return BaseState;
 523 }
 524
 525 static bool isIntrinsic(const CallBase &Call, Intrinsic::ID ID) {
 526   const Function *CF = Call.getCalledFunction();
 527   return CF && CF->isIntrinsic() && CF->getIntrinsicID() == ID;
 528 }
 529
 530 static bool isSehScopeEnd(const CallBase &Call) {
 531   return isIntrinsic(Call, Intrinsic::seh_scope_end);
 532 }
 533
 534 static bool isSehScopeBegin(const CallBase &Call) {
 535   return isIntrinsic(Call, Intrinsic::seh_scope_begin);
 536 }
 537
 538 // Calculate the state a call-site is in.
 539 int WinEHStatePass::getStateForCall(
 540     DenseMap<BasicBlock *, ColorVector> &BlockColors, WinEHFuncInfo &FuncInfo,
 541     CallBase &Call) {
 542   if (auto *II = dyn_cast<InvokeInst>(&Call)) {
 543     if (isSehScopeEnd(*II)) {
 544       return getBaseStateForBB(BlockColors, FuncInfo, II->getNormalDest());
 545     }
 546
 547     // Look up the state number of the EH pad this unwinds to.
 548     assert(FuncInfo.InvokeStateMap.count(II) && "invoke has no state!");
 549     return FuncInfo.InvokeStateMap[II];
 550   }
 551   // Possibly throwing call instructions have no actions to take after
 552   // an unwind. Ensure they are in the -1 state.
 553   return getBaseStateForBB(BlockColors, FuncInfo, Call.getParent());
 554 }
 555
 556 // Calculate the intersection of all the FinalStates for a BasicBlock's
 557 // predecessors.
 558 static int getPredState(DenseMap<BasicBlock *, int> &FinalStates, Function &F,
 559                         int ParentBaseState, BasicBlock *BB) {
 560   // The entry block has no predecessors but we know that the prologue always
 561   // sets us up with a fixed state.
 562   if (&F.getEntryBlock() == BB)
 563     return ParentBaseState;
 564
 565   // This is an EH Pad, conservatively report this basic block as overdefined.
 566   if (BB->isEHPad())
 567     return OverdefinedState;
 568
 569   int CommonState = OverdefinedState;
 570   for (BasicBlock *PredBB : predecessors(BB)) {
 571     // We didn't manage to get a state for one of these predecessors,
 572     // conservatively report this basic block as overdefined.
 573     auto PredEndState = FinalStates.find(PredBB);
 574     if (PredEndState == FinalStates.end())
 575       return OverdefinedState;
 576
 577     // This code is reachable via exceptional control flow,
 578     // conservatively report this basic block as overdefined.
 579     if (isa<CatchReturnInst>(PredBB->getTerminator()))
 580       return OverdefinedState;
 581
 582     int PredState = PredEndState->second;
 583     assert(PredState != OverdefinedState &&
 584            "overdefined BBs shouldn't be in FinalStates");
 585     if (CommonState == OverdefinedState)
 586       CommonState = PredState;
 587
 588     // At least two predecessors have different FinalStates,
 589     // conservatively report this basic block as overdefined.
 590     if (CommonState != PredState)
 591       return OverdefinedState;
 592   }
 593
 594   return CommonState;
 595 }
 596
 597 // Calculate the intersection of all the InitialStates for a BasicBlock's
 598 // successors.
 599 static int getSuccState(DenseMap<BasicBlock *, int> &InitialStates, Function &F,
 600                         int ParentBaseState, BasicBlock *BB) {
 601   // This block rejoins normal control flow,
 602   // conservatively report this basic block as overdefined.
 603   if (isa<CatchReturnInst>(BB->getTerminator()))
 604     return OverdefinedState;
 605
 606   int CommonState = OverdefinedState;
 607   for (BasicBlock *SuccBB : successors(BB)) {
 608     // We didn't manage to get a state for one of these predecessors,
 609     // conservatively report this basic block as overdefined.
 610     auto SuccStartState = InitialStates.find(SuccBB);
 611     if (SuccStartState == InitialStates.end())
 612       return OverdefinedState;
 613
 614     // This is an EH Pad, conservatively report this basic block as overdefined.
 615     if (SuccBB->isEHPad())
 616       return OverdefinedState;
 617
 618     int SuccState = SuccStartState->second;
 619     assert(SuccState != OverdefinedState &&
 620            "overdefined BBs shouldn't be in FinalStates");
 621     if (CommonState == OverdefinedState)
 622       CommonState = SuccState;
 623
 624     // At least two successors have different InitialStates,
 625     // conservatively report this basic block as overdefined.
 626     if (CommonState != SuccState)
 627       return OverdefinedState;
 628   }
 629
 630   return CommonState;
 631 }
 632
 633 bool WinEHStatePass::isStateStoreNeeded(EHPersonality Personality,
 634                                         CallBase &Call) {
 635   if (isSehScopeBegin(Call) || isSehScopeEnd(Call)) {
 636     return true;
 637   }
 638
 639   // If the function touches memory, it needs a state store.
 640   if (isAsynchronousEHPersonality(Personality))
 641     return !Call.doesNotAccessMemory();
 642
 643   // If the function throws, it needs a state store.
 644   return !Call.doesNotThrow();
 645 }
 646
 647 void WinEHStatePass::addStateStores(Function &F, WinEHFuncInfo &FuncInfo) {
 648   // Mark the registration node. The backend needs to know which alloca it is so
 649   // that it can recover the original frame pointer.
 650   IRBuilder<> Builder(RegNode->getNextNode());
 651   Value *RegNodeI8 = Builder.CreateBitCast(RegNode, Builder.getPtrTy());
 652   Builder.CreateIntrinsic(Intrinsic::x86_seh_ehregnode, {}, {RegNodeI8});
 653
 654   if (EHGuardNode) {
 655     IRBuilder<> Builder(EHGuardNode->getNextNode());
 656     Value *EHGuardNodeI8 =
 657         Builder.CreateBitCast(EHGuardNode, Builder.getPtrTy());
 658     Builder.CreateIntrinsic(Intrinsic::x86_seh_ehguard, {}, {EHGuardNodeI8});
 659   }
 660
 661   // Calculate state numbers.
 662   if (isAsynchronousEHPersonality(Personality))
 663     calculateSEHStateNumbers(&F, FuncInfo);
 664   else
 665     calculateWinCXXEHStateNumbers(&F, FuncInfo);
 666
 667   // Iterate all the instructions and emit state number stores.
 668   DenseMap<BasicBlock *, ColorVector> BlockColors = colorEHFunclets(F);
 669   ReversePostOrderTraversal<Function *> RPOT(&F);
 670
 671   // InitialStates yields the state of the first call-site for a BasicBlock.
 672   DenseMap<BasicBlock *, int> InitialStates;
 673   // FinalStates yields the state of the last call-site for a BasicBlock.
 674   DenseMap<BasicBlock *, int> FinalStates;
 675   // Worklist used to revisit BasicBlocks with indeterminate
 676   // Initial/Final-States.
 677   std::deque<BasicBlock *> Worklist;
 678   // Fill in InitialStates and FinalStates for BasicBlocks with call-sites.
 679   for (BasicBlock *BB : RPOT) {
 680     int InitialState = OverdefinedState;
 681     int FinalState;
 682     if (&F.getEntryBlock() == BB)
 683       InitialState = FinalState = ParentBaseState;
 684     for (Instruction &I : *BB) {
 685       auto *Call = dyn_cast<CallBase>(&I);
 686       if (!Call || !isStateStoreNeeded(Personality, *Call))
 687         continue;
 688
 689       int State = getStateForCall(BlockColors, FuncInfo, *Call);
 690       if (InitialState == OverdefinedState)
 691         InitialState = State;
 692       FinalState = State;
 693     }
 694     // No call-sites in this basic block? That's OK, we will come back to these
 695     // in a later pass.
 696     if (InitialState == OverdefinedState) {
 697       Worklist.push_back(BB);
 698       continue;
 699     }
 700     LLVM_DEBUG(dbgs() << "X86WinEHState: " << BB->getName()
 701                       << " InitialState=" << InitialState << '\n');
 702     LLVM_DEBUG(dbgs() << "X86WinEHState: " << BB->getName()
 703                       << " FinalState=" << FinalState << '\n');
 704     InitialStates.insert({BB, InitialState});
 705     FinalStates.insert({BB, FinalState});
 706   }
 707
 708   // Try to fill-in InitialStates and FinalStates which have no call-sites.
 709   while (!Worklist.empty()) {
 710     BasicBlock *BB = Worklist.front();
 711     Worklist.pop_front();
 712     // This BasicBlock has already been figured out, nothing more we can do.
 713     if (InitialStates.count(BB) != 0)
 714       continue;
 715
 716     int PredState = getPredState(FinalStates, F, ParentBaseState, BB);
 717     if (PredState == OverdefinedState)
 718       continue;
 719
 720     // We successfully inferred this BasicBlock's state via it's predecessors;
 721     // enqueue it's successors to see if we can infer their states.
 722     InitialStates.insert({BB, PredState});
 723     FinalStates.insert({BB, PredState});
 724     for (BasicBlock *SuccBB : successors(BB))
 725       Worklist.push_back(SuccBB);
 726   }
 727
 728   // Try to hoist stores from successors.
 729   for (BasicBlock *BB : RPOT) {
 730     int SuccState = getSuccState(InitialStates, F, ParentBaseState, BB);
 731     if (SuccState == OverdefinedState)
 732       continue;
 733
 734     // Update our FinalState to reflect the common InitialState of our
 735     // successors.
 736     FinalStates.insert({BB, SuccState});
 737   }
 738
 739   // Finally, insert state stores before call-sites which transition us to a new
 740   // state.
 741   for (BasicBlock *BB : RPOT) {
 742     auto &BBColors = BlockColors[BB];
 743     BasicBlock *FuncletEntryBB = BBColors.front();
 744     if (isa<CleanupPadInst>(FuncletEntryBB->getFirstNonPHI()))
 745       continue;
 746
 747     int PrevState = getPredState(FinalStates, F, ParentBaseState, BB);
 748     LLVM_DEBUG(dbgs() << "X86WinEHState: " << BB->getName()
 749                       << " PrevState=" << PrevState << '\n');
 750
 751     for (Instruction &I : *BB) {
 752       auto *Call = dyn_cast<CallBase>(&I);
 753       if (!Call || !isStateStoreNeeded(Personality, *Call))
 754         continue;
 755
 756       int State = getStateForCall(BlockColors, FuncInfo, *Call);
 757       if (State != PrevState)
 758         insertStateNumberStore(&I, State);
 759       PrevState = State;
 760     }
 761
 762     // We might have hoisted a state store into this block, emit it now.
 763     auto EndState = FinalStates.find(BB);
 764     if (EndState != FinalStates.end())
 765       if (EndState->second != PrevState)
 766         insertStateNumberStore(BB->getTerminator(), EndState->second);
 767   }
 768
 769   SmallVector<CallBase *, 1> SetJmp3Calls;
 770   for (BasicBlock *BB : RPOT) {
 771     for (Instruction &I : *BB) {
 772       auto *Call = dyn_cast<CallBase>(&I);
 773       if (!Call)
 774         continue;
 775       if (Call->getCalledOperand()->stripPointerCasts() !=
 776           SetJmp3.getCallee()->stripPointerCasts())
 777         continue;
 778
 779       SetJmp3Calls.push_back(Call);
 780     }
 781   }
 782
 783   for (CallBase *Call : SetJmp3Calls) {
 784     auto &BBColors = BlockColors[Call->getParent()];
 785     BasicBlock *FuncletEntryBB = BBColors.front();
 786     bool InCleanup = isa<CleanupPadInst>(FuncletEntryBB->getFirstNonPHI());
 787
 788     IRBuilder<> Builder(Call);
 789     Value *State;
 790     if (InCleanup) {
 791       Value *StateField = Builder.CreateStructGEP(RegNode->getAllocatedType(),
 792                                                   RegNode, StateFieldIndex);
 793       State = Builder.CreateLoad(Builder.getInt32Ty(), StateField);
 794     } else {
 795       State = Builder.getInt32(getStateForCall(BlockColors, FuncInfo, *Call));
 796     }
 797     rewriteSetJmpCall(Builder, F, *Call, State);
 798   }
 799 }
 800
 801 void WinEHStatePass::insertStateNumberStore(Instruction *IP, int State) {
 802   IRBuilder<> Builder(IP);
 803   Value *StateField = Builder.CreateStructGEP(RegNode->getAllocatedType(),
 804                                               RegNode, StateFieldIndex);
 805   Builder.CreateStore(Builder.getInt32(State), StateField);
 806 }
 807
 808 void WinEHStatePass::updateEspForInAllocas(Function &F) {
 809   for (BasicBlock &BB : F) {
 810     for (Instruction &I : BB) {
 811       if (auto *Alloca = dyn_cast<AllocaInst>(&I)) {
 812         if (Alloca->isStaticAlloca())
 813           continue;
 814         IRBuilder<> Builder(Alloca->getNextNonDebugInstruction());
 815         // SavedESP = llvm.stacksave()
 816         Value *SP = Builder.CreateStackSave();
 817         Builder.CreateStore(SP, Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
 818       }
 819
 820       if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
 821         if (II->getIntrinsicID() != Intrinsic::stackrestore)
 822           continue;
 823         IRBuilder<> Builder(II->getNextNonDebugInstruction());
 824         // SavedESP = llvm.stacksave()
 825         Value *SP = Builder.CreateStackSave();
 826         Builder.CreateStore(SP, Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
 827       }
 828     }
 829   }
 830 }