1 files changed, 70 insertions, 10 deletions

diff --git a/include/llvm/Analysis/RegionInfo.h b/include/llvm/Analysis/RegionInfo.h
index b098eea..188d11c 100644
--- a/include/llvm/Analysis/RegionInfo.h
+++ b/include/llvm/Analysis/RegionInfo.h
@@ -473,25 +473,85 @@ 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 <bool IsConst>
+  class block_iterator_wrapper
+    : public df_iterator<typename conditional<IsConst,
+                                              const BasicBlock,
+                                              BasicBlock>::type*> {
+    typedef df_iterator<typename conditional<IsConst,
+                                             const BasicBlock,
+                                             BasicBlock>::type*>
+      super;
+  public:
+    typedef block_iterator_wrapper<IsConst> Self;
+    typedef typename super::pointer pointer;
+
+    // Construct the begin iterator.
+    block_iterator_wrapper(pointer Entry, pointer Exit) : super(df_begin(Entry))
+    {
+      // Mark the exit of the region as visited, so that the children of the
+      // exit and the exit itself, i.e. the block outside the region will never
+      // be visited.
+      super::Visited.insert(Exit);
+    }
+
+    // Construct the end iterator.
+    block_iterator_wrapper() : super(df_end<pointer>((BasicBlock *)0)) {}
+
+    /*implicit*/ block_iterator_wrapper(super I) : super(I) {}
+
+    // FIXME: Even a const_iterator returns a non-const BasicBlock pointer.
+    //        This was introduced for backwards compatibility, but should
+    //        be removed as soon as all users are fixed.
+    BasicBlock *operator*() const {
+      return const_cast<BasicBlock*>(super::operator*());
+    }
+  };
+
+  typedef block_iterator_wrapper<false> block_iterator;
+  typedef block_iterator_wrapper<true>  const_block_iterator;
+
+  block_iterator block_begin() {
+   return block_iterator(getEntry(), getExit());
+  }
 
-  block_iterator block_begin();
-  block_iterator block_end();
+  block_iterator block_end() {
+   return block_iterator();
+  }
 
-  const_block_iterator block_begin() const;
-  const_block_iterator block_end() const;
+  const_block_iterator block_begin() const {
+    return const_block_iterator(getEntry(), getExit());
+  }
+  const_block_iterator block_end() const {
+    return const_block_iterator();
+  }
   //@}
 
   /// @name Element Iterators