[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

diff --git a/include/llvm/IR/BasicBlock.h b/include/llvm/IR/BasicBlock.h
index c917b1f2cada1616ee0ca93eaa9bdb49634b1dc8..235cb57cfd09a0a1e18824244996b91c3fdba973 100644 (file)
--- a/include/llvm/IR/BasicBlock.h
+++ b/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/lib/IR/BasicBlock.cpp b/lib/IR/BasicBlock.cpp
index 90ca21ab91f8fcfcaf3ae1e8f0a8ea509dcb91da..1f8659d4e2caefe87fc26915b4cb761d5c718b87 100644 (file)
--- a/lib/IR/BasicBlock.cpp
+++ b/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/unittests/IR/BasicBlockTest.cpp b/unittests/IR/BasicBlockTest.cpp
new file mode 100644 (file)
index 0000000..5bf434d
--- /dev/null
+++ b/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/unittests/IR/CMakeLists.txt b/unittests/IR/CMakeLists.txt
index 750f638c7a42ff23e6ffd5fc628512e354be5e5e..6734de8e2d950d26061b95ae1f99282d8b6aea08 100644 (file)
--- a/unittests/IR/CMakeLists.txt
+++ b/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