[IR] Add an iterator and range accessor for the PHI nodes of a basic

block.

This allows writing much more natural and readable range based for loops
directly over the PHI nodes. It also takes advantage of the same tricks
for terminating the sequence as the hand coded versions.

I've replaced one example of this mostly to showcase the difference and
I've added a unit test to make sure the facilities really work the way
they're intended. I want to use this inside of SimpleLoopUnswitch but it
seems generally nice.

Differential Revision: https://reviews.llvm.org/D33533

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@303964 91177308-0d34-0410-b5e6-96231b3b80d8

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

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);
       }
     }
   }

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.

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