Modify df_iterator to support post-order actions

[david2050]

Summary

This makes a change to the state used to maintain visited information for depth first iterator. We know assume a method "completed(...)" which is called after all children of a node have been visited. In all existing cases, this method does nothing so this patch has no functional changes. It will however allow a client to distinguish back from cross edges in a DFS tree.

[mehdi_amini]

This looks fine to me as well.
(I don't know how you'll figure back-edge with it, but I'll see in the next patch)

[david2050]

For illustration, this is what the client code looks like (abstracted a bit) for at least one back edge in a loop.
(Note this code is still order dependent in that the specific edges selected as back edges is order dependent for irreducible loops but is well defined otherwise).

    // This stores state for the depth-first iterator. In addition
    // to recording which nodes have been visited we also record whether
    // a node is currently on the "stack" of active ancestors of the current
    // node.
    class DFState {
      DenseMap<BasicBlock *, bool> Status;

    public:
      typedef typename DenseMap<BasicBlock *, bool>::iterator iterator;

      DFState(unsigned N) { Status.reserve(N); }
      std::pair<iterator, bool> insert(BasicBlock *BB) {
        return Status.insert(std::make_pair(BB, true));
      }
      unsigned count(BasicBlock *BB) { return Status.count(BB); }
      // Invoked after we have visited all children of a node.
      void completed(BasicBlock *BB) { Status[BB] = false; }
      // Return true if \p BB is currently on the active stack
      // of ancestors.
      bool onStack(BasicBlock *BB) {
        auto Iter = Status.find(BB);
        return Iter != Status.end() && Iter->second;
      }
    } State(F.size());

    // Iterate over blocks in depth-first pre-order and
    // treat all edges to a block already seen as loop back edges
    // and mark the branch live it if there is a back edge.
    typedef df_iterator<BasicBlock *, DFState, true> Iterator;
    auto Iter = Iterator::begin(&F.getEntryBlock(), State);
    auto End = Iterator::end(&F.getEntryBlock(), State);
    for (; Iter != End; ++Iter) {
      auto *BB = *Iter;

      for (auto Succ : successors(BB))
        if (State.onStack(Succ)) {
          // back edge....
        }
    }
}

[david2050]

Any additional suggestions for this diff?

llvm/trunk/include/llvm/ADT/DepthFirstIterator.h

 template<class SetType>
 class df_iterator_storage<SetType, true> {
 public:
   df_iterator_storage(SetType &VSet) : Visited(VSet) {}
   df_iterator_storage(const df_iterator_storage &S) : Visited(S.Visited) {}
   SetType &Visited;
 };
 
+// The visited stated for the iteration is a simple set augmented with
+// one more method, completed, which is invoked when all children of a
+// node have been processed. It is intended to distinguish of back and
+// cross edges in the spanning tree but is not used in the common case.␓
+template <typename NodeRef, unsigned SmallSize=8>
+struct df_iterator_default_set : public llvm::SmallPtrSet<NodeRef, SmallSize> {
+  typedef llvm::SmallPtrSet<NodeRef, SmallSize>  BaseSet;
+  typedef typename BaseSet::iterator iterator;
+  std::pair<iterator,bool> insert(NodeRef N) { return BaseSet::insert(N) ; }
+  template <typename IterT>
+  void insert(IterT Begin, IterT End) { BaseSet::insert(Begin,End); }
+
+  void completed(NodeRef) { }
+};
+
 // Generic Depth First Iterator
 template <class GraphT,
           class SetType =
-              llvm::SmallPtrSet<typename GraphTraits<GraphT>::NodeRef, 8>,
+              df_iterator_default_set<typename GraphTraits<GraphT>::NodeRef>,
           bool ExtStorage = false, class GT = GraphTraits<GraphT>>
 class df_iterator
     : public std::iterator<std::forward_iterator_tag, typename GT::NodeRef>,
       public df_iterator_storage<SetType, ExtStorage> {
   typedef std::iterator<std::forward_iterator_tag, typename GT::NodeRef> super;
 
   typedef typename GT::NodeRef NodeRef;
   typedef typename GT::ChildIteratorType ChildItTy;
inline df_iterator(NodeRef Node) {
     this->Visited.insert(Node);
     VisitStack.push_back(StackElement(Node, None));
   }
   inline df_iterator() {
     // End is when stack is empty
   }
   inline df_iterator(NodeRef Node, SetType &S)
       : df_iterator_storage<SetType, ExtStorage>(S) {
-    if (!S.count(Node)) {
+    if (this->Visited.insert(Node).second)
       VisitStack.push_back(StackElement(Node, None));
-      this->Visited.insert(Node);
-    }
   }
   inline df_iterator(SetType &S)
     : df_iterator_storage<SetType, ExtStorage>(S) {
     // End is when stack is empty
   }
 
   inline void toNext() {
     do {
do {
         NodeRef Next = *(*Opt)++;
         // Has our next sibling been visited?
         if (this->Visited.insert(Next).second) {
           // No, do it now.
           VisitStack.push_back(StackElement(Next, None));
           return;
         }
       }
+      this->Visited.completed(Node);
       
       // Oops, ran out of successors... go up a level on the stack.
       VisitStack.pop_back();
     } while (!VisitStack.empty());
   }
 
 public:
   typedef typename super::pointer pointer;
 

 template <class T, class SetTy>
 iterator_range<df_ext_iterator<T, SetTy>> depth_first_ext(const T& G,
                                                           SetTy &S) {
   return make_range(df_ext_begin(G, S), df_ext_end(G, S));
 }
 
 // Provide global definitions of inverse depth first iterators...
 template <class T,
-          class SetTy = llvm::SmallPtrSet<typename GraphTraits<T>::NodeRef, 8>,
+          class SetTy =
+              df_iterator_default_set<typename GraphTraits<T>::NodeRef>,
           bool External = false>
 struct idf_iterator : public df_iterator<Inverse<T>, SetTy, External> {
   idf_iterator(const df_iterator<Inverse<T>, SetTy, External> &V)
     : df_iterator<Inverse<T>, SetTy, External>(V) {}
 };
 
 template <class T>
 idf_iterator<T> idf_begin(const T& G) {

llvm/trunk/include/llvm/Analysis/LoopInfoImpl.h

 template<class BlockT, class LoopT>
 void LoopBase<BlockT, LoopT>::verifyLoop() const {
 #ifndef NDEBUG
   assert(!Blocks.empty() && "Loop header is missing");
 
   // Setup for using a depth-first iterator to visit every block in the loop.
   SmallVector<BlockT*, 8> ExitBBs;
   getExitBlocks(ExitBBs);
-  llvm::SmallPtrSet<BlockT*, 8> VisitSet;
+  df_iterator_default_set<BlockT*> VisitSet;
   VisitSet.insert(ExitBBs.begin(), ExitBBs.end());
-  df_ext_iterator<BlockT*, llvm::SmallPtrSet<BlockT*, 8> >
+  df_ext_iterator<BlockT*, df_iterator_default_set<BlockT*>>
     BI = df_ext_begin(getHeader(), VisitSet),
     BE = df_ext_end(getHeader(), VisitSet);
 
   // Keep track of the number of BBs visited.
   unsigned NumVisited = 0;
 
   // Check the individual blocks.
   for ( ; BI != BE; ++BI) {

llvm/trunk/include/llvm/Analysis/RegionInfo.h

#endif
   //@}
 
   /// @name Element Iterators
   ///
   /// These iterators iterate over all BasicBlock and subregion RegionNodes that
   /// are direct children of this Region. It does not iterate over any
   /// RegionNodes that are also element of a subregion of this Region.
   //@{
-  typedef df_iterator<RegionNodeT *, SmallPtrSet<RegionNodeT *, 8>, false,
-                      GraphTraits<RegionNodeT *>> element_iterator;
-
-  typedef df_iterator<const RegionNodeT *, SmallPtrSet<const RegionNodeT *, 8>,
-                      false,
-                      GraphTraits<const RegionNodeT *>> const_element_iterator;
+  typedef df_iterator<RegionNodeT *, df_iterator_default_set<RegionNodeT *>,
+                      false, GraphTraits<RegionNodeT *>>
+      element_iterator;
+
+  typedef df_iterator<const RegionNodeT *,
+                      df_iterator_default_set<const RegionNodeT *>, false,
+                      GraphTraits<const RegionNodeT *>>
+      const_element_iterator;
 
   element_iterator element_begin();
   element_iterator element_end();
   iterator_range<element_iterator> elements() {
     return make_range(element_begin(), element_end());
   }
 
   const_element_iterator element_begin() const;

llvm/trunk/include/llvm/Analysis/RegionIterator.h

template <> struct GraphTraits<RegionT *> : public GraphTraits<NodeT *> {    \
     }                                                                          \
     static nodes_iterator nodes_end(RegionT *R) {                              \
       return nodes_iterator::end(getEntryNode(R));                             \
     }                                                                          \
   };                                                                           \
   template <>                                                                  \
   struct GraphTraits<FlatIt<RegionT *>>                                        \
       : public GraphTraits<FlatIt<NodeT *>> {                                  \
-    typedef df_iterator<NodeRef, SmallPtrSet<NodeRef, 8>, false,               \
+    typedef df_iterator<NodeRef, df_iterator_default_set<NodeRef>, false,      \
                         GraphTraits<FlatIt<NodeRef>>>                          \
         nodes_iterator;                                                        \
     static NodeRef getEntryNode(RegionT *R) {                                  \
       return R->getBBNode(R->getEntry());                                      \
     }                                                                          \
     static nodes_iterator nodes_begin(RegionT *R) {                            \
       return nodes_iterator::begin(getEntryNode(R));                           \
     }                                                                          \
     static nodes_iterator nodes_end(RegionT *R) {                              \
       return nodes_iterator::end(getEntryNode(R));                             \
     }                                                                          \
   }
 
 RegionNodeGraphTraits(RegionNode, BasicBlock, Region);
 RegionNodeGraphTraits(const RegionNode, BasicBlock, Region);
 
 RegionGraphTraits(Region, RegionNode);
 RegionGraphTraits(const Region, const RegionNode);
 
 template <> struct GraphTraits<RegionInfo*>
   : public GraphTraits<FlatIt<RegionNode*> > {
-  typedef df_iterator<NodeRef, SmallPtrSet<NodeRef, 8>, false,
+  typedef df_iterator<NodeRef, df_iterator_default_set<NodeRef>, false,
                       GraphTraits<FlatIt<NodeRef>>>
       nodes_iterator;
 
   static NodeRef getEntryNode(RegionInfo *RI) {
     return GraphTraits<FlatIt<Region*> >::getEntryNode(RI->getTopLevelRegion());
   }
   static nodes_iterator nodes_begin(RegionInfo* RI) {
     return nodes_iterator::begin(getEntryNode(RI));
   }
   static nodes_iterator nodes_end(RegionInfo *RI) {
     return nodes_iterator::end(getEntryNode(RI));
   }
 };
 
 template <> struct GraphTraits<RegionInfoPass*>
   : public GraphTraits<RegionInfo *> {
-  typedef df_iterator<NodeRef, SmallPtrSet<NodeRef, 8>, false,
+  typedef df_iterator<NodeRef, df_iterator_default_set<NodeRef>, false,
                       GraphTraits<FlatIt<NodeRef>>>
       nodes_iterator;
 
   static NodeRef getEntryNode(RegionInfoPass *RI) {
     return GraphTraits<RegionInfo*>::getEntryNode(&RI->getRegionInfo());
   }
   static nodes_iterator nodes_begin(RegionInfoPass* RI) {
     return GraphTraits<RegionInfo*>::nodes_begin(&RI->getRegionInfo());

llvm/trunk/include/llvm/CodeGen/MachineRegionInfo.h

 RegionNodeGraphTraits(MachineRegionNode, MachineBasicBlock, MachineRegion);
 RegionNodeGraphTraits(const MachineRegionNode, MachineBasicBlock, MachineRegion);
 
 RegionGraphTraits(MachineRegion, MachineRegionNode);
 RegionGraphTraits(const MachineRegion, const MachineRegionNode);
 
 template <> struct GraphTraits<MachineRegionInfo*>
   : public GraphTraits<FlatIt<MachineRegionNode*> > {
-  typedef df_iterator<NodeRef, SmallPtrSet<NodeRef, 8>, false,
+  typedef df_iterator<NodeRef, df_iterator_default_set<NodeRef>, false,
                       GraphTraits<FlatIt<NodeRef>>>
       nodes_iterator;
 
   static NodeRef getEntryNode(MachineRegionInfo *RI) {
     return GraphTraits<FlatIt<MachineRegion*> >::getEntryNode(RI->getTopLevelRegion());
   }
   static nodes_iterator nodes_begin(MachineRegionInfo* RI) {
     return nodes_iterator::begin(getEntryNode(RI));
   }
   static nodes_iterator nodes_end(MachineRegionInfo *RI) {
     return nodes_iterator::end(getEntryNode(RI));
   }
 };
 
 template <> struct GraphTraits<MachineRegionInfoPass*>
   : public GraphTraits<MachineRegionInfo *> {
-  typedef df_iterator<NodeRef, SmallPtrSet<NodeRef, 8>, false,
+  typedef df_iterator<NodeRef, df_iterator_default_set<NodeRef>, false,
                       GraphTraits<FlatIt<NodeRef>>>
       nodes_iterator;
 
   static NodeRef getEntryNode(MachineRegionInfoPass *RI) {
     return GraphTraits<MachineRegionInfo*>::getEntryNode(&RI->getRegionInfo());
   }
   static nodes_iterator nodes_begin(MachineRegionInfoPass* RI) {
     return GraphTraits<MachineRegionInfo*>::nodes_begin(&RI->getRegionInfo());

llvm/trunk/include/llvm/IR/Dominators.h

 
 //===-------------------------------------
 // DominatorTree GraphTraits specializations so the DominatorTree can be
 // iterable by generic graph iterators.
 
 template <class Node, class ChildIterator> struct DomTreeGraphTraitsBase {
   typedef Node *NodeRef;
   typedef ChildIterator ChildIteratorType;
-  typedef df_iterator<Node *, SmallPtrSet<NodeRef, 8>> nodes_iterator;
+  typedef df_iterator<Node *, df_iterator_default_set<Node*>> nodes_iterator;
 
   static NodeRef getEntryNode(NodeRef N) { return N; }
   static ChildIteratorType child_begin(NodeRef N) { return N->begin(); }
   static ChildIteratorType child_end(NodeRef N) { return N->end(); }
 
   static nodes_iterator nodes_begin(NodeRef N) {
     return df_begin(getEntryNode(N));
   }

llvm/trunk/lib/CodeGen/LiveIntervalAnalysis.cpp

void LiveIntervals::pruneValue(LiveRange &LR, SlotIndex Kill,
 
   // VNI is live out of KillMBB.
   LR.removeSegment(Kill, MBBEnd);
   if (EndPoints) EndPoints->push_back(MBBEnd);
 
   // Find all blocks that are reachable from KillMBB without leaving VNI's live
   // range. It is possible that KillMBB itself is reachable, so start a DFS
   // from each successor.
-  typedef SmallPtrSet<MachineBasicBlock*, 9> VisitedTy;
+  typedef df_iterator_default_set<MachineBasicBlock*,9> VisitedTy;
   VisitedTy Visited;
   for (MachineBasicBlock::succ_iterator
        SuccI = KillMBB->succ_begin(), SuccE = KillMBB->succ_end();
        SuccI != SuccE; ++SuccI) {
     for (df_ext_iterator<MachineBasicBlock*, VisitedTy>
          I = df_ext_begin(*SuccI, Visited), E = df_ext_end(*SuccI, Visited);
          I != E;) {
       MachineBasicBlock *MBB = *I;

llvm/trunk/lib/CodeGen/LiveVariables.cpp

bool LiveVariables::runOnMachineFunction(MachineFunction &mf) {
 
   analyzePHINodes(mf);
 
   // Calculate live variable information in depth first order on the CFG of the
   // function.  This guarantees that we will see the definition of a virtual
   // register before its uses due to dominance properties of SSA (except for PHI
   // nodes, which are treated as a special case).
   MachineBasicBlock *Entry = &MF->front();
-  SmallPtrSet<MachineBasicBlock*,16> Visited;
+  df_iterator_default_set<MachineBasicBlock*,16> Visited;
 
   for (MachineBasicBlock *MBB : depth_first_ext(Entry, Visited)) {
     runOnBlock(MBB, NumRegs);
 
     PhysRegDef.assign(NumRegs, nullptr);
     PhysRegUse.assign(NumRegs, nullptr);
   }
 

llvm/trunk/lib/CodeGen/MachineVerifier.cpp

 /// by a FrameDestroy <n>, stack adjustments are identical on all
 /// CFG edges to a merge point, and frame is destroyed at end of a return block.
 void MachineVerifier::verifyStackFrame() {
   unsigned FrameSetupOpcode   = TII->getCallFrameSetupOpcode();
   unsigned FrameDestroyOpcode = TII->getCallFrameDestroyOpcode();
 
   SmallVector<StackStateOfBB, 8> SPState;
   SPState.resize(MF->getNumBlockIDs());
-  SmallPtrSet<const MachineBasicBlock*, 8> Reachable;
+  df_iterator_default_set<const MachineBasicBlock*> Reachable;
 
   // Visit the MBBs in DFS order.
   for (df_ext_iterator<const MachineFunction*,
-                       SmallPtrSet<const MachineBasicBlock*, 8> >
+                       df_iterator_default_set<const MachineBasicBlock*> >
        DFI = df_ext_begin(MF, Reachable), DFE = df_ext_end(MF, Reachable);
        DFI != DFE; ++DFI) {
     const MachineBasicBlock *MBB = *DFI;
 
     StackStateOfBB BBState;
     // Check the exit state of the DFS stack predecessor.
     if (DFI.getPathLength() >= 2) {
       const MachineBasicBlock *StackPred = DFI.getPath(DFI.getPathLength() - 2);

llvm/trunk/lib/CodeGen/PrologEpilogInserter.cpp

 ///
 void PEI::replaceFrameIndices(MachineFunction &Fn) {
   const TargetFrameLowering &TFI = *Fn.getSubtarget().getFrameLowering();
   if (!TFI.needsFrameIndexResolution(Fn)) return;
 
   // Store SPAdj at exit of a basic block.
   SmallVector<int, 8> SPState;
   SPState.resize(Fn.getNumBlockIDs());
-  SmallPtrSet<MachineBasicBlock*, 8> Reachable;
+  df_iterator_default_set<MachineBasicBlock*> Reachable;
 
   // Iterate over the reachable blocks in DFS order.
   for (auto DFI = df_ext_begin(&Fn, Reachable), DFE = df_ext_end(&Fn, Reachable);
        DFI != DFE; ++DFI) {
     int SPAdj = 0;
     // Check the exit state of the DFS stack predecessor.
     if (DFI.getPathLength() >= 2) {
       MachineBasicBlock *StackPred = DFI.getPath(DFI.getPathLength() - 2);

llvm/trunk/lib/CodeGen/UnreachableBlockElim.cpp

 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Pass.h"
 #include "llvm/Target/TargetInstrInfo.h"
 using namespace llvm;
 
 static bool eliminateUnreachableBlock(Function &F) {
-  SmallPtrSet<BasicBlock*, 8> Reachable;
+  df_iterator_default_set<BasicBlock*> Reachable;
 
   // Mark all reachable blocks.
   for (BasicBlock *BB : depth_first_ext(&F, Reachable))
     (void)BB/* Mark all reachable blocks */;
 
   // Loop over all dead blocks, remembering them and deleting all instructions
   // in them.
   std::vector<BasicBlock*> DeadBlocks;

 
 void UnreachableMachineBlockElim::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addPreserved<MachineLoopInfo>();
   AU.addPreserved<MachineDominatorTree>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 bool UnreachableMachineBlockElim::runOnMachineFunction(MachineFunction &F) {
-  SmallPtrSet<MachineBasicBlock*, 8> Reachable;
+  df_iterator_default_set<MachineBasicBlock*> Reachable;
   bool ModifiedPHI = false;
 
   MMI = getAnalysisIfAvailable<MachineModuleInfo>();
   MachineDominatorTree *MDT = getAnalysisIfAvailable<MachineDominatorTree>();
   MachineLoopInfo *MLI = getAnalysisIfAvailable<MachineLoopInfo>();
 
   // Mark all reachable blocks.
   for (MachineBasicBlock *BB : depth_first_ext(&F, Reachable))

llvm/trunk/lib/Target/X86/X86FloatingPoint.cpp

bool FPS::runOnMachineFunction(MachineFunction &MF) {
 
   // Prepare cross-MBB liveness.
   bundleCFG(MF);
 
   StackTop = 0;
 
   // Process the function in depth first order so that we process at least one
   // of the predecessors for every reachable block in the function.
-  SmallPtrSet<MachineBasicBlock*, 8> Processed;
+  df_iterator_default_set<MachineBasicBlock*> Processed;
   MachineBasicBlock *Entry = &MF.front();
 
   bool Changed = false;
   for (MachineBasicBlock *BB : depth_first_ext(Entry, Processed))
     Changed |= processBasicBlock(MF, *BB);
 
   // Process any unreachable blocks in arbitrary order now.
   if (MF.size() != Processed.size())

llvm/trunk/lib/Transforms/IPO/ArgumentPromotion.cpp

static bool isSafeToPromoteArgument(Argument *Arg, bool isByValOrInAlloca,
 
   // Okay, now we know that the argument is only used by load instructions and
   // it is safe to unconditionally perform all of them. Use alias analysis to
   // check to see if the pointer is guaranteed to not be modified from entry of
   // the function to each of the load instructions.
 
   // Because there could be several/many load instructions, remember which
   // blocks we know to be transparent to the load.
-  SmallPtrSet<BasicBlock*, 16> TranspBlocks;
+  df_iterator_default_set<BasicBlock*, 16> TranspBlocks;
 
   for (LoadInst *Load : Loads) {
     // Check to see if the load is invalidated from the start of the block to
     // the load itself.
     BasicBlock *BB = Load->getParent();
 
     MemoryLocation Loc = MemoryLocation::get(Load);
     if (AAR.canInstructionRangeModRef(BB->front(), *Load, Loc, MRI_Mod))

llvm/trunk/unittests/ADT/DepthFirstIteratorTest.cpp

template <typename T> struct CountedSet {
   int InsertVisited = 0;
 
   std::pair<iterator, bool> insert(const T &Item) {
     InsertVisited++;
     return S.insert(Item);
   }
 
   size_t count(const T &Item) const { return S.count(Item); }
+  
+  void completed(T) { }
 };
 
 template <typename T> class df_iterator_storage<CountedSet<T>, true> {
 public:
   df_iterator_storage(CountedSet<T> &VSet) : Visited(VSet) {}
 
   CountedSet<T> &Visited;
 };