4 files changed, 130 insertions, 7 deletions

diff --git a/llvm/include/llvm/IR/BasicBlock.h b/llvm/include/llvm/IR/BasicBlock.h
index c917b1f2cad..235cb57cfd0 100644
--- a/llvm/include/llvm/IR/BasicBlock.h
+++ b/llvm/include/llvm/IR/BasicBlock.h
@@ -33,6 +33,7 @@ class Function;
 class LandingPadInst;
 class LLVMContext;
 class Module;
+class PHINode;
 class TerminatorInst;
 class ValueSymbolTable;
 
@@ -261,6 +262,50 @@ public:
   inline const Instruction       &back() const { return InstList.back();  }
   inline       Instruction       &back()       { return InstList.back();  }
 
+  /// Iterator to walk just the phi nodes in the basic block.
+  template <typename PHINodeT = PHINode, typename BBIteratorT = iterator>
+  class phi_iterator_impl
+      : public iterator_facade_base<phi_iterator_impl<PHINodeT, BBIteratorT>,
+                                    std::forward_iterator_tag, PHINodeT> {
+    friend BasicBlock;
+
+    PHINodeT *PN;
+
+    phi_iterator_impl(PHINodeT *PN) : PN(PN) {}
+
+  public:
+    // Allow default construction to build variables, but this doesn't build
+    // a useful iterator.
+    phi_iterator_impl() = default;
+
+    // Allow conversion between instantiations where valid.
+    template <typename PHINodeU, typename BBIteratorU>
+    phi_iterator_impl(const phi_iterator_impl<PHINodeU, BBIteratorU> &Arg)
+        : PN(Arg.PN) {}
+
+    bool operator==(const phi_iterator_impl &Arg) const { return PN == Arg.PN; }
+
+    PHINodeT &operator*() const { return *PN; }
+
+    using phi_iterator_impl::iterator_facade_base::operator++;
+    phi_iterator_impl &operator++() {
+      assert(PN && "Cannot increment the end iterator!");
+      PN = dyn_cast<PHINodeT>(std::next(BBIteratorT(PN)));
+      return *this;
+    }
+  };
+  typedef phi_iterator_impl<> phi_iterator;
+  typedef phi_iterator_impl<const PHINode, BasicBlock::const_iterator>
+      const_phi_iterator;
+
+  /// Returns a range that iterates over the phis in the basic block.
+  ///
+  /// Note that this cannot be used with basic blocks that have no terminator.
+  iterator_range<const_phi_iterator> phis() const {
+    return const_cast<BasicBlock *>(this)->phis();
+  }
+  iterator_range<phi_iterator> phis();
+
   /// \brief Return the underlying instruction list container.
   ///
   /// Currently you need to access the underlying instruction list container
diff --git a/llvm/lib/IR/BasicBlock.cpp b/llvm/lib/IR/BasicBlock.cpp
index 90ca21ab91f..1f8659d4e2c 100644
--- a/llvm/lib/IR/BasicBlock.cpp
+++ b/llvm/lib/IR/BasicBlock.cpp
@@ -263,6 +263,10 @@ const BasicBlock *BasicBlock::getUniqueSuccessor() const {
   return SuccBB;
 }
 
+iterator_range<BasicBlock::phi_iterator> BasicBlock::phis() {
+  return make_range<phi_iterator>(dyn_cast<PHINode>(&front()), nullptr);
+}
+
 /// This method is used to notify a BasicBlock that the
 /// specified Predecessor of the block is no longer able to reach it.  This is
 /// actually not used to update the Predecessor list, but is actually used to
@@ -389,13 +393,11 @@ BasicBlock *BasicBlock::splitBasicBlock(iterator I, const Twine &BBName) {
     // Loop over any phi nodes in the basic block, updating the BB field of
     // incoming values...
     BasicBlock *Successor = *I;
-    PHINode *PN;
-    for (BasicBlock::iterator II = Successor->begin();
-         (PN = dyn_cast<PHINode>(II)); ++II) {
-      int IDX = PN->getBasicBlockIndex(this);
-      while (IDX != -1) {
-        PN->setIncomingBlock((unsigned)IDX, New);
-        IDX = PN->getBasicBlockIndex(this);
+    for (auto &PN : Successor->phis()) {
+      int Idx = PN.getBasicBlockIndex(this);
+      while (Idx != -1) {
+        PN.setIncomingBlock((unsigned)Idx, New);
+        Idx = PN.getBasicBlockIndex(this);
       }
     }
   }
diff --git a/llvm/unittests/IR/BasicBlockTest.cpp b/llvm/unittests/IR/BasicBlockTest.cpp
new file mode 100644
index 00000000000..5bf434de961
--- /dev/null
+++ b/llvm/unittests/IR/BasicBlockTest.cpp
@@ -0,0 +1,75 @@
+//===- llvm/unittest/IR/BasicBlockTest.cpp - BasicBlock unit tests --------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/NoFolder.h"
+#include "gmock/gmock-matchers.h"
+#include "gtest/gtest.h"
+#include <memory>
+
+namespace llvm {
+namespace {
+
+TEST(BasicBlockTest, PhiRange) {
+  LLVMContext Context;
+
+  // Create the main block.
+  std::unique_ptr<BasicBlock> BB(BasicBlock::Create(Context));
+
+  // Create some predecessors of it.
+  std::unique_ptr<BasicBlock> BB1(BasicBlock::Create(Context));
+  BranchInst::Create(BB.get(), BB1.get());
+  std::unique_ptr<BasicBlock> BB2(BasicBlock::Create(Context));
+  BranchInst::Create(BB.get(), BB2.get());
+
+  // Make it a cycle.
+  auto *BI = BranchInst::Create(BB.get(), BB.get());
+
+  // Now insert some PHI nodes.
+  auto *Int32Ty = Type::getInt32Ty(Context);
+  auto *P1 = PHINode::Create(Int32Ty, /*NumReservedValues*/ 3, "phi.1", BI);
+  auto *P2 = PHINode::Create(Int32Ty, /*NumReservedValues*/ 3, "phi.2", BI);
+  auto *P3 = PHINode::Create(Int32Ty, /*NumReservedValues*/ 3, "phi.3", BI);
+
+  // Some non-PHI nodes.
+  auto *Sum = BinaryOperator::CreateAdd(P1, P2, "sum", BI);
+
+  // Now wire up the incoming values that are interesting.
+  P1->addIncoming(P2, BB.get());
+  P2->addIncoming(P1, BB.get());
+  P3->addIncoming(Sum, BB.get());
+
+  // Finally, let's iterate them, which is the thing we're trying to test.
+  // We'll use this to wire up the rest of the incoming values.
+  for (auto &PN : BB->phis()) {
+    PN.addIncoming(UndefValue::get(Int32Ty), BB1.get());
+    PN.addIncoming(UndefValue::get(Int32Ty), BB2.get());
+  }
+
+  // Test that we can use const iterators and generally that the iterators
+  // behave like iterators.
+  BasicBlock::const_phi_iterator CI;
+  CI = CI = BB->phis().begin();
+  EXPECT_NE(CI, BB->phis().end());
+
+  // And iterate a const range.
+  for (const auto &PN : const_cast<const BasicBlock *>(BB.get())->phis()) {
+    EXPECT_EQ(BB.get(), PN.getIncomingBlock(0));
+    EXPECT_EQ(BB1.get(), PN.getIncomingBlock(1));
+    EXPECT_EQ(BB2.get(), PN.getIncomingBlock(2));
+  }
+}
+
+} // End anonymous namespace.
+} // End llvm namespace.
diff --git a/llvm/unittests/IR/CMakeLists.txt b/llvm/unittests/IR/CMakeLists.txt
index 750f638c7a4..6734de8e2d9 100644
--- a/llvm/unittests/IR/CMakeLists.txt
+++ b/llvm/unittests/IR/CMakeLists.txt
@@ -8,6 +8,7 @@ set(LLVM_LINK_COMPONENTS
 set(IRSources
   AsmWriterTest.cpp
   AttributesTest.cpp
+  BasicBlockTest.cpp
   ConstantRangeTest.cpp
   ConstantsTest.cpp
   DebugInfoTest.cpp