[WinEH] Restore states after SEH scope end

Author: momo5502

llvm/lib/Target/X86/X86WinEHState.cpp

@@ -608,10 +608,17 @@ static int getSuccState(DenseMap<BasicBlock *, int> &InitialStates, Function &F,
   return CommonState;
 }
 
-static bool isSehScopeBegin(const CallBase &Call) {
+static bool isIntrinsic(const CallBase &Call, Intrinsic::ID ID) {
   const Function *CF = Call.getCalledFunction();
-  return CF && CF->isIntrinsic() &&
-         CF->getIntrinsicID() == Intrinsic::seh_scope_begin;
+  return CF && CF->isIntrinsic() && CF->getIntrinsicID() == ID;
+}
+
+static bool isSehScopeEnd(const CallBase &Call) {
+  return isIntrinsic(Call, Intrinsic::seh_scope_end);
+}
+
+static bool isSehScopeBegin(const CallBase &Call) {
+  return isIntrinsic(Call, Intrinsic::seh_scope_begin);
 }
 
 bool WinEHStatePass::isStateStoreNeeded(EHPersonality Personality,
@@ -656,6 +663,8 @@ void WinEHStatePass::addStateStores(Function &F, WinEHFuncInfo &FuncInfo) {
   DenseMap<BasicBlock *, int> InitialStates;
   // FinalStates yields the state of the last call-site for a BasicBlock.
   DenseMap<BasicBlock *, int> FinalStates;
+  // SEH scope end target blocks
+  SmallPtrSet<BasicBlock *, 4> ScopeEndBlocks;
   // Worklist used to revisit BasicBlocks with indeterminate
   // Initial/Final-States.
   std::deque<BasicBlock *> Worklist;
@@ -667,7 +676,15 @@ void WinEHStatePass::addStateStores(Function &F, WinEHFuncInfo &FuncInfo) {
       InitialState = FinalState = ParentBaseState;
     for (Instruction &I : *BB) {
       auto *Call = dyn_cast<CallBase>(&I);
-      if (!Call || !isStateStoreNeeded(Personality, *Call))
+      if (!Call)
+        continue;
+
+      if (isSehScopeEnd(*Call)) {
+        auto *Invoke = cast<InvokeInst>(Call);
+        ScopeEndBlocks.insert(Invoke->getNormalDest());
+      }
+
+      if (!isStateStoreNeeded(Personality, *Call))
         continue;
 
       int State = getStateForCall(BlockColors, FuncInfo, *Call);
@@ -720,6 +737,12 @@ void WinEHStatePass::addStateStores(Function &F, WinEHFuncInfo &FuncInfo) {
     FinalStates.insert({BB, SuccState});
   }
 
+  // Insert state restores after SEH scope ends
+  for (BasicBlock *BB : ScopeEndBlocks) {
+    int state = getBaseStateForBB(BlockColors, FuncInfo, BB);
+    insertStateNumberStore(BB->getFirstNonPHI(), state);
+  }
+
   // Finally, insert state stores before call-sites which transition us to a new
   // state.
   for (BasicBlock *BB : RPOT) {