[UnJ] Use SmallPtrSets for block collections NFC

[dmgreen]

Summary

We no longer care about the order of blocks in these collections,
so can change to SmallPtrSets, making contains checks quicker.

[SjoerdMeijer]

For readability, you could consider creating a type alias for "SmallPtrSet<BasicBlock *, 4>", but that's just a nit and this looks like some straightforward refactoring to me.

[dmgreen]

Thanks

lib/Transforms/Utils/LoopUnrollAndJam.cpp

 #include "llvm/Transforms/Utils/UnrollLoop.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-unroll-and-jam"
 
 STATISTIC(NumUnrolledAndJammed, "Number of loops unroll and jammed");
 STATISTIC(NumCompletelyUnrolledAndJammed, "Number of loops unroll and jammed");
 
-static bool containsBB(std::vector<BasicBlock *> &V, BasicBlock *BB) {
-  return std::find(V.begin(), V.end(), BB) != V.end();
-}
-
 // Partition blocks in an outer/inner loop pair into blocks before and after
 // the loop
-static bool partitionOuterLoopBlocks(Loop *L, Loop *SubLoop,
-                                     std::vector<BasicBlock *> &ForeBlocks,
-                                     std::vector<BasicBlock *> &SubLoopBlocks,
-                                     std::vector<BasicBlock *> &AftBlocks,
-                                     DominatorTree *DT) {
+static bool partitionOuterLoopBlocks(
+    Loop *L, Loop *SubLoop, SmallPtrSet<BasicBlock *, 4> &ForeBlocks,
+    SmallPtrSet<BasicBlock *, 4> &SubLoopBlocks,
+    SmallPtrSet<BasicBlock *, 4> &AftBlocks, DominatorTree *DT) {
   BasicBlock *SubLoopLatch = SubLoop->getLoopLatch();
-  SubLoopBlocks = SubLoop->getBlocks();
+  SubLoopBlocks.insert(SubLoop->block_begin(), SubLoop->block_end());
 
   for (BasicBlock *BB : L->blocks()) {
     if (!SubLoop->contains(BB)) {
       if (DT->dominates(SubLoopLatch, BB))
-        AftBlocks.push_back(BB);
+        AftBlocks.insert(BB);
       else
-        ForeBlocks.push_back(BB);
+        ForeBlocks.insert(BB);
     }
   }
 
   // Check that all blocks in ForeBlocks together dominate the subloop
   // TODO: This might ideally be done better with a dominator/postdominators.
   BasicBlock *SubLoopPreHeader = SubLoop->getLoopPreheader();
   for (BasicBlock *BB : ForeBlocks) {
     if (BB == SubLoopPreHeader)
       continue;
     TerminatorInst *TI = BB->getTerminator();
     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
-      if (!containsBB(ForeBlocks, TI->getSuccessor(i)))
+      if (!ForeBlocks.count(TI->getSuccessor(i)))
         return false;
   }
 
   return true;
 }
 
 // Move the phi operands of Header from Latch out of AftBlocks to InsertLoc.
 static void
 moveHeaderPhiOperandsToForeBlocks(BasicBlock *Header, BasicBlock *Latch,
                                   Instruction *InsertLoc,
-                                  std::vector<BasicBlock *> &AftBlocks) {
+                                  SmallPtrSetImpl<BasicBlock *> &AftBlocks) {
   // We need to ensure we move the instructions in the correct order,
   // starting with the earliest required instruction and moving forward.
   std::vector<Instruction *> Worklist;
   std::vector<Instruction *> Visited;
   for (auto &Phi : Header->phis()) {
     Value *V = Phi.getIncomingValueForBlock(Latch);
     if (Instruction *I = dyn_cast<Instruction>(V))
       Worklist.push_back(I);
   }
 
   while (!Worklist.empty()) {
     Instruction *I = Worklist.back();
     Worklist.pop_back();
-    if (!containsBB(AftBlocks, I->getParent()))
+    if (!AftBlocks.count(I->getParent()))
       continue;
 
     Visited.push_back(I);
     for (auto &U : I->operands())
       if (Instruction *II = dyn_cast<Instruction>(U))
         Worklist.push_back(II);
   }
 
llvm::UnrollAndJamLoop(Loop *L, unsigned Count, unsigned TripCount,
   assert(BI && !BI->isUnconditional());
   bool ContinueOnTrue = L->contains(BI->getSuccessor(0));
   BasicBlock *LoopExit = BI->getSuccessor(ContinueOnTrue);
   bool SubLoopContinueOnTrue = SubLoop->contains(
       SubLoop->getLoopLatch()->getTerminator()->getSuccessor(0));
 
   // Partition blocks in an outer/inner loop pair into blocks before and after
   // the loop
-  std::vector<BasicBlock *> SubLoopBlocks;
-  std::vector<BasicBlock *> ForeBlocks;
-  std::vector<BasicBlock *> AftBlocks;
+  SmallPtrSet<BasicBlock *, 4> SubLoopBlocks;
+  SmallPtrSet<BasicBlock *, 4> ForeBlocks;
+  SmallPtrSet<BasicBlock *, 4> AftBlocks;
   partitionOuterLoopBlocks(L, SubLoop, ForeBlocks, SubLoopBlocks, AftBlocks,
                            DT);
 
   // We keep track of the entering/first and exiting/last block of each of
   // Fore/SubLoop/Aft in each iteration. This helps make the stapling up of
   // blocks easier.
   std::vector<BasicBlock *> ForeBlocksFirst;
   std::vector<BasicBlock *> ForeBlocksLast;
for (unsigned It = 1; It != Count; ++It) {
     // Maps Blocks[It] -> Blocks[It-1]
     DenseMap<Value *, Value *> PrevItValueMap;
 
     for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
       ValueToValueMapTy VMap;
       BasicBlock *New = CloneBasicBlock(*BB, VMap, "." + Twine(It));
       Header->getParent()->getBasicBlockList().push_back(New);
 
-      if (containsBB(ForeBlocks, *BB)) {
+      if (ForeBlocks.count(*BB)) {
         L->addBasicBlockToLoop(New, *LI);
 
         if (*BB == ForeBlocksFirst[0])
           ForeBlocksFirst.push_back(New);
         if (*BB == ForeBlocksLast[0])
           ForeBlocksLast.push_back(New);
-      } else if (containsBB(SubLoopBlocks, *BB)) {
+      } else if (SubLoopBlocks.count(*BB)) {
         SubLoop->addBasicBlockToLoop(New, *LI);
 
         if (*BB == SubLoopBlocksFirst[0])
           SubLoopBlocksFirst.push_back(New);
         if (*BB == SubLoopBlocksLast[0])
           SubLoopBlocksLast.push_back(New);
-      } else if (containsBB(AftBlocks, *BB)) {
+      } else if (AftBlocks.count(*BB)) {
         L->addBasicBlockToLoop(New, *LI);
 
         if (*BB == AftBlocksFirst[0])
           AftBlocksFirst.push_back(New);
         if (*BB == AftBlocksLast[0])
           AftBlocksLast.push_back(New);
       } else {
         llvm_unreachable("BB being cloned should be in Fore/Sub/Aft");
#endif
   // Update LoopInfo if the loop is completely removed.
   if (CompletelyUnroll)
     LI->erase(L);
 
   return CompletelyUnroll ? LoopUnrollResult::FullyUnrolled
                           : LoopUnrollResult::PartiallyUnrolled;
 }
 
-static bool getLoadsAndStores(std::vector<BasicBlock *> &Blocks,
+static bool getLoadsAndStores(SmallPtrSetImpl<BasicBlock *> &Blocks,
                               SmallVector<Value *, 4> &MemInstr) {
   // Scan the BBs and collect legal loads and stores.
   // Returns false if non-simple loads/stores are found.
   for (BasicBlock *BB : Blocks) {
     for (Instruction &I : *BB) {
       if (auto *Ld = dyn_cast<LoadInst>(&I)) {
         if (!Ld->isSimple())
           return false;
for (Value *J : Later) {
             return false;
         }
       }
     }
   }
   return true;
 }
 
-static bool checkDependencies(Loop *L, std::vector<BasicBlock *> &ForeBlocks,
-                              std::vector<BasicBlock *> &SubLoopBlocks,
-                              std::vector<BasicBlock *> &AftBlocks,
+static bool checkDependencies(Loop *L,
+                              SmallPtrSetImpl<BasicBlock *> &ForeBlocks,
+                              SmallPtrSetImpl<BasicBlock *> &SubLoopBlocks,
+                              SmallPtrSetImpl<BasicBlock *> &AftBlocks,
                               DependenceInfo &DI) {
   // Get all loads/store pairs for each blocks
   SmallVector<Value *, 4> ForeMemInstr;
   SmallVector<Value *, 4> SubLoopMemInstr;
   SmallVector<Value *, 4> AftMemInstr;
   if (!getLoadsAndStores(ForeBlocks, ForeMemInstr) ||
       !getLoadsAndStores(SubLoopBlocks, SubLoopMemInstr) ||
       !getLoadsAndStores(AftBlocks, AftMemInstr))
if (Latch != Exit)
     return false;
   if (SubLoopLatch != SubLoopExit)
     return false;
 
   if (Header->hasAddressTaken() || SubLoopHeader->hasAddressTaken())
     return false;
 
   // Split blocks into Fore/SubLoop/Aft based on dominators
-  std::vector<BasicBlock *> SubLoopBlocks;
-  std::vector<BasicBlock *> ForeBlocks;
-  std::vector<BasicBlock *> AftBlocks;
+  SmallPtrSet<BasicBlock *, 4> SubLoopBlocks;
+  SmallPtrSet<BasicBlock *, 4> ForeBlocks;
+  SmallPtrSet<BasicBlock *, 4> AftBlocks;
   if (!partitionOuterLoopBlocks(L, SubLoop, ForeBlocks, SubLoopBlocks,
                                 AftBlocks, &DT))
     return false;
 
   // Aft blocks may need to move instructions to fore blocks, which becomes more
   // difficult if there are multiple (potentially conditionally executed)
   // blocks. For now we just exclude loops with multiple aft blocks.
   if (AftBlocks.size() != 1)
if (Instruction *I = dyn_cast<Instruction>(V))
       Worklist.push_back(I);
   }
   while (!Worklist.empty()) {
     Instruction *I = Worklist.back();
     Worklist.pop_back();
     if (Visited.insert(I).second) {
       if (SubLoop->contains(I->getParent()))
         return false;
-      if (containsBB(AftBlocks, I->getParent())) {
+      if (AftBlocks.count(I->getParent())) {
         // If we hit a phi node in afts we know we are done (probably LCSSA)
         if (isa<PHINode>(I))
           return false;
         if (I->mayHaveSideEffects() || I->mayReadOrWriteMemory())
           return false;
         for (auto &U : I->operands())
           if (Instruction *II = dyn_cast<Instruction>(U))
             Worklist.push_back(II);