10b57cec5SDimitry Andric //===--------- LoopIterator.h - Iterate over loop blocks --------*- C++ -*-===//
20b57cec5SDimitry Andric //
30b57cec5SDimitry Andric // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
40b57cec5SDimitry Andric // See https://llvm.org/LICENSE.txt for license information.
50b57cec5SDimitry Andric // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
60b57cec5SDimitry Andric //
70b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
80b57cec5SDimitry Andric // This file defines iterators to visit the basic blocks within a loop.
90b57cec5SDimitry Andric //
100b57cec5SDimitry Andric // These iterators currently visit blocks within subloops as well.
110b57cec5SDimitry Andric // Unfortunately we have no efficient way of summarizing loop exits which would
120b57cec5SDimitry Andric // allow skipping subloops during traversal.
130b57cec5SDimitry Andric //
140b57cec5SDimitry Andric // If you want to visit all blocks in a loop and don't need an ordered traveral,
150b57cec5SDimitry Andric // use Loop::block_begin() instead.
160b57cec5SDimitry Andric //
170b57cec5SDimitry Andric // This is intentionally designed to work with ill-formed loops in which the
180b57cec5SDimitry Andric // backedge has been deleted. The only prerequisite is that all blocks
190b57cec5SDimitry Andric // contained within the loop according to the most recent LoopInfo analysis are
200b57cec5SDimitry Andric // reachable from the loop header.
210b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
220b57cec5SDimitry Andric 
230b57cec5SDimitry Andric #ifndef LLVM_ANALYSIS_LOOPITERATOR_H
240b57cec5SDimitry Andric #define LLVM_ANALYSIS_LOOPITERATOR_H
250b57cec5SDimitry Andric 
260b57cec5SDimitry Andric #include "llvm/ADT/PostOrderIterator.h"
270b57cec5SDimitry Andric #include "llvm/Analysis/LoopInfo.h"
280b57cec5SDimitry Andric 
290b57cec5SDimitry Andric namespace llvm {
300b57cec5SDimitry Andric 
310b57cec5SDimitry Andric class LoopBlocksTraversal;
320b57cec5SDimitry Andric 
330b57cec5SDimitry Andric // A traits type that is intended to be used in graph algorithms. The graph
340b57cec5SDimitry Andric // traits starts at the loop header, and traverses the BasicBlocks that are in
350b57cec5SDimitry Andric // the loop body, but not the loop header. Since the loop header is skipped,
360b57cec5SDimitry Andric // the back edges are excluded.
370b57cec5SDimitry Andric //
380b57cec5SDimitry Andric // TODO: Explore the possibility to implement LoopBlocksTraversal in terms of
390b57cec5SDimitry Andric //       LoopBodyTraits, so that insertEdge doesn't have to be specialized.
400b57cec5SDimitry Andric struct LoopBodyTraits {
410b57cec5SDimitry Andric   using NodeRef = std::pair<const Loop *, BasicBlock *>;
420b57cec5SDimitry Andric 
430b57cec5SDimitry Andric   // This wraps a const Loop * into the iterator, so we know which edges to
440b57cec5SDimitry Andric   // filter out.
450b57cec5SDimitry Andric   class WrappedSuccIterator
460b57cec5SDimitry Andric       : public iterator_adaptor_base<
470b57cec5SDimitry Andric             WrappedSuccIterator, succ_iterator,
480b57cec5SDimitry Andric             typename std::iterator_traits<succ_iterator>::iterator_category,
490b57cec5SDimitry Andric             NodeRef, std::ptrdiff_t, NodeRef *, NodeRef> {
500b57cec5SDimitry Andric     using BaseT = iterator_adaptor_base<
510b57cec5SDimitry Andric         WrappedSuccIterator, succ_iterator,
520b57cec5SDimitry Andric         typename std::iterator_traits<succ_iterator>::iterator_category,
530b57cec5SDimitry Andric         NodeRef, std::ptrdiff_t, NodeRef *, NodeRef>;
540b57cec5SDimitry Andric 
550b57cec5SDimitry Andric     const Loop *L;
560b57cec5SDimitry Andric 
570b57cec5SDimitry Andric   public:
WrappedSuccIteratorLoopBodyTraits580b57cec5SDimitry Andric     WrappedSuccIterator(succ_iterator Begin, const Loop *L)
590b57cec5SDimitry Andric         : BaseT(Begin), L(L) {}
600b57cec5SDimitry Andric 
610b57cec5SDimitry Andric     NodeRef operator*() const { return {L, *I}; }
620b57cec5SDimitry Andric   };
630b57cec5SDimitry Andric 
640b57cec5SDimitry Andric   struct LoopBodyFilter {
operatorLoopBodyTraits::LoopBodyFilter650b57cec5SDimitry Andric     bool operator()(NodeRef N) const {
660b57cec5SDimitry Andric       const Loop *L = N.first;
670b57cec5SDimitry Andric       return N.second != L->getHeader() && L->contains(N.second);
680b57cec5SDimitry Andric     }
690b57cec5SDimitry Andric   };
700b57cec5SDimitry Andric 
710b57cec5SDimitry Andric   using ChildIteratorType =
720b57cec5SDimitry Andric       filter_iterator<WrappedSuccIterator, LoopBodyFilter>;
730b57cec5SDimitry Andric 
getEntryNodeLoopBodyTraits740b57cec5SDimitry Andric   static NodeRef getEntryNode(const Loop &G) { return {&G, G.getHeader()}; }
750b57cec5SDimitry Andric 
child_beginLoopBodyTraits760b57cec5SDimitry Andric   static ChildIteratorType child_begin(NodeRef Node) {
770b57cec5SDimitry Andric     return make_filter_range(make_range<WrappedSuccIterator>(
780b57cec5SDimitry Andric                                  {succ_begin(Node.second), Node.first},
790b57cec5SDimitry Andric                                  {succ_end(Node.second), Node.first}),
800b57cec5SDimitry Andric                              LoopBodyFilter{})
810b57cec5SDimitry Andric         .begin();
820b57cec5SDimitry Andric   }
830b57cec5SDimitry Andric 
child_endLoopBodyTraits840b57cec5SDimitry Andric   static ChildIteratorType child_end(NodeRef Node) {
850b57cec5SDimitry Andric     return make_filter_range(make_range<WrappedSuccIterator>(
860b57cec5SDimitry Andric                                  {succ_begin(Node.second), Node.first},
870b57cec5SDimitry Andric                                  {succ_end(Node.second), Node.first}),
880b57cec5SDimitry Andric                              LoopBodyFilter{})
890b57cec5SDimitry Andric         .end();
900b57cec5SDimitry Andric   }
910b57cec5SDimitry Andric };
920b57cec5SDimitry Andric 
930b57cec5SDimitry Andric /// Store the result of a depth first search within basic blocks contained by a
940b57cec5SDimitry Andric /// single loop.
950b57cec5SDimitry Andric ///
960b57cec5SDimitry Andric /// TODO: This could be generalized for any CFG region, or the entire CFG.
970b57cec5SDimitry Andric class LoopBlocksDFS {
980b57cec5SDimitry Andric public:
990b57cec5SDimitry Andric   /// Postorder list iterators.
1000b57cec5SDimitry Andric   typedef std::vector<BasicBlock*>::const_iterator POIterator;
1010b57cec5SDimitry Andric   typedef std::vector<BasicBlock*>::const_reverse_iterator RPOIterator;
1020b57cec5SDimitry Andric 
1030b57cec5SDimitry Andric   friend class LoopBlocksTraversal;
1040b57cec5SDimitry Andric 
1050b57cec5SDimitry Andric private:
1060b57cec5SDimitry Andric   Loop *L;
1070b57cec5SDimitry Andric 
1080b57cec5SDimitry Andric   /// Map each block to its postorder number. A block is only mapped after it is
1090b57cec5SDimitry Andric   /// preorder visited by DFS. It's postorder number is initially zero and set
1100b57cec5SDimitry Andric   /// to nonzero after it is finished by postorder traversal.
1110b57cec5SDimitry Andric   DenseMap<BasicBlock*, unsigned> PostNumbers;
1120b57cec5SDimitry Andric   std::vector<BasicBlock*> PostBlocks;
1130b57cec5SDimitry Andric 
1140b57cec5SDimitry Andric public:
LoopBlocksDFS(Loop * Container)1150b57cec5SDimitry Andric   LoopBlocksDFS(Loop *Container) :
1160b57cec5SDimitry Andric     L(Container), PostNumbers(NextPowerOf2(Container->getNumBlocks())) {
1170b57cec5SDimitry Andric     PostBlocks.reserve(Container->getNumBlocks());
1180b57cec5SDimitry Andric   }
1190b57cec5SDimitry Andric 
getLoop()1200b57cec5SDimitry Andric   Loop *getLoop() const { return L; }
1210b57cec5SDimitry Andric 
1220b57cec5SDimitry Andric   /// Traverse the loop blocks and store the DFS result.
1235f757f3fSDimitry Andric   void perform(const LoopInfo *LI);
1240b57cec5SDimitry Andric 
1250b57cec5SDimitry Andric   /// Return true if postorder numbers are assigned to all loop blocks.
isComplete()1260b57cec5SDimitry Andric   bool isComplete() const { return PostBlocks.size() == L->getNumBlocks(); }
1270b57cec5SDimitry Andric 
1280b57cec5SDimitry Andric   /// Iterate over the cached postorder blocks.
beginPostorder()1290b57cec5SDimitry Andric   POIterator beginPostorder() const {
1300b57cec5SDimitry Andric     assert(isComplete() && "bad loop DFS");
1310b57cec5SDimitry Andric     return PostBlocks.begin();
1320b57cec5SDimitry Andric   }
endPostorder()1330b57cec5SDimitry Andric   POIterator endPostorder() const { return PostBlocks.end(); }
1340b57cec5SDimitry Andric 
1350b57cec5SDimitry Andric   /// Reverse iterate over the cached postorder blocks.
beginRPO()1360b57cec5SDimitry Andric   RPOIterator beginRPO() const {
1370b57cec5SDimitry Andric     assert(isComplete() && "bad loop DFS");
1380b57cec5SDimitry Andric     return PostBlocks.rbegin();
1390b57cec5SDimitry Andric   }
endRPO()1400b57cec5SDimitry Andric   RPOIterator endRPO() const { return PostBlocks.rend(); }
1410b57cec5SDimitry Andric 
1420b57cec5SDimitry Andric   /// Return true if this block has been preorder visited.
hasPreorder(BasicBlock * BB)1430b57cec5SDimitry Andric   bool hasPreorder(BasicBlock *BB) const { return PostNumbers.count(BB); }
1440b57cec5SDimitry Andric 
1450b57cec5SDimitry Andric   /// Return true if this block has a postorder number.
hasPostorder(BasicBlock * BB)1460b57cec5SDimitry Andric   bool hasPostorder(BasicBlock *BB) const {
1470b57cec5SDimitry Andric     DenseMap<BasicBlock*, unsigned>::const_iterator I = PostNumbers.find(BB);
1480b57cec5SDimitry Andric     return I != PostNumbers.end() && I->second;
1490b57cec5SDimitry Andric   }
1500b57cec5SDimitry Andric 
1510b57cec5SDimitry Andric   /// Get a block's postorder number.
getPostorder(BasicBlock * BB)1520b57cec5SDimitry Andric   unsigned getPostorder(BasicBlock *BB) const {
1530b57cec5SDimitry Andric     DenseMap<BasicBlock*, unsigned>::const_iterator I = PostNumbers.find(BB);
1540b57cec5SDimitry Andric     assert(I != PostNumbers.end() && "block not visited by DFS");
1550b57cec5SDimitry Andric     assert(I->second && "block not finished by DFS");
1560b57cec5SDimitry Andric     return I->second;
1570b57cec5SDimitry Andric   }
1580b57cec5SDimitry Andric 
1590b57cec5SDimitry Andric   /// Get a block's reverse postorder number.
getRPO(BasicBlock * BB)1600b57cec5SDimitry Andric   unsigned getRPO(BasicBlock *BB) const {
1610b57cec5SDimitry Andric     return 1 + PostBlocks.size() - getPostorder(BB);
1620b57cec5SDimitry Andric   }
1630b57cec5SDimitry Andric 
clear()1640b57cec5SDimitry Andric   void clear() {
1650b57cec5SDimitry Andric     PostNumbers.clear();
1660b57cec5SDimitry Andric     PostBlocks.clear();
1670b57cec5SDimitry Andric   }
1680b57cec5SDimitry Andric };
1690b57cec5SDimitry Andric 
1700b57cec5SDimitry Andric /// Wrapper class to LoopBlocksDFS that provides a standard begin()/end()
1710b57cec5SDimitry Andric /// interface for the DFS reverse post-order traversal of blocks in a loop body.
1720b57cec5SDimitry Andric class LoopBlocksRPO {
1730b57cec5SDimitry Andric private:
1740b57cec5SDimitry Andric   LoopBlocksDFS DFS;
1750b57cec5SDimitry Andric 
1760b57cec5SDimitry Andric public:
LoopBlocksRPO(Loop * Container)1770b57cec5SDimitry Andric   LoopBlocksRPO(Loop *Container) : DFS(Container) {}
1780b57cec5SDimitry Andric 
1790b57cec5SDimitry Andric   /// Traverse the loop blocks and store the DFS result.
perform(const LoopInfo * LI)1805f757f3fSDimitry Andric   void perform(const LoopInfo *LI) {
1810b57cec5SDimitry Andric     DFS.perform(LI);
1820b57cec5SDimitry Andric   }
1830b57cec5SDimitry Andric 
1840b57cec5SDimitry Andric   /// Reverse iterate over the cached postorder blocks.
begin()1850b57cec5SDimitry Andric   LoopBlocksDFS::RPOIterator begin() const { return DFS.beginRPO(); }
end()1860b57cec5SDimitry Andric   LoopBlocksDFS::RPOIterator end() const { return DFS.endRPO(); }
1870b57cec5SDimitry Andric };
1880b57cec5SDimitry Andric 
1890b57cec5SDimitry Andric /// Specialize po_iterator_storage to record postorder numbers.
1900b57cec5SDimitry Andric template<> class po_iterator_storage<LoopBlocksTraversal, true> {
1910b57cec5SDimitry Andric   LoopBlocksTraversal &LBT;
1920b57cec5SDimitry Andric public:
po_iterator_storage(LoopBlocksTraversal & lbs)1930b57cec5SDimitry Andric   po_iterator_storage(LoopBlocksTraversal &lbs) : LBT(lbs) {}
1940b57cec5SDimitry Andric   // These functions are defined below.
195bdd1243dSDimitry Andric   bool insertEdge(std::optional<BasicBlock *> From, BasicBlock *To);
1960b57cec5SDimitry Andric   void finishPostorder(BasicBlock *BB);
1970b57cec5SDimitry Andric };
1980b57cec5SDimitry Andric 
1990b57cec5SDimitry Andric /// Traverse the blocks in a loop using a depth-first search.
2000b57cec5SDimitry Andric class LoopBlocksTraversal {
2010b57cec5SDimitry Andric public:
2020b57cec5SDimitry Andric   /// Graph traversal iterator.
2030b57cec5SDimitry Andric   typedef po_iterator<BasicBlock*, LoopBlocksTraversal, true> POTIterator;
2040b57cec5SDimitry Andric 
2050b57cec5SDimitry Andric private:
2060b57cec5SDimitry Andric   LoopBlocksDFS &DFS;
2075f757f3fSDimitry Andric   const LoopInfo *LI;
2080b57cec5SDimitry Andric 
2090b57cec5SDimitry Andric public:
LoopBlocksTraversal(LoopBlocksDFS & Storage,const LoopInfo * LInfo)2105f757f3fSDimitry Andric   LoopBlocksTraversal(LoopBlocksDFS &Storage, const LoopInfo *LInfo) :
2110b57cec5SDimitry Andric     DFS(Storage), LI(LInfo) {}
2120b57cec5SDimitry Andric 
2130b57cec5SDimitry Andric   /// Postorder traversal over the graph. This only needs to be done once.
2140b57cec5SDimitry Andric   /// po_iterator "automatically" calls back to visitPreorder and
2150b57cec5SDimitry Andric   /// finishPostorder to record the DFS result.
begin()2160b57cec5SDimitry Andric   POTIterator begin() {
2170b57cec5SDimitry Andric     assert(DFS.PostBlocks.empty() && "Need clear DFS result before traversing");
2180b57cec5SDimitry Andric     assert(DFS.L->getNumBlocks() && "po_iterator cannot handle an empty graph");
2190b57cec5SDimitry Andric     return po_ext_begin(DFS.L->getHeader(), *this);
2200b57cec5SDimitry Andric   }
end()2210b57cec5SDimitry Andric   POTIterator end() {
2220b57cec5SDimitry Andric     // po_ext_end interface requires a basic block, but ignores its value.
2230b57cec5SDimitry Andric     return po_ext_end(DFS.L->getHeader(), *this);
2240b57cec5SDimitry Andric   }
2250b57cec5SDimitry Andric 
2260b57cec5SDimitry Andric   /// Called by po_iterator upon reaching a block via a CFG edge. If this block
2270b57cec5SDimitry Andric   /// is contained in the loop and has not been visited, then mark it preorder
2280b57cec5SDimitry Andric   /// visited and return true.
2290b57cec5SDimitry Andric   ///
2300b57cec5SDimitry Andric   /// TODO: If anyone is interested, we could record preorder numbers here.
visitPreorder(BasicBlock * BB)2310b57cec5SDimitry Andric   bool visitPreorder(BasicBlock *BB) {
2320b57cec5SDimitry Andric     if (!DFS.L->contains(LI->getLoopFor(BB)))
2330b57cec5SDimitry Andric       return false;
2340b57cec5SDimitry Andric 
2350b57cec5SDimitry Andric     return DFS.PostNumbers.insert(std::make_pair(BB, 0)).second;
2360b57cec5SDimitry Andric   }
2370b57cec5SDimitry Andric 
2380b57cec5SDimitry Andric   /// Called by po_iterator each time it advances, indicating a block's
2390b57cec5SDimitry Andric   /// postorder.
finishPostorder(BasicBlock * BB)2400b57cec5SDimitry Andric   void finishPostorder(BasicBlock *BB) {
2410b57cec5SDimitry Andric     assert(DFS.PostNumbers.count(BB) && "Loop DFS skipped preorder");
2420b57cec5SDimitry Andric     DFS.PostBlocks.push_back(BB);
2430b57cec5SDimitry Andric     DFS.PostNumbers[BB] = DFS.PostBlocks.size();
2440b57cec5SDimitry Andric   }
2450b57cec5SDimitry Andric };
2460b57cec5SDimitry Andric 
insertEdge(std::optional<BasicBlock * > From,BasicBlock * To)2470b57cec5SDimitry Andric inline bool po_iterator_storage<LoopBlocksTraversal, true>::insertEdge(
248bdd1243dSDimitry Andric     std::optional<BasicBlock *> From, BasicBlock *To) {
2490b57cec5SDimitry Andric   return LBT.visitPreorder(To);
2500b57cec5SDimitry Andric }
2510b57cec5SDimitry Andric 
2520b57cec5SDimitry Andric inline void po_iterator_storage<LoopBlocksTraversal, true>::
finishPostorder(BasicBlock * BB)2530b57cec5SDimitry Andric finishPostorder(BasicBlock *BB) {
2540b57cec5SDimitry Andric   LBT.finishPostorder(BB);
2550b57cec5SDimitry Andric }
2560b57cec5SDimitry Andric 
2570b57cec5SDimitry Andric } // End namespace llvm
2580b57cec5SDimitry Andric 
2590b57cec5SDimitry Andric #endif
260