Insert state restore before SEH scope end

llvm/lib/Target/X86/X86WinEHState.cpp

@@ -513,11 +513,28 @@ int WinEHStatePass::getBaseStateForBB(
   return BaseState;
 }
 
+static bool isIntrinsic(const CallBase &Call, Intrinsic::ID ID) {
+  const Function *CF = Call.getCalledFunction();
+  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);
+}
+
 // Calculate the state a call-site is in.
 int WinEHStatePass::getStateForCall(
     DenseMap<BasicBlock *, ColorVector> &BlockColors, WinEHFuncInfo &FuncInfo,
     CallBase &Call) {
   if (auto *II = dyn_cast<InvokeInst>(&Call)) {
+    if (isSehScopeEnd(*II)) {
+      return getBaseStateForBB(BlockColors, FuncInfo, II->getNormalDest());
+    }
+
     // 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];
@@ -604,22 +621,9 @@ static int getSuccState(DenseMap<BasicBlock *, int> &InitialStates, Function &F,
   return CommonState;
 }
 
-static bool isIntrinsic(const CallBase &Call, Intrinsic::ID ID) {
-  const Function *CF = Call.getCalledFunction();
-  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,
                                         CallBase &Call) {
-  if (isSehScopeBegin(Call)) {
+  if (isSehScopeBegin(Call) || isSehScopeEnd(Call)) {
     return true;
   }
 
@@ -659,8 +663,6 @@ 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;
@@ -675,11 +677,6 @@ void WinEHStatePass::addStateStores(Function &F, WinEHFuncInfo &FuncInfo) {
       if (!Call)
         continue;
 
-      if (isSehScopeEnd(*Call)) {
-        auto *Invoke = cast<InvokeInst>(Call);
-        ScopeEndBlocks.insert(Invoke->getNormalDest());
-      }
-
       if (!isStateStoreNeeded(Personality, *Call))
         continue;
 
@@ -733,12 +730,6 @@ 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) {

llvm/test/CodeGen/WinEH/wineh-scope-statenumbering.ll

@@ -19,13 +19,12 @@ entry:
 
 invoke.cont:
   store i32 1, ptr inttoptr (i32 1 to ptr), align 4
+  ; CHECK:  store i32 -1, ptr %10, align 4
+  ; CHECK-NEXT:  invoke void @llvm.seh.scope.end()
   invoke void @llvm.seh.scope.end()
           to label %invoke.cont1 unwind label %ehcleanup
 
 invoke.cont1:
-  ; CHECK: invoke.cont1:
-  ; CHECK-NEXT:%10 = getelementptr inbounds nuw %CXXExceptionRegistration, ptr %0, i32 0, i32 2
-  ; CHECK-NEXT:  store i32 -1, ptr %10, align 4
   call x86_thiscallcc void @"??1Destructor@@QAE@XZ"(ptr noundef nonnull align 4 dereferenceable(4) %x) #1
   ret void