3 files changed, 9 insertions, 23 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
index b62bd2bd63e..be5a7034be4 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
@@ -271,6 +271,8 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
     }
   }
 
+  WinEHFuncInfo &EHInfo = *MF->getWinEHFuncInfo();
+
   // Mark landing pad blocks.
   SmallVector<const LandingPadInst *, 4> LPads;
   for (BB = Fn->begin(); BB != EB; ++BB) {
@@ -289,7 +291,6 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
     return;
 
   // Calculate state numbers if we haven't already.
-  WinEHFuncInfo &EHInfo = *MF->getWinEHFuncInfo();
   if (Personality == EHPersonality::MSVC_CXX)
     calculateWinCXXEHStateNumbers(&fn, EHInfo);
   else if (isAsynchronousEHPersonality(Personality))
@@ -297,8 +298,6 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
   else if (Personality == EHPersonality::CoreCLR)
     calculateClrEHStateNumbers(&fn, EHInfo);
 
-  calculateCatchReturnSuccessorColors(&fn, EHInfo);
-
   // Map all BB references in the WinEH data to MBBs.
   for (WinEHTryBlockMapEntry &TBME : EHInfo.TryBlockMap) {
     for (WinEHHandlerType &H : TBME.HandlerArray) {
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index e446a934554..45ae39af760 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -1205,8 +1205,13 @@ void SelectionDAGBuilder::visitCatchRet(const CatchReturnInst &I) {
   // Figure out the funclet membership for the catchret's successor.
   // This will be used by the FuncletLayout pass to determine how to order the
   // BB's.
-  WinEHFuncInfo *EHInfo = DAG.getMachineFunction().getWinEHFuncInfo();
-  const BasicBlock *SuccessorColor = EHInfo->CatchRetSuccessorColorMap[&I];
+  // A 'catchret' returns to the outer scope's color.
+  Value *ParentPad = I.getParentPad();
+  const BasicBlock *SuccessorColor;
+  if (isa<ConstantTokenNone>(ParentPad))
+    SuccessorColor = &FuncInfo.Fn->getEntryBlock();
+  else
+    SuccessorColor = cast<Instruction>(ParentPad)->getParent();
   assert(SuccessorColor && "No parent funclet for catchret!");
   MachineBasicBlock *SuccessorColorMBB = FuncInfo.MBBMap[SuccessorColor];
   assert(SuccessorColorMBB && "No MBB for SuccessorColor!");
diff --git a/llvm/lib/CodeGen/WinEHPrepare.cpp b/llvm/lib/CodeGen/WinEHPrepare.cpp
index 2426c27d43d..f112b1bb3b9 100644
--- a/llvm/lib/CodeGen/WinEHPrepare.cpp
+++ b/llvm/lib/CodeGen/WinEHPrepare.cpp
@@ -664,24 +664,6 @@ void WinEHPrepare::colorFunclets(Function &F) {
   }
 }
 
-void llvm::calculateCatchReturnSuccessorColors(const Function *Fn,
-                                               WinEHFuncInfo &FuncInfo) {
-  for (const BasicBlock &BB : *Fn) {
-    const auto *CatchRet = dyn_cast<CatchReturnInst>(BB.getTerminator());
-    if (!CatchRet)
-      continue;
-    // A 'catchret' returns to the outer scope's color.
-    Value *ParentPad = CatchRet->getParentPad();
-    const BasicBlock *Color;
-    if (isa<ConstantTokenNone>(ParentPad))
-      Color = &Fn->getEntryBlock();
-    else
-      Color = cast<Instruction>(ParentPad)->getParent();
-    // Record the catchret successor's funclet membership.
-    FuncInfo.CatchRetSuccessorColorMap[CatchRet] = Color;
-  }
-}
-
 void WinEHPrepare::demotePHIsOnFunclets(Function &F) {
   // Strip PHI nodes off of EH pads.
   SmallVector<PHINode *, 16> PHINodes;