4 files changed, 87 insertions, 18 deletions

diff --git a/include/llvm/Analysis/RegionInfo.h b/include/llvm/Analysis/RegionInfo.h
index b098eeaa3d..eae94e757b 100644
--- a/include/llvm/Analysis/RegionInfo.h
+++ b/include/llvm/Analysis/RegionInfo.h
@@ -473,19 +473,71 @@ public:
   const_iterator end() const { return children.end(); }
   //@}
 
-  /// @name BasicBlock Iterators
+  /// @name BasicBlock Node Iterators
   ///
   /// These iterators iterate over all BasicBlock RegionNodes that are
-  /// contained in this Region. The iterator also iterates over BasicBlocks
-  /// that are elements of a subregion of this Region. It is therefore called a
-  /// flat iterator.
+  /// contained in this Region. The iterator also iterates over BasicBlock
+  /// RegionNodes that are elements of a subregion of this Region. It is
+  /// therefore called a flat iterator.
   //@{
   typedef df_iterator<RegionNode*, SmallPtrSet<RegionNode*, 8>, false,
-                      GraphTraits<FlatIt<RegionNode*> > > block_iterator;
+                      GraphTraits<FlatIt<RegionNode*> > > block_node_iterator;
 
   typedef df_iterator<const RegionNode*, SmallPtrSet<const RegionNode*, 8>,
                       false, GraphTraits<FlatIt<const RegionNode*> > >
-            const_block_iterator;
+            const_block_node_iterator;
+
+  block_node_iterator block_node_begin();
+  block_node_iterator block_node_end();
+
+  const_block_node_iterator block_node_begin() const;
+  const_block_node_iterator block_node_end() const;
+  //@}
+
+  /// @name BasicBlock Iterators
+  ///
+  /// These iterators iterate over all BasicBlocks that are contained in this
+  /// Region. The iterator also iterates over BasicBlocks that are elements of
+  /// a subregion of this Region. It is therefore called a flat iterator.
+  //@{
+  template <typename RegionNodeItT>
+  class block_iterator_wrapper
+    : public std::iterator<std::forward_iterator_tag, BasicBlock, ptrdiff_t> {
+    typedef std::iterator<std::forward_iterator_tag, BasicBlock, ptrdiff_t>
+      super;
+
+    RegionNodeItT Iter;
+
+  public:
+    typedef block_iterator_wrapper<RegionNodeItT> Self;
+    typedef typename super::pointer pointer;
+
+    block_iterator_wrapper(RegionNodeItT Iter) : Iter(Iter) {}
+
+    bool operator==(const Self &RHS) const { return Iter == RHS.Iter; }
+    bool operator!=(const Self &RHS) const { return Iter != RHS.Iter; }
+    pointer operator*() const {
+      return (*Iter)->template getNodeAs<BasicBlock>();
+    }
+
+    Self& operator++() {
+      ++Iter;
+      return *this;
+    }
+    Self operator++(int) {
+      Self tmp = *this;
+      ++*this;
+      return tmp;
+    }
+
+    const Self &operator=(const Self &I) {
+      Iter = I.Iter;
+      return *this;
+    }
+  };
+  typedef block_iterator_wrapper<block_node_iterator> block_iterator;
+  typedef block_iterator_wrapper<const_block_node_iterator>
+    const_block_iterator;
 
   block_iterator block_begin();
   block_iterator block_end();
diff --git a/lib/Analysis/RegionInfo.cpp b/lib/Analysis/RegionInfo.cpp
index b507b1e340..5f4458b402 100644
--- a/lib/Analysis/RegionInfo.cpp
+++ b/lib/Analysis/RegionInfo.cpp
@@ -47,7 +47,7 @@ static cl::opt<enum Region::PrintStyle> printStyle("print-region-style",
   cl::values(
     clEnumValN(Region::PrintNone, "none",  "print no details"),
     clEnumValN(Region::PrintBB, "bb",
-               "print regions in detail with block_iterator"),
+               "print regions in detail with block_node_iterator"),
     clEnumValN(Region::PrintRN, "rn",
                "print regions in detail with element_iterator"),
     clEnumValEnd));
@@ -246,22 +246,38 @@ void Region::verifyRegionNest() const {
   verifyRegion();
 }
 
-Region::block_iterator Region::block_begin() {
+Region::block_node_iterator Region::block_node_begin() {
   return GraphTraits<FlatIt<Region*> >::nodes_begin(this);
 }
 
-Region::block_iterator Region::block_end() {
+Region::block_node_iterator Region::block_node_end() {
   return GraphTraits<FlatIt<Region*> >::nodes_end(this);
 }
 
-Region::const_block_iterator Region::block_begin() const {
+Region::const_block_node_iterator Region::block_node_begin() const {
   return GraphTraits<FlatIt<const Region*> >::nodes_begin(this);
 }
 
-Region::const_block_iterator Region::block_end() const {
+Region::const_block_node_iterator Region::block_node_end() const {
   return GraphTraits<FlatIt<const Region*> >::nodes_end(this);
 }
 
+Region::block_iterator Region::block_begin() {
+  return block_node_begin();
+}
+
+Region::block_iterator Region::block_end() {
+  return block_node_end();
+}
+
+Region::const_block_iterator Region::block_begin() const {
+  return block_node_begin();
+}
+
+Region::const_block_iterator Region::block_end() const {
+  return block_node_end();
+}
+
 Region::element_iterator Region::element_begin() {
   return GraphTraits<Region*>::nodes_begin(this);
 }
@@ -425,7 +441,9 @@ void Region::print(raw_ostream &OS, bool print_tree, unsigned level,
     OS.indent(level*2 + 2);
 
     if (Style == PrintBB) {
-      for (const_block_iterator I = block_begin(), E = block_end(); I!=E; ++I)
+      for (const_block_node_iterator I = block_node_begin(),
+                                     E = block_node_end();
+           I != E; ++I)
         OS << **I << ", "; // TODO: remove the last ","
     } else if (Style == PrintRN) {
       for (const_element_iterator I = element_begin(), E = element_end(); I!=E; ++I)
diff --git a/lib/Analysis/RegionPass.cpp b/lib/Analysis/RegionPass.cpp
index 3a3529baf9..c97b5ebd7d 100644
--- a/lib/Analysis/RegionPass.cpp
+++ b/lib/Analysis/RegionPass.cpp
@@ -195,7 +195,8 @@ public:
 
   virtual bool runOnRegion(Region *R, RGPassManager &RGM) {
     Out << Banner;
-    for (Region::block_iterator I = R->block_begin(), E = R->block_end();
+    for (Region::block_node_iterator I = R->block_node_begin(),
+                                     E = R->block_node_end();
          I != E; ++I)
       (*I)->getEntry()->print(Out);
 
diff --git a/lib/Analysis/RegionPrinter.cpp b/lib/Analysis/RegionPrinter.cpp
index a1730b0a3c..8b23cc7042 100644
--- a/lib/Analysis/RegionPrinter.cpp
+++ b/lib/Analysis/RegionPrinter.cpp
@@ -122,13 +122,11 @@ struct DOTGraphTraits<RegionInfo*> : public DOTGraphTraits<RegionNode*> {
     RegionInfo *RI = R->getRegionInfo();
 
     for (Region::const_block_iterator BI = R->block_begin(),
-         BE = R->block_end(); BI != BE; ++BI) {
-      BasicBlock *BB = (*BI)->getNodeAs<BasicBlock>();
-      if (RI->getRegionFor(BB) == R)
+         BE = R->block_end(); BI != BE; ++BI)
+      if (RI->getRegionFor(*BI) == R)
         O.indent(2 * (depth + 1)) << "Node"
-          << static_cast<const void*>(RI->getTopLevelRegion()->getBBNode(BB))
+          << static_cast<const void*>(RI->getTopLevelRegion()->getBBNode(*BI))
           << ";\n";
-    }
 
     O.indent(2 * depth) << "}\n";
   }