3 files changed, 65 insertions, 13 deletions

diff --git a/llvm/include/llvm/IR/Dominators.h b/llvm/include/llvm/IR/Dominators.h
index def91e73eb1..9be6acc3359 100644
--- a/llvm/include/llvm/IR/Dominators.h
+++ b/llvm/include/llvm/IR/Dominators.h
@@ -66,6 +66,7 @@ public:
     return End;
   }
 
+  /// Check if this is the only edge between Start and End.
   bool isSingleEdge() const;
 };
 
@@ -143,6 +144,11 @@ public:
   bool dominates(const Instruction *Def, const Use &U) const;
   bool dominates(const Instruction *Def, const Instruction *User) const;
   bool dominates(const Instruction *Def, const BasicBlock *BB) const;
+
+  /// Return true if an edge dominates a use.
+  ///
+  /// If BBE is not a unique edge between start and end of the edge, it can
+  /// never dominate the use.
   bool dominates(const BasicBlockEdge &BBE, const Use &U) const;
   bool dominates(const BasicBlockEdge &BBE, const BasicBlock *BB) const;
 
diff --git a/llvm/lib/IR/Dominators.cpp b/llvm/lib/IR/Dominators.cpp
index 44948cc5831..37e735251fd 100644
--- a/llvm/lib/IR/Dominators.cpp
+++ b/llvm/lib/IR/Dominators.cpp
@@ -150,12 +150,6 @@ bool DominatorTree::dominates(const Instruction *Def,
 
 bool DominatorTree::dominates(const BasicBlockEdge &BBE,
                               const BasicBlock *UseBB) const {
-  // Assert that we have a single edge. We could handle them by simply
-  // returning false, but since isSingleEdge is linear on the number of
-  // edges, the callers can normally handle them more efficiently.
-  assert(BBE.isSingleEdge() &&
-         "This function is not efficient in handling multiple edges");
-
   // If the BB the edge ends in doesn't dominate the use BB, then the
   // edge also doesn't.
   const BasicBlock *Start = BBE.getStart();
@@ -188,11 +182,17 @@ bool DominatorTree::dominates(const BasicBlockEdge &BBE,
   // trivially dominates itself, so we only have to find if it dominates the
   // other predecessors. Since the only way out of X is via NormalDest, X can
   // only properly dominate a node if NormalDest dominates that node too.
+  int IsDuplicateEdge = 0;
   for (const_pred_iterator PI = pred_begin(End), E = pred_end(End);
        PI != E; ++PI) {
     const BasicBlock *BB = *PI;
-    if (BB == Start)
+    if (BB == Start) {
+      // If there are multiple edges between Start and End, by definition they
+      // can't dominate anything.
+      if (IsDuplicateEdge++)
+        return false;
       continue;
+    }
 
     if (!dominates(End, BB))
       return false;
@@ -201,12 +201,6 @@ bool DominatorTree::dominates(const BasicBlockEdge &BBE,
 }
 
 bool DominatorTree::dominates(const BasicBlockEdge &BBE, const Use &U) const {
-  // Assert that we have a single edge. We could handle them by simply
-  // returning false, but since isSingleEdge is linear on the number of
-  // edges, the callers can normally handle them more efficiently.
-  assert(BBE.isSingleEdge() &&
-         "This function is not efficient in handling multiple edges");
-
   Instruction *UserInst = cast<Instruction>(U.getUser());
   // A PHI in the end of the edge is dominated by it.
   PHINode *PN = dyn_cast<PHINode>(UserInst);
diff --git a/llvm/unittests/IR/DominatorTreeTest.cpp b/llvm/unittests/IR/DominatorTreeTest.cpp
index d2062839a73..232f0cbd4ed 100644
--- a/llvm/unittests/IR/DominatorTreeTest.cpp
+++ b/llvm/unittests/IR/DominatorTreeTest.cpp
@@ -257,3 +257,55 @@ TEST(DominatorTree, Unreachable) {
         DT->verifyDomTree();
       });
 }
+
+TEST(DominatorTree, NonUniqueEdges) {
+  StringRef ModuleString =
+      "define i32 @f(i32 %i, i32 *%p) {\n"
+      "bb0:\n"
+      "   store i32 %i, i32 *%p\n"
+      "   switch i32 %i, label %bb2 [\n"
+      "     i32 0, label %bb1\n"
+      "     i32 1, label %bb1\n"
+      "   ]\n"
+      " bb1:\n"
+      "   ret i32 1\n"
+      " bb2:\n"
+      "   ret i32 4\n"
+      "}\n";
+
+  // Parse the module.
+  LLVMContext Context;
+  std::unique_ptr<Module> M = makeLLVMModule(Context, ModuleString);
+
+  runWithDomTree(
+      *M, "f",
+      [&](Function &F, DominatorTree *DT, DominatorTreeBase<BasicBlock> *PDT) {
+        Function::iterator FI = F.begin();
+
+        BasicBlock *BB0 = &*FI++;
+        BasicBlock *BB1 = &*FI++;
+        BasicBlock *BB2 = &*FI++;
+
+        const TerminatorInst *TI = BB0->getTerminator();
+        assert(TI->getNumSuccessors() == 3 && "Switch has three successors");
+
+        BasicBlockEdge Edge_BB0_BB2(BB0, TI->getSuccessor(0));
+        assert(Edge_BB0_BB2.getEnd() == BB2 &&
+               "Default label is the 1st successor");
+
+        BasicBlockEdge Edge_BB0_BB1_a(BB0, TI->getSuccessor(1));
+        assert(Edge_BB0_BB1_a.getEnd() == BB1 && "BB1 is the 2nd successor");
+
+        BasicBlockEdge Edge_BB0_BB1_b(BB0, TI->getSuccessor(2));
+        assert(Edge_BB0_BB1_b.getEnd() == BB1 && "BB1 is the 3rd successor");
+
+        EXPECT_TRUE(DT->dominates(Edge_BB0_BB2, BB2));
+        EXPECT_FALSE(DT->dominates(Edge_BB0_BB2, BB1));
+
+        EXPECT_FALSE(DT->dominates(Edge_BB0_BB1_a, BB1));
+        EXPECT_FALSE(DT->dominates(Edge_BB0_BB1_b, BB1));
+
+        EXPECT_FALSE(DT->dominates(Edge_BB0_BB1_a, BB2));
+        EXPECT_FALSE(DT->dominates(Edge_BB0_BB1_b, BB2));
+      });
+}