[ADT] Refactor post order traversal (#191047)

Currently, po_iterator holds the traversal state. This makes copying
and moving po_iterator fairly expensive and the code cannot be optimized
away in several cases (most of it isn't even inlined in a default
build).

Therefore, refactor post-order traversal to hold the state in a wrapper
class with cheap iterators. Additionally, replace po_storage base class
with a CRTP implementation where users can provide their own storage.

Benefits:

- Performance in stage2-O3 improves by 0.19% instructions:u and even
  more substantially in cycles/wall-time.

- Users that use a custom storage/iteration limitation can do so in a
  more clean way by subclassing PostIteratorTraversalBase. See e.g.
  LoopBlocksTraversal.

- For graphs with block numbers, reserving can now be implemented
  reasonably easy (not done yet).

Implications:

- PostOrderTraversal::iterator is no longer a forward iterator. This
  property was never really used, though.

- PostOrderTraversal must be live while iterators are live. For typical
  uses (for (X x : post_order(...))), this is no problem, but could end
  up being problematic if the iterator is wrapped (e.g.
  for (X x : make_filter_range(post_order(...), ...)) -- problematic,
  because make_filter_range doesn't preserve the range but only the two
  iterators, which become invalid as the for loop is entered). This is a
  limitation of the way LLVM implements ranges.

clang/include/clang/Analysis/Analyses/PostOrderCFGView.h

@@ -29,32 +29,15 @@ class PostOrderCFGView : public ManagedAnalysis {
 
 public:
   /// Implements a set of CFGBlocks using a BitVector.
-  ///
-  /// This class contains a minimal interface, primarily dictated by the SetType
-  /// template parameter of the llvm::po_iterator template, as used with
-  /// external storage. We also use this set to keep track of which CFGBlocks we
-  /// visit during the analysis.
   class CFGBlockSet {
     llvm::BitVector VisitedBlockIDs;
 
   public:
-    // po_iterator requires this iterator, but the only interface needed is the
-    // value_type type.
-    struct iterator { using value_type = const CFGBlock *; };
-
     CFGBlockSet() = default;
     CFGBlockSet(const CFG *G) : VisitedBlockIDs(G->getNumBlockIDs(), false) {}
 
     /// Set the bit associated with a particular CFGBlock.
-    /// This is the important method for the SetType template parameter.
     std::pair<std::nullopt_t, bool> insert(const CFGBlock *Block) {
-      // Note that insert() is called by po_iterator, which doesn't check to
-      // make sure that Block is non-null.  Moreover, the CFGBlock iterator will
-      // occasionally hand out null pointers for pruned edges, so we catch those
-      // here.
-      if (!Block)
-        return std::make_pair(std::nullopt,
-                              false); // if an edge is trivially false.
       if (VisitedBlockIDs.test(Block->getBlockID()))
         return std::make_pair(std::nullopt, false);
       VisitedBlockIDs.set(Block->getBlockID());
@@ -70,41 +53,6 @@ public:
   };
 
 private:
-  // The CFG orders the blocks of loop bodies before those of loop successors
-  // (both numerically, and in the successor order of the loop condition
-  // block). So, RPO necessarily reverses that order, placing the loop successor
-  // *before* the loop body. For many analyses, particularly those that converge
-  // to a fixpoint, this results in potentially significant extra work because
-  // loop successors will necessarily need to be reconsidered once the algorithm
-  // has reached a fixpoint on the loop body.
-  //
-  // This definition of CFG graph traits reverses the order of children, so that
-  // loop bodies will come first in an RPO.
-  struct CFGLoopBodyFirstTraits {
-    using NodeRef = const ::clang::CFGBlock *;
-    using ChildIteratorType = ::clang::CFGBlock::const_succ_reverse_iterator;
-
-    static ChildIteratorType child_begin(NodeRef N) { return N->succ_rbegin(); }
-    static ChildIteratorType child_end(NodeRef N) { return N->succ_rend(); }
-
-    using nodes_iterator = ::clang::CFG::const_iterator;
-
-    static NodeRef getEntryNode(const ::clang::CFG *F) {
-      return &F->getEntry();
-    }
-
-    static nodes_iterator nodes_begin(const ::clang::CFG *F) {
-      return F->nodes_begin();
-    }
-
-    static nodes_iterator nodes_end(const ::clang::CFG *F) {
-      return F->nodes_end();
-    }
-
-    static unsigned size(const ::clang::CFG *F) { return F->size(); }
-  };
-  using po_iterator =
-      llvm::po_iterator<const CFG *, CFGBlockSet, true, CFGLoopBodyFirstTraits>;
   std::vector<const CFGBlock *> Blocks;
 
   using BlockOrderTy = llvm::DenseMap<const CFGBlock *, unsigned>;

clang/lib/Analysis/PostOrderCFGView.cpp

@@ -20,12 +20,40 @@ void PostOrderCFGView::anchor() {}
 
 PostOrderCFGView::PostOrderCFGView(const CFG *cfg) {
   Blocks.reserve(cfg->getNumBlockIDs());
-  CFGBlockSet BSet(cfg);
 
-  for (po_iterator I = po_iterator::begin(cfg, BSet),
-                   E = po_iterator::end(cfg, BSet); I != E; ++I) {
-    BlockOrder[*I] = Blocks.size() + 1;
-    Blocks.push_back(*I);
+  // The CFG orders the blocks of loop bodies before those of loop successors
+  // (both numerically, and in the successor order of the loop condition
+  // block). So, RPO necessarily reverses that order, placing the loop successor
+  // *before* the loop body. For many analyses, particularly those that converge
+  // to a fixpoint, this results in potentially significant extra work because
+  // loop successors will necessarily need to be reconsidered once the algorithm
+  // has reached a fixpoint on the loop body.
+  //
+  // This definition of CFG graph traits reverses the order of children, so that
+  // loop bodies will come first in an RPO.
+  struct CFGLoopBodyFirstTraits {
+    using NodeRef = const ::clang::CFGBlock *;
+    using ChildIteratorType = ::clang::CFGBlock::const_succ_reverse_iterator;
+
+    static ChildIteratorType child_begin(NodeRef N) { return N->succ_rbegin(); }
+    static ChildIteratorType child_end(NodeRef N) { return N->succ_rend(); }
+  };
+
+  struct POTraversal
+      : llvm::PostOrderTraversalBase<POTraversal, CFGLoopBodyFirstTraits> {
+    CFGBlockSet BSet;
+
+    POTraversal(const CFG *cfg) : BSet(cfg) { this->init(&cfg->getEntry()); }
+    bool insertEdge(std::optional<const CFGBlock *>, const CFGBlock *To) {
+      if (!To)
+        return false;
+      return BSet.insert(To).second;
+    }
+  };
+
+  for (const CFGBlock *Block : POTraversal(cfg)) {
+    BlockOrder[Block] = Blocks.size() + 1;
+    Blocks.push_back(Block);
   }
 }
 

llvm/include/llvm/ADT/PostOrderIterator.h

@@ -19,78 +19,12 @@
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/iterator_range.h"
 #include <iterator>
 #include <optional>
-#include <set>
 #include <type_traits>
 #include <utility>
 
 namespace llvm {
-
-// The po_iterator_storage template provides access to the set of already
-// visited nodes during the po_iterator's depth-first traversal.
-//
-// The default implementation simply contains a set of visited nodes, while
-// the External=true version uses a reference to an external set.
-//
-// It is possible to prune the depth-first traversal in several ways:
-//
-// - When providing an external set that already contains some graph nodes,
-//   those nodes won't be visited again. This is useful for restarting a
-//   post-order traversal on a graph with nodes that aren't dominated by a
-//   single node.
-//
-// - By providing a custom SetType class, unwanted graph nodes can be excluded
-//   by having the insert() function return false. This could for example
-//   confine a CFG traversal to blocks in a specific loop.
-//
-// - Finally, by specializing the po_iterator_storage template itself, graph
-//   edges can be pruned by returning false in the insertEdge() function. This
-//   could be used to remove loop back-edges from the CFG seen by po_iterator.
-//
-// A specialized po_iterator_storage class can observe both the pre-order and
-// the post-order. The insertEdge() function is called in a pre-order, while
-// the finishPostorder() function is called just before the po_iterator moves
-// on to the next node.
-
-/// Default po_iterator_storage implementation with an internal set object.
-template<class SetType, bool External>
-class po_iterator_storage {
-  SetType Visited;
-
-public:
-  // Return true if edge destination should be visited.
-  template <typename NodeRef>
-  bool insertEdge(std::optional<NodeRef> From, NodeRef To) {
-    return Visited.insert(To).second;
-  }
-
-  // Called after all children of BB have been visited.
-  template <typename NodeRef> void finishPostorder(NodeRef BB) {}
-};
-
-/// Specialization of po_iterator_storage that references an external set.
-template<class SetType>
-class po_iterator_storage<SetType, true> {
-  SetType &Visited;
-
-public:
-  po_iterator_storage(SetType &VSet) : Visited(VSet) {}
-  po_iterator_storage(const po_iterator_storage &S) : Visited(S.Visited) {}
-
-  // Return true if edge destination should be visited, called with From = 0 for
-  // the root node.
-  // Graph edges can be pruned by specializing this function.
-  template <class NodeRef>
-  bool insertEdge(std::optional<NodeRef> From, NodeRef To) {
-    return Visited.insert(To).second;
-  }
-
-  // Called after all children of BB have been visited.
-  template <class NodeRef> void finishPostorder(NodeRef BB) {}
-};
-
 namespace po_detail {
 
 template <typename NodeRef> class NumberSet {
@@ -120,180 +54,200 @@ using DefaultSet =
 
 } // namespace po_detail
 
-template <class GraphT, class SetType = po_detail::DefaultSet<GraphT>,
-          bool ExtStorage = false, class GT = GraphTraits<GraphT>>
-class po_iterator : public po_iterator_storage<SetType, ExtStorage> {
-public:
-  // When External storage is used we are not multi-pass safe.
-  using iterator_category =
-      std::conditional_t<ExtStorage, std::input_iterator_tag,
-                         std::forward_iterator_tag>;
-  using value_type = typename GT::NodeRef;
-  using difference_type = std::ptrdiff_t;
-  using pointer = value_type *;
-  using reference = const value_type &;
-
-private:
-  using NodeRef = typename GT::NodeRef;
-  using ChildItTy = typename GT::ChildIteratorType;
+/// CRTP base class for post-order traversal. Storage for visited nodes must be
+/// provided by the sub-class, which must implement insertEdge(). Due to CRTP
+/// limitations, the sub-class must call init() with the start node before
+/// traversing; not calling init results in an empty iterator.
+///
+/// Sub-classes can observe the post-order traversal with finishPostorder(),
+/// which is called before the iterator moves to the next node, and also the
+/// pre-order traversal with insertEdge().
+///
+/// Unwanted graph nodes (e.g. from a previous traversal) can be skipped by
+/// returning false from insertEdge().
+///
+/// This class only supports a single traversal of the graph.
+template <typename DerivedT, typename GraphTraits>
+class PostOrderTraversalBase {
+  using NodeRef = typename GraphTraits::NodeRef;
+  using ChildItTy = typename GraphTraits::ChildIteratorType;
 
   /// Used to maintain the ordering.
   /// First element is basic block pointer, second is iterator for the next
   /// child to visit, third is the end iterator.
   SmallVector<std::tuple<NodeRef, ChildItTy, ChildItTy>, 8> VisitStack;
 
-  po_iterator(NodeRef BB) {
-    this->insertEdge(std::optional<NodeRef>(), BB);
-    VisitStack.emplace_back(BB, GT::child_begin(BB), GT::child_end(BB));
-    traverseChild();
-  }
+public:
+  class iterator {
+    friend class PostOrderTraversalBase;
 
-  po_iterator() = default; // End is when stack is empty.
+  public:
+    using iterator_category = std::input_iterator_tag;
+    using value_type = NodeRef;
+    using difference_type = std::ptrdiff_t;
+    using pointer = value_type *;
+    using reference = NodeRef;
 
-  po_iterator(NodeRef BB, SetType &S)
-      : po_iterator_storage<SetType, ExtStorage>(S) {
-    if (this->insertEdge(std::optional<NodeRef>(), BB)) {
-      VisitStack.emplace_back(BB, GT::child_begin(BB), GT::child_end(BB));
+  private:
+    DerivedT *POT = nullptr;
+    NodeRef V = nullptr;
+
+  public:
+    iterator() = default;
+
+  private:
+    iterator(DerivedT &POT, value_type V) : POT(&POT), V(V) {}
+
+  public:
+    bool operator==(const iterator &X) const { return V == X.V; }
+    bool operator!=(const iterator &X) const { return !(*this == X); }
+
+    reference operator*() const { return V; }
+    pointer operator->() const { return &V; }
+
+    iterator &operator++() { // Preincrement
+      V = POT->next();
+      return *this;
+    }
+
+    iterator operator++(int) { // Postincrement
+      iterator tmp = *this;
+      ++*this;
+      return tmp;
+    }
+  };
+
+protected:
+  PostOrderTraversalBase() = default;
+
+  DerivedT *derived() { return static_cast<DerivedT *>(this); }
+
+  /// Initialize post-order traversal at given start node.
+  void init(NodeRef Start) {
+    if (derived()->insertEdge(std::optional<NodeRef>(), Start)) {
+      VisitStack.emplace_back(Start, GraphTraits::child_begin(Start),
+                              GraphTraits::child_end(Start));
       traverseChild();
     }
   }
 
-  po_iterator(SetType &S)
-      : po_iterator_storage<SetType, ExtStorage>(S) {
-  } // End is when stack is empty.
-
+private:
   void traverseChild() {
     while (true) {
       auto &Entry = VisitStack.back();
       if (std::get<1>(Entry) == std::get<2>(Entry))
         break;
       NodeRef BB = *std::get<1>(Entry)++;
-      if (this->insertEdge(std::optional<NodeRef>(std::get<0>(Entry)), BB)) {
+      if (derived()->insertEdge(std::optional<NodeRef>(std::get<0>(Entry)),
+                                BB)) {
         // If the block is not visited...
-        VisitStack.emplace_back(BB, GT::child_begin(BB), GT::child_end(BB));
+        VisitStack.emplace_back(BB, GraphTraits::child_begin(BB),
+                                GraphTraits::child_end(BB));
       }
     }
   }
 
-public:
-  // Provide static "constructors"...
-  static po_iterator begin(const GraphT &G) {
-    return po_iterator(GT::getEntryNode(G));
-  }
-  static po_iterator end(const GraphT &G) { return po_iterator(); }
-
-  static po_iterator begin(const GraphT &G, SetType &S) {
-    return po_iterator(GT::getEntryNode(G), S);
-  }
-  static po_iterator end(const GraphT &G, SetType &S) { return po_iterator(S); }
-
-  bool operator==(const po_iterator &x) const {
-    return VisitStack == x.VisitStack;
-  }
-  bool operator!=(const po_iterator &x) const { return !(*this == x); }
-
-  reference operator*() const { return std::get<0>(VisitStack.back()); }
-
-  // This is a nonstandard operator-> that dereferences the pointer an extra
-  // time... so that you can actually call methods ON the BasicBlock, because
-  // the contained type is a pointer.  This allows BBIt->getTerminator() f.e.
-  //
-  NodeRef operator->() const { return **this; }
-
-  po_iterator &operator++() { // Preincrement
-    this->finishPostorder(std::get<0>(VisitStack.back()));
+  NodeRef next() {
+    derived()->finishPostorder(std::get<0>(VisitStack.back()));
     VisitStack.pop_back();
     if (!VisitStack.empty())
       traverseChild();
-    return *this;
+    return !VisitStack.empty() ? std::get<0>(VisitStack.back()) : nullptr;
   }
 
-  po_iterator operator++(int) { // Postincrement
-    po_iterator tmp = *this;
-    ++*this;
-    return tmp;
+public:
+  iterator begin() {
+    if (VisitStack.empty())
+      return iterator(); // We don't even want to see the start node.
+    return iterator(*derived(), std::get<0>(VisitStack.back()));
+  }
+  iterator end() { return iterator(); }
+
+  // Methods that are intended to be overridden by sub-classes.
+
+  /// Add edge and return whether To should be visited. From is nullopt for the
+  /// root node.
+  bool insertEdge(std::optional<NodeRef> From, NodeRef To);
+
+  /// Callback just before the iterator moves to the next block.
+  void finishPostorder(NodeRef) {}
+};
+
+/// Post-order traversal of a graph. Note: the traversal state is stored in this
+/// class, not in the iterators -- the lifetime of PostOrderTraversal must
+/// exceed the lifetime of the iterators. Special care must be taken with
+/// range-based for-loops in combination with LLVM ranges:
+///
+///   // Fine:
+///   for (BasicBlock *BB : post_order(F)) { ... }
+///
+///   // Problematic! Lifetime of PostOrderTraversal ends before the loop is
+///   // entered, because make_filter_range only stores the iterators but not
+///   // the range object itself.
+///   for (BasicBlock *BB : make_filter_range(post_order(F), ...)) { ... }
+///   // Fixed:
+///   auto POT = post_order(F);
+///   for (BasicBlock *BB : make_filter_range(POT, ...)) { ... }
+///
+/// This class only supports a single traversal of the graph.
+template <typename GraphT, typename SetType = po_detail::DefaultSet<GraphT>>
+class PostOrderTraversal
+    : public PostOrderTraversalBase<PostOrderTraversal<GraphT, SetType>,
+                                    GraphTraits<GraphT>> {
+  using NodeRef = typename GraphTraits<GraphT>::NodeRef;
+
+  SetType Visited;
+
+public:
+  /// Default constructor for an empty traversal.
+  PostOrderTraversal() = default;
+
+  /// Post-order traversal of the graph starting at the root node using an
+  /// internal storage.
+  PostOrderTraversal(const GraphT &G) {
+    this->init(GraphTraits<GraphT>::getEntryNode(G));
+  }
+
+  bool insertEdge(std::optional<NodeRef> From, NodeRef To) {
+    return Visited.insert(To).second;
+  }
+};
+
+/// Post-order traversal of the graph starting at the root node using an
+/// external storage. This can be used to keep track of visited nodes after the
+/// traversal and to skip nodes that are already contained in the set. See
+/// PostOrderTraversal for usage restrictions.
+template <typename GraphT, typename SetType>
+class PostOrderExtTraversal
+    : public PostOrderTraversalBase<PostOrderExtTraversal<GraphT, SetType>,
+                                    GraphTraits<GraphT>> {
+  using NodeRef = typename GraphTraits<GraphT>::NodeRef;
+
+  SetType &Visited;
+
+public:
+  PostOrderExtTraversal(const GraphT &G, SetType &S) : Visited(S) {
+    this->init(GraphTraits<GraphT>::getEntryNode(G));
+  }
+
+  bool insertEdge(std::optional<NodeRef> From, NodeRef To) {
+    return Visited.insert(To).second;
   }
 };
 
 // Provide global constructors that automatically figure out correct types...
 //
-template <class T>
-po_iterator<T> po_begin(const T &G) { return po_iterator<T>::begin(G); }
-template <class T>
-po_iterator<T> po_end  (const T &G) { return po_iterator<T>::end(G); }
-
-template <class T> iterator_range<po_iterator<T>> post_order(const T &G) {
-  return make_range(po_begin(G), po_end(G));
+/// Post-order traversal of a graph. Note: this returns a PostOrderTraversal,
+/// not an iterator range; \see PostOrderTraversal.
+template <class T> auto post_order(const T &G) {
+  return PostOrderTraversal<T>(G);
+}
+template <class T, class SetType> auto post_order_ext(const T &G, SetType &S) {
+  return PostOrderExtTraversal<T, SetType>(G, S);
 }
-
-// Provide global definitions of external postorder iterators...
-template <class T, class SetType = std::set<typename GraphTraits<T>::NodeRef>>
-struct po_ext_iterator : po_iterator<T, SetType, true> {
-  po_ext_iterator(const po_iterator<T, SetType, true> &V) :
-  po_iterator<T, SetType, true>(V) {}
-};
-
 template <class T, class SetType>
-po_ext_iterator<T, SetType> po_ext_begin(const T &G, SetType &S) {
-  return po_ext_iterator<T, SetType>::begin(G, S);
-}
-
-template <class T, class SetType>
-po_ext_iterator<T, SetType> po_ext_end(const T &G, SetType &S) {
-  return po_ext_iterator<T, SetType>::end(G, S);
-}
-
-template <class T, class SetType>
-iterator_range<po_ext_iterator<T, SetType>> post_order_ext(const T &G, SetType &S) {
-  return make_range(po_ext_begin(G, S), po_ext_end(G, S));
-}
-
-// Provide global definitions of inverse post order iterators...
-template <class T, class SetType = std::set<typename GraphTraits<T>::NodeRef>,
-          bool External = false>
-struct ipo_iterator : po_iterator<Inverse<T>, SetType, External> {
-  ipo_iterator(const po_iterator<Inverse<T>, SetType, External> &V) :
-     po_iterator<Inverse<T>, SetType, External> (V) {}
-};
-
-template <class T>
-ipo_iterator<T> ipo_begin(const T &G) {
-  return ipo_iterator<T>::begin(G);
-}
-
-template <class T>
-ipo_iterator<T> ipo_end(const T &G){
-  return ipo_iterator<T>::end(G);
-}
-
-template <class T>
-iterator_range<ipo_iterator<T>> inverse_post_order(const T &G) {
-  return make_range(ipo_begin(G), ipo_end(G));
-}
-
-// Provide global definitions of external inverse postorder iterators...
-template <class T, class SetType = std::set<typename GraphTraits<T>::NodeRef>>
-struct ipo_ext_iterator : ipo_iterator<T, SetType, true> {
-  ipo_ext_iterator(const ipo_iterator<T, SetType, true> &V) :
-    ipo_iterator<T, SetType, true>(V) {}
-  ipo_ext_iterator(const po_iterator<Inverse<T>, SetType, true> &V) :
-    ipo_iterator<T, SetType, true>(V) {}
-};
-
-template <class T, class SetType>
-ipo_ext_iterator<T, SetType> ipo_ext_begin(const T &G, SetType &S) {
-  return ipo_ext_iterator<T, SetType>::begin(G, S);
-}
-
-template <class T, class SetType>
-ipo_ext_iterator<T, SetType> ipo_ext_end(const T &G, SetType &S) {
-  return ipo_ext_iterator<T, SetType>::end(G, S);
-}
-
-template <class T, class SetType>
-iterator_range<ipo_ext_iterator<T, SetType>>
-inverse_post_order_ext(const T &G, SetType &S) {
-  return make_range(ipo_ext_begin(G, S), ipo_ext_end(G, S));
+auto inverse_post_order_ext(const T &G, SetType &S) {
+  return PostOrderExtTraversal<Inverse<T>, SetType>(G, S);
 }
 
 //===--------------------------------------------------------------------===//
@@ -331,7 +285,7 @@ class ReversePostOrderTraversal {
   VecTy Blocks; // Block list in normal PO order
 
   void Initialize(const GraphT &G) {
-    std::copy(po_begin(G), po_end(G), std::back_inserter(Blocks));
+    llvm::copy(post_order(G), std::back_inserter(Blocks));
   }
 
 public:

llvm/include/llvm/Analysis/LoopIterator.h

@@ -186,57 +186,41 @@ public:
   LoopBlocksDFS::RPOIterator end() const { return DFS.endRPO(); }
 };
 
-/// 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(std::optional<BasicBlock *> From, BasicBlock *To);
-  void finishPostorder(BasicBlock *BB);
-};
-
 /// Traverse the blocks in a loop using a depth-first search.
-class LoopBlocksTraversal {
-public:
-  /// Graph traversal iterator.
-  typedef po_iterator<BasicBlock*, LoopBlocksTraversal, true> POTIterator;
-
-private:
+class LoopBlocksTraversal
+    : public PostOrderTraversalBase<LoopBlocksTraversal,
+                                    GraphTraits<Function *>> {
   LoopBlocksDFS &DFS;
   const LoopInfo *LI;
 
 public:
-  LoopBlocksTraversal(LoopBlocksDFS &Storage, const LoopInfo *LInfo) :
-    DFS(Storage), LI(LInfo) {}
+  LoopBlocksTraversal(LoopBlocksDFS &Storage, const LoopInfo *LInfo)
+      : DFS(Storage), LI(LInfo) {}
 
   /// Postorder traversal over the graph. This only needs to be done once.
-  /// po_iterator "automatically" calls back to visitPreorder and
+  /// PostOrderTraversalBase "automatically" calls back to insertEdge and
   /// finishPostorder to record the DFS result.
-  POTIterator begin() {
+  iterator begin() {
     assert(DFS.PostBlocks.empty() && "Need clear DFS result before traversing");
-    assert(DFS.L->getNumBlocks() && "po_iterator cannot handle an empty graph");
-    return po_ext_begin(DFS.L->getHeader(), *this);
-  }
-  POTIterator end() {
-    // po_ext_end interface requires a basic block, but ignores its value.
-    return po_ext_end(DFS.L->getHeader(), *this);
+    assert(DFS.L->getNumBlocks() && "cannot handle an empty graph");
+    init(DFS.L->getHeader());
+    return PostOrderTraversalBase::begin();
   }
+  iterator end() { return PostOrderTraversalBase::end(); }
 
-  /// Called by po_iterator upon reaching a block via a CFG edge. If this block
-  /// is contained in the loop and has not been visited, then mark it preorder
-  /// visited and return true.
+  /// Called upon reaching a block via a CFG edge. If this block is contained
+  /// in the loop and has not been visited, then mark it preorder visited and
+  /// return true (i.e., traverse the edge).
   ///
   /// TODO: If anyone is interested, we could record preorder numbers here.
-  bool visitPreorder(BasicBlock *BB) {
+  bool insertEdge(std::optional<BasicBlock *> /*From*/, BasicBlock *BB) {
     if (!DFS.L->contains(LI->getLoopFor(BB)))
       return false;
 
     return DFS.PostNumbers.insert(std::make_pair(BB, 0)).second;
   }
 
-  /// Called by po_iterator each time it advances, indicating a block's
-  /// postorder.
+  /// Called each time the iterator advances, indicating a block's postorder.
   void finishPostorder(BasicBlock *BB) {
     assert(DFS.PostNumbers.count(BB) && "Loop DFS skipped preorder");
     DFS.PostBlocks.push_back(BB);
@@ -244,16 +228,6 @@ public:
   }
 };
 
-inline bool po_iterator_storage<LoopBlocksTraversal, true>::insertEdge(
-    std::optional<BasicBlock *> From, BasicBlock *To) {
-  return LBT.visitPreorder(To);
-}
-
-inline void po_iterator_storage<LoopBlocksTraversal, true>::
-finishPostorder(BasicBlock *BB) {
-  LBT.finishPostorder(BB);
-}
-
 } // End namespace llvm
 
 #endif

llvm/lib/Analysis/LoopInfo.cpp

@@ -1284,8 +1284,6 @@ PreservedAnalyses LoopVerifierPass::run(Function &F,
 /// visit blocks during the initial traversal.
 void LoopBlocksDFS::perform(const LoopInfo *LI) {
   LoopBlocksTraversal Traversal(*this, LI);
-  for (LoopBlocksTraversal::POTIterator POI = Traversal.begin(),
-                                        POE = Traversal.end();
-       POI != POE; ++POI)
+  for ([[maybe_unused]] BasicBlock *BB : Traversal)
     ;
 }

llvm/lib/CodeGen/MachineTraceMetrics.cpp

@@ -482,15 +482,19 @@ struct LoopBounds {
 
 } // end anonymous namespace
 
-// Specialize po_iterator_storage in order to prune the post-order traversal so
-// it is limited to the current loop and doesn't traverse the loop back edges.
-template <> class llvm::po_iterator_storage<LoopBounds, true> {
+// Restrict the post-order traversal to the current loop and don't traverse the
+// loop back edges.
+template <typename GraphT>
+class LoopBoundsPostOrderTraversal
+    : public PostOrderTraversalBase<LoopBoundsPostOrderTraversal<GraphT>,
+                                    GraphTraits<GraphT>> {
   LoopBounds &LB;
 
 public:
-  po_iterator_storage(LoopBounds &lb) : LB(lb) {}
-
-  void finishPostorder(const MachineBasicBlock*) {}
+  LoopBoundsPostOrderTraversal(const MachineBasicBlock *Start, LoopBounds &LB)
+      : LB(LB) {
+    this->init(Start);
+  }
 
   bool insertEdge(std::optional<const MachineBasicBlock *> From,
                   const MachineBasicBlock *To) {
@@ -525,7 +529,9 @@ void MachineTraceMetrics::Ensemble::computeTrace(const MachineBasicBlock *MBB) {
   // Run an upwards post-order search for the trace start.
   Bounds.Downward = false;
   Bounds.Visited.clear();
-  for (const auto *I : inverse_post_order_ext(MBB, Bounds)) {
+  for (const auto *I :
+       LoopBoundsPostOrderTraversal<Inverse<const MachineBasicBlock *>>(
+           MBB, Bounds)) {
     LLVM_DEBUG(dbgs() << "  pred for " << printMBBReference(*I) << ": ");
     TraceBlockInfo &TBI = BlockInfo[I->getNumber()];
     // All the predecessors have been visited, pick the preferred one.
@@ -543,7 +549,8 @@ void MachineTraceMetrics::Ensemble::computeTrace(const MachineBasicBlock *MBB) {
   // Run a downwards post-order search for the trace end.
   Bounds.Downward = true;
   Bounds.Visited.clear();
-  for (const auto *I : post_order_ext(MBB, Bounds)) {
+  for (const auto *I :
+       LoopBoundsPostOrderTraversal<const MachineBasicBlock *>(MBB, Bounds)) {
     LLVM_DEBUG(dbgs() << "  succ for " << printMBBReference(*I) << ": ");
     TraceBlockInfo &TBI = BlockInfo[I->getNumber()];
     // All the successors have been visited, pick the preferred one.

llvm/lib/Target/SPIRV/SPIRVUtils.h

@@ -21,6 +21,7 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/TypedPointerType.h"
 #include <queue>
+#include <set>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>

llvm/lib/Transforms/Vectorize/VPlanCFG.h

@@ -261,15 +261,14 @@ vp_depth_first_shallow(const VPBlockBase *G) {
 
 /// Returns an iterator range to traverse the graph starting at \p G in
 /// post order. The iterator won't traverse through region blocks.
-inline iterator_range<
-    po_iterator<VPBlockShallowTraversalWrapper<VPBlockBase *>>>
+inline PostOrderTraversal<VPBlockShallowTraversalWrapper<VPBlockBase *>>
 vp_post_order_shallow(VPBlockBase *G) {
   return post_order(VPBlockShallowTraversalWrapper<VPBlockBase *>(G));
 }
 
 /// Returns an iterator range to traverse the graph starting at \p G in
 /// post order while traversing through region blocks.
-inline iterator_range<po_iterator<VPBlockDeepTraversalWrapper<VPBlockBase *>>>
+inline PostOrderTraversal<VPBlockDeepTraversalWrapper<VPBlockBase *>>
 vp_post_order_deep(VPBlockBase *G) {
   return post_order(VPBlockDeepTraversalWrapper<VPBlockBase *>(G));
 }

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -734,8 +734,9 @@ static bool isDeadRecipe(VPRecipeBase &R) {
 }
 
 void VPlanTransforms::removeDeadRecipes(VPlan &Plan) {
-  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
-           vp_post_order_deep(Plan.getEntry()))) {
+  PostOrderTraversal<VPBlockDeepTraversalWrapper<VPBlockBase *>> POT(
+      Plan.getEntry());
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(POT)) {
     // The recipes in the block are processed in reverse order, to catch chains
     // of dead recipes.
     for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) {
@@ -2707,8 +2708,9 @@ static void licm(VPlan &Plan) {
   // Sink recipes with no users inside the vector loop region if all users are
   // in the same exit block of the region.
   // TODO: Extend to sink recipes from inner loops.
-  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
-           vp_post_order_shallow(LoopRegion->getEntry()))) {
+  PostOrderTraversal<VPBlockShallowTraversalWrapper<VPBlockBase *>> POT(
+      LoopRegion->getEntry());
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(POT)) {
     for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) {
       if (cannotHoistOrSinkRecipe(R))
         continue;

llvm/lib/Transforms/Vectorize/VPlanUtils.h

@@ -269,8 +269,7 @@ public:
 
   /// Return an iterator range over \p Range which only includes \p BlockTy
   /// blocks. The accesses are casted to \p BlockTy.
-  template <typename BlockTy, typename T>
-  static auto blocksOnly(const T &Range) {
+  template <typename BlockTy, typename T> static auto blocksOnly(T &&Range) {
     // Create BaseTy with correct const-ness based on BlockTy.
     using BaseTy = std::conditional_t<std::is_const<BlockTy>::value,
                                       const VPBlockBase, VPBlockBase>;

llvm/unittests/ADT/PostOrderIteratorTest.cpp

@@ -23,30 +23,22 @@ namespace {
 
 // Whether we're able to compile
 TEST(PostOrderIteratorTest, Compiles) {
-  using ExtSetTy = SmallPtrSet<void *, 4>;
-
-  // Tests that template specializations are kept up to date
-  void *Null = nullptr;
-  po_iterator_storage<std::set<void *>, false> PIS;
-  PIS.insertEdge(std::optional<void *>(), Null);
-  ExtSetTy Ext;
-  po_iterator_storage<ExtSetTy, true> PISExt(Ext);
-  PIS.insertEdge(std::optional<void *>(), Null);
-
-  // Test above, but going through po_iterator (which inherits from template
-  // base)
   Graph<6> G;
   using NodeType = Graph<6>::NodeType;
   NodeType *NullNode = nullptr;
-  auto PI = po_end(G);
+
+  auto PI = post_order(G);
   PI.insertEdge(std::optional<NodeType *>(), NullNode);
-  auto PIExt = po_ext_end(G, Ext);
+
+  using ExtSetTy = SmallPtrSet<void *, 4>;
+  ExtSetTy Ext;
+  auto PIExt = post_order_ext(G, Ext);
   PIExt.insertEdge(std::optional<NodeType *>(), NullNode);
 }
 
-static_assert(
-    std::is_convertible_v<decltype(*std::declval<po_iterator<Graph<3>>>()),
-                          po_iterator<Graph<3>>::reference>);
+static_assert(std::is_convertible_v<
+              decltype(*post_order(std::declval<Graph<3>>()).begin()),
+              PostOrderTraversal<Graph<3>>::iterator::reference>);
 
 // Test post-order and reverse post-order traversals for simple graph type.
 TEST(PostOrderIteratorTest, PostOrderAndReversePostOrderTraverrsal) {
@@ -83,34 +75,4 @@ TEST(PostOrderIteratorTest, PostOrderAndReversePostOrderTraverrsal) {
   EXPECT_EQ(1, FromIterator[4]);
   EXPECT_EQ(4, FromIterator[5]);
 }
-
-// po_iterator should be (at-least) a forward-iterator
-static_assert(std::is_base_of_v<std::forward_iterator_tag,
-                                po_iterator<Graph<4>>::iterator_category>);
-
-// po_ext_iterator cannot provide multi-pass guarantee, therefore its only
-// an input-iterator
-static_assert(std::is_same_v<po_ext_iterator<Graph<4>>::iterator_category,
-                             std::input_iterator_tag>);
-
-TEST(PostOrderIteratorTest, MultiPassSafeWithInternalSet) {
-  Graph<4> G;
-  G.AddEdge(0, 1);
-  G.AddEdge(1, 2);
-  G.AddEdge(1, 3);
-
-  std::array<decltype(G)::NodeType *, 4> NodesFirstPass, NodesSecondPass;
-
-  auto B = po_begin(G), E = po_end(G);
-
-  std::size_t I = 0;
-  for (auto It = B; It != E; ++It)
-    NodesFirstPass[I++] = *It;
-
-  I = 0;
-  for (auto It = B; It != E; ++It)
-    NodesSecondPass[I++] = *It;
-
-  EXPECT_EQ(NodesFirstPass, NodesSecondPass);
-}
 }

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

@@ -546,7 +546,7 @@ TEST_F(VPBasicBlockTest, TraversingIteratorTest) {
 
     // Post-order.
     FromIterator.clear();
-    FromIterator.append(po_begin(Start), po_end(Start));
+    copy(post_order(Start), std::back_inserter(FromIterator));
     EXPECT_EQ(10u, FromIterator.size());
     EXPECT_EQ(VPBB2, FromIterator[0]);
     EXPECT_EQ(R1BB3, FromIterator[1]);
@@ -597,7 +597,7 @@ TEST_F(VPBasicBlockTest, TraversingIteratorTest) {
 
     // Post-order.
     FromIterator.clear();
-    FromIterator.append(po_begin(Start), po_end(Start));
+    copy(post_order(Start), std::back_inserter(FromIterator));
     EXPECT_EQ(5u, FromIterator.size());
     EXPECT_EQ(R2BB2, FromIterator[0]);
     EXPECT_EQ(R2BB1, FromIterator[1]);
@@ -678,8 +678,9 @@ TEST_F(VPBasicBlockTest, TraversingIteratorTest) {
 
     // Post-order, const VPRegionBlocks only.
     VPBlockDeepTraversalWrapper<const VPBlockBase *> StartConst(VPBB1);
-    SmallVector<const VPRegionBlock *> FromIteratorVPRegion(
-        VPBlockUtils::blocksOnly<const VPRegionBlock>(post_order(StartConst)));
+    SmallVector<const VPRegionBlock *> FromIteratorVPRegion;
+    copy(VPBlockUtils::blocksOnly<const VPRegionBlock>(post_order(StartConst)),
+         std::back_inserter(FromIteratorVPRegion));
     EXPECT_EQ(3u, FromIteratorVPRegion.size());
     EXPECT_EQ(R3, FromIteratorVPRegion[0]);
     EXPECT_EQ(R2, FromIteratorVPRegion[1]);

mlir/include/mlir/IR/Iterators.h

@@ -102,21 +102,23 @@ struct ReverseDominanceIterator {
   }
 
   static auto makeIterable(Region &region) {
-    Block *null = nullptr;
-    if (SkipGraphRegion && !mayHaveSSADominance(region)) {
-      // Skip graph regions.
-      return llvm::make_pointee_range(
-          llvm::make_range(llvm::po_end(null), llvm::po_end(null)));
-    }
+    struct Traversal : llvm::PostOrderTraversal<Region *> {
+      using BaseT = llvm::PostOrderTraversal<Region *>;
+      using BaseT::PostOrderTraversal; // Re-use constructors.
+
+      // Walk API expects Block references instead of pointers.
+      using iterator = llvm::pointee_iterator<typename BaseT::iterator>;
+      iterator begin() { return iterator(BaseT::begin()); }
+      iterator end() { return iterator(BaseT::end()); }
+    };
+
+    // Skip graph regions.
+    if ((SkipGraphRegion && !mayHaveSSADominance(region)) || region.empty())
+      return Traversal();
 
     // Create post-order iterator. Blocks are enumerated according to their
     // successor relationship.
-    auto it = region.empty()
-                  ? llvm::make_range(llvm::po_end(null), llvm::po_end(null))
-                  : llvm::post_order(&region.front());
-
-    // Walk API expects Block references instead of pointers.
-    return llvm::make_pointee_range(it);
+    return Traversal(&region);
   }
 };
 } // namespace mlir