2 files changed, 72 insertions, 3 deletions

diff --git a/llvm/lib/CodeGen/WinEHPrepare.cpp b/llvm/lib/CodeGen/WinEHPrepare.cpp
index 56648538f98..122d2fa2ce7 100644
--- a/llvm/lib/CodeGen/WinEHPrepare.cpp
+++ b/llvm/lib/CodeGen/WinEHPrepare.cpp
@@ -306,6 +306,7 @@ static void calculateExplicitCXXStateNumbers(WinEHFuncInfo &FuncInfo,
     BasicBlock *CleanupBlock = CEPI->getCleanupPad()->getParent();
     calculateExplicitCXXStateNumbers(FuncInfo, *CleanupBlock, ParentState);
     // Anything unwinding through CleanupEndPadInst is in ParentState.
+    FuncInfo.EHPadStateMap[FirstNonPHI] = ParentState;
     for (const BasicBlock *PredBlock : predecessors(&BB))
       if ((PredBlock = getEHPadFromPredecessor(PredBlock)))
         calculateExplicitCXXStateNumbers(FuncInfo, *PredBlock, ParentState);
@@ -614,12 +615,20 @@ colorFunclets(Function &F, SmallVectorImpl<BasicBlock *> &EntryBlocks,
         !isa<CleanupEndPadInst>(VisitingHead)) {
       // Mark this as a funclet head as a member of itself.
       FuncletBlocks[Visiting].insert(Visiting);
-      // Queue exits with the parent color.
+      // Queue exits (i.e. successors of rets/endpads) with the parent color.
+      // Skip any exits that are catchendpads, since the parent color must then
+      // represent one of the catches chained to that catchendpad, but the
+      // catchendpad should get the color of the common parent of all its
+      // chained catches (i.e. the grandparent color of the current pad).
+      // We don't need to worry abou catchendpads going unvisited, since the
+      // catches chained to them must have unwind edges to them by which we will
+      // visit them.
       for (User *U : VisitingHead->users()) {
         if (auto *Exit = dyn_cast<TerminatorInst>(U)) {
           for (BasicBlock *Succ : successors(Exit->getParent()))
-            if (BlockColors[Succ].insert(Color).second)
-              Worklist.push_back({Succ, Color});
+            if (!isa<CatchEndPadInst>(*Succ->getFirstNonPHI()))
+              if (BlockColors[Succ].insert(Color).second)
+                Worklist.push_back({Succ, Color});
         }
       }
       // Handle CatchPad specially since its successors need different colors.
diff --git a/llvm/test/CodeGen/WinEH/wineh-cloning.ll b/llvm/test/CodeGen/WinEH/wineh-cloning.ll
index ed5e149d658..e6074f29a3a 100644
--- a/llvm/test/CodeGen/WinEH/wineh-cloning.ll
+++ b/llvm/test/CodeGen/WinEH/wineh-cloning.ll
@@ -486,6 +486,66 @@ unreachable:
   unreachable
 }
 
+define void @test14() personality i32 (...)* @__CxxFrameHandler3 {
+entry:
+  invoke void @f()
+    to label %exit unwind label %catch1.pad
+catch1.pad:
+  %catch1 = catchpad [i32 1]
+    to label %catch1.body unwind label %catch2.pad
+catch1.body:
+  invoke void @h(i32 1)
+    to label %catch1.body2 unwind label %catch.end
+catch1.body2:
+  invoke void @f()
+    to label %catch1.ret unwind label %cleanup1.pad
+cleanup1.pad:
+  %cleanup1 = cleanuppad []
+  call void @f()
+  cleanupret %cleanup1 unwind label %catch.end
+catch1.ret:
+  catchret %catch1 to label %exit
+catch2.pad:
+  %catch2 = catchpad [i32 2]
+    to label %catch2.body unwind label %catch.end
+catch2.body:
+  invoke void @h(i32 2)
+    to label %catch2.body2 unwind label %catch.end
+catch2.body2:
+  invoke void @f()
+    to label %catch2.ret unwind label %cleanup2.pad
+cleanup2.pad:
+  %cleanup2 = cleanuppad []
+  call void @f()
+  cleanupret %cleanup2 unwind label %catch.end
+catch2.ret:
+  catchret %catch2 to label %exit
+catch.end:
+  catchendpad unwind to caller
+exit:
+  ret void
+}
+; Make sure we don't clone the catchendpad even though the
+; cleanupendpads targeting it would naively imply that it
+; should get their respective parent colors (catch1 and catch2),
+; as well as its properly getting the root function color.  The
+; references from the invokes ensure that if we did make clones
+; for each catch, they'd be reachable, as those invokes would get
+; rewritten
+; CHECK-LABEL: define void @test14()
+; CHECK-NOT:  catchendpad
+; CHECK:      invoke void @h(i32 1)
+; CHECK-NEXT:   unwind label %catch.end
+; CHECK-NOT:  catchendpad
+; CHECK:      invoke void @h(i32 2)
+; CHECK-NEXT:   unwind label %catch.end
+; CHECK-NOT:   catchendpad
+; CHECK:     catch.end:
+; CHECK-NEXT:  catchendpad
+; CHECK-NOT:   catchendpad
+
+;; Debug info (from test12)
+
 ; Make sure the DISubprogram doesn't get cloned
 ; CHECK-LABEL: !llvm.module.flags
 ; CHECK-NOT: !DISubprogram