[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:

Sure, but I think in the pariticular case of the sample in my issue (https://godbolt.org/z/8Gbqd6Gha) where the try ends on a return, I can at least see why there is no seh.try.end emitted.
Because where could you place it? 
The return still has to be inside of the try (because in the case of async eh we are instruction precise i think and a ret could still cause a segfault - which can be handled by the catch) and emitting it after would end up in a bogus control flow / killing the terminator.

This also brings up another issue i see in my sample:
An add instruction is still being hoisted out of the try range - even if you apply my fix:
https://godbolt.org/z/nnTW1vYz6

I think this could be a bug in machine-sink, which probably isn't informed that it cannot move code out of async-eh ranges.


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