[llvm] [WinEH] Emit state stores before SEH scopes (PR #116546)

[Maurice Heumann via llvm-commits]

https://github.com/momo5502 updated https://github.com/llvm/llvm-project/pull/116546

>From 6ede329985d56c3a66f9ea4042c714849a00cf41 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann at gmail.com>
Date: Sun, 17 Nov 2024 15:52:01 +0100
Subject: [PATCH 1/2] [WinEH] Emit state stores before SEH scopes

---
 llvm/lib/Target/X86/X86WinEHState.cpp | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/llvm/lib/Target/X86/X86WinEHState.cpp b/llvm/lib/Target/X86/X86WinEHState.cpp
index ef212736730114..e6aad3fa2a4fb1 100644
--- a/llvm/lib/Target/X86/X86WinEHState.cpp
+++ b/llvm/lib/Target/X86/X86WinEHState.cpp
@@ -604,8 +604,18 @@ static int getSuccState(DenseMap<BasicBlock *, int> &InitialStates, Function &F,
   return CommonState;
 }
 
+static bool isSehScopeBegin(const CallBase &Call) {
+  const Function *CF = Call.getCalledFunction();
+  return CF && CF->isIntrinsic() &&
+         CF->getIntrinsicID() == Intrinsic::seh_scope_begin;
+}
+
 bool WinEHStatePass::isStateStoreNeeded(EHPersonality Personality,
                                         CallBase &Call) {
+  if (isSehScopeBegin(Call)) {
+    return true;
+  }
+
   // If the function touches memory, it needs a state store.
   if (isAsynchronousEHPersonality(Personality))
     return !Call.doesNotAccessMemory();

>From 282680b4a789298fdcc2dfe693a2a3c47ab7fe7b Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann at gmail.com>
Date: Sun, 17 Nov 2024 17:23:50 +0100
Subject: [PATCH 2/2] [WinEH] Restore states after SEH scope end

---
 llvm/lib/Target/X86/X86WinEHState.cpp | 31 +++++++++++++++++++++++----
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/X86/X86WinEHState.cpp b/llvm/lib/Target/X86/X86WinEHState.cpp
index e6aad3fa2a4fb1..d5bcbd2a08217d 100644
--- a/llvm/lib/Target/X86/X86WinEHState.cpp
+++ b/llvm/lib/Target/X86/X86WinEHState.cpp
@@ -604,10 +604,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,
@@ -652,6 +659,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;
@@ -663,7 +672,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);
@@ -716,6 +733,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) {