[llvm] [Inline][WinEH] Fix try scopes leaking to caller on inline (PR #164170)

[via llvm-commits]

================
@@ -2210,6 +2211,135 @@ inlineRetainOrClaimRVCalls(CallBase &CB, objcarc::ARCInstKind RVCallKind,
   }
 }
 
+// Determine SEH try scopes and order them by dominance.
+static SmallVector<llvm::InvokeInst *, 1>
+GetOrderedSehTryScopes(Function *Func, DominatorTree &DT) {
+  SmallVector<llvm::InvokeInst *, 1> Scopes{};
+  bool DTCalculated = false;
+
+  for (auto &BB : *Func) {
+    auto *TI = BB.getTerminator();
+    if (auto *II = dyn_cast<InvokeInst>(TI)) {
+      auto *Call = cast<CallBase>(II);
+      const auto *Fn = Call->getCalledFunction();
+      if (!Fn || !Fn->isIntrinsic())
+        continue;
+
+      if (Fn->getIntrinsicID() != Intrinsic::seh_try_begin)
+        continue;
+
+      if (!DTCalculated) {
+        DT.recalculate(*Func);
+        DTCalculated = true;
+      }
+
+      auto InsertIt = Scopes.end();
+      if (!Scopes.empty())
+        for (auto ScopeIt = Scopes.begin(); ScopeIt != Scopes.end(); ++ScopeIt)
+          if (DT.dominates((*ScopeIt)->getParent(), &BB))
+            InsertIt = std::next(ScopeIt);
+
+      Scopes.insert(InsertIt, II);
+    }
+  }
+
+  return Scopes;
+}
+
+// Find, if present, the outermost unterminated try scope for the input block.
+static llvm::InvokeInst *GetOutermostUnterminatedTryScopeBegin(
+    llvm::BasicBlock *ReturnBlock, SmallVector<llvm::InvokeInst *, 1> &Scopes,
+    DominatorTree &DT) {
+  llvm::InvokeInst *UnterminatedScope{nullptr};
+
+  for (auto ScopeIt = Scopes.rbegin(); ScopeIt != Scopes.rend(); ++ScopeIt) {
+    auto *Invoke = *ScopeIt;
+
+    // Return might not be in a try scope.
+    if (!DT.dominates(Invoke->getParent(), ReturnBlock))
+      continue;
+
+    // If there is a catch which connects to the return, the try scope is
+    // considered to be terminated.
+    if (isPotentiallyReachable(Invoke->getUnwindDest(), ReturnBlock, nullptr,
+                               &DT))
+      continue;
+
+    UnterminatedScope = Invoke;
+  }
+
+  return UnterminatedScope;
+}
+
+struct UnterminatedTryScope {
+  llvm::BasicBlock *ReturnBlock;
+  llvm::BasicBlock *TryStart;
+  llvm::BasicBlock *UnwindDestination;
+};
+
+// For Windows -EHa there may be the case of try scopes spanning over a return
+// instruction. If no other return out of the function is given (e.g. by letting
+// all other blocks terminate with unreachable) the try scope is considered
+// unterminated upon inlining. To avoid leaking the try scopes over to a caller
+// we need to add a terminator for the outermost unterminated try scope.
+static SmallVector<UnterminatedTryScope, 1>
+GetUnterminatedTryScopes(Function *CalledFunc) {
+  SmallVector<UnterminatedTryScope, 1> UnterminatedScopes;
+  DominatorTree DT;
+  auto Scopes = GetOrderedSehTryScopes(CalledFunc, DT);
+  if (Scopes.empty())
+    return UnterminatedScopes;
+
+  SmallVector<ReturnInst *, 8> Returns;
+  for (auto &BB : *CalledFunc)
+    if (ReturnInst *RI = dyn_cast<ReturnInst>(BB.getTerminator()))
+      Returns.push_back(RI);
+
+  for (auto *RI : Returns) {
+    auto *Block = RI->getParent();
+    auto *OutermostScope =
+        GetOutermostUnterminatedTryScopeBegin(Block, Scopes, DT);
+    if (!OutermostScope)
+      continue;
+
+    UnterminatedScopes.push_back(
+        {Block, OutermostScope->getParent(), OutermostScope->getUnwindDest()});
+  }
+
+  return UnterminatedScopes;
+}
+
+// Insert terminator for unterminated try scopes.
+static void HandleUnterminatedTryScopes(
+    Function *Caller,
+    SmallVector<UnterminatedTryScope, 1> &UnterminatedTryScopes,
+    ValueToValueMapTy &VMap) {
+  for (auto &Scope : UnterminatedTryScopes) {
+    auto *ReturnBlock = cast<BasicBlock>(VMap[Scope.ReturnBlock]);
+    auto *TryStart = cast<BasicBlock>(VMap[Scope.TryStart]);
+    auto *UnwindDestination = cast<BasicBlock>(VMap[Scope.UnwindDestination]);
+    // did not survive (partial) inlining - ignore
+    if (!ReturnBlock || !TryStart || !UnwindDestination)
+      continue;
+
+    auto *Mod = (llvm::Module *)Caller->getParent();
+    auto *SehTryEndFn =
+        Intrinsic::getOrInsertDeclaration(Mod, Intrinsic::seh_try_end);
+    auto *BB = ReturnBlock->splitBasicBlockBefore(ReturnBlock->getTerminator(),
+                                                  "try.end");
+    BB->getTerminator()->eraseFromParent();
+    IRBuilder<>(BB).CreateInvoke(SehTryEndFn, ReturnBlock, UnwindDestination);
----------------
MuellerMP wrote:

I guess this is a bit of an exaggerated example but there exists the possibility to catch a segfault on a ret right here:
https://godbolt.org/z/s13aqhfro

Now i believe this example should also hold for mutlithreaded cases, where Thread A changes the memory protection of the ret instruction while Thread B will execute it next - thus causing unwinding (catchable with EHa - as long as the stack is still valid).

https://github.com/llvm/llvm-project/pull/164170