Allow for customized graph edge pruning in PostOrderIterator.h

Make it possible to prune individual graph edges from a post-order
traversal by specializing the po_iterator_storage template. Previously,
it was only possible to prune full graph nodes. Edge pruning makes it
possible to remove loop back-edges, for example.

Also replace the existing DFSetTraits customization hook with a
po_iterator_storage method for observing the post-order. DFSetTraits was
only used by LoopIterator.h which now provides a po_iterator_storage
specialization.

Thanks to Sean and Chandler for reviewing.

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

1 files changed, 19 insertions, 23 deletions

diff --git a/include/llvm/Analysis/LoopIterator.h b/include/llvm/Analysis/LoopIterator.h
index 269ac80740..68f25f74bc 100644
--- a/include/llvm/Analysis/LoopIterator.h
+++ b/include/llvm/Analysis/LoopIterator.h
@@ -109,6 +109,16 @@ public:
   }
 };
 
+/// Specialize po_iterator_storage to record postorder numbers.
+template<> class po_iterator_storage<LoopBlocksTraversal, true> {
+  LoopBlocksTraversal &LBT;
+public:
+  po_iterator_storage(LoopBlocksTraversal &lbs) : LBT(lbs) {}
+  // These functions are defined below.
+  bool insertEdge(BasicBlock *From, BasicBlock *To);
+  void finishPostorder(BasicBlock *BB);
+};
+
 /// Traverse the blocks in a loop using a depth-first search.
 class LoopBlocksTraversal {
 public:
@@ -155,31 +165,17 @@ public:
     DFS.PostBlocks.push_back(BB);
     DFS.PostNumbers[BB] = DFS.PostBlocks.size();
   }
-
-  //===----------------------------------------------------------------------
-  // Implement part of the std::set interface for the purpose of driving the
-  // generic po_iterator.
-
-  /// Return true if the block is outside the loop or has already been visited.
-  /// Sorry if this is counterintuitive.
-  bool count(BasicBlock *BB) const {
-    return !DFS.L->contains(LI->getLoopFor(BB)) || DFS.PostNumbers.count(BB);
-  }
-
-  /// If this block is contained in the loop and has not been visited, return
-  /// true and assign a preorder number. This is a proxy for visitPreorder
-  /// called by POIterator.
-  bool insert(BasicBlock *BB) {
-    return visitPreorder(BB);
-  }
 };
 
-/// Specialize DFSetTraits to record postorder numbers.
-template<> struct DFSetTraits<LoopBlocksTraversal> {
-  static void finishPostorder(BasicBlock *BB, LoopBlocksTraversal& LBT) {
-    LBT.finishPostorder(BB);
-  }
-};
+inline bool po_iterator_storage<LoopBlocksTraversal, true>::
+insertEdge(BasicBlock *From, BasicBlock *To) {
+  return LBT.visitPreorder(To);
+}
+
+inline void po_iterator_storage<LoopBlocksTraversal, true>::
+finishPostorder(BasicBlock *BB) {
+  LBT.finishPostorder(BB);
+}
 
 } // End namespace llvm