[STLExtras] Add distance() for ranges, pred_size(), and succ_size()

This commit adds a wrapper for std::distance() which works with ranges.
As it would be a common case to write `distance(predecessors(BB))`, this
also introduces `pred_size()` and `succ_size()` helpers to make that
easier to write.

Differential Revision: https://reviews.llvm.org/D46668

llvm-svn: 332057

llvm/include/llvm/ADT/STLExtras.h

@@ -1026,6 +1026,13 @@ void erase_if(Container &C, UnaryPredicate P) {
   C.erase(remove_if(C, P), C.end());
 }
 
+/// Wrapper function around std::distance which works with ranges.
+template <typename R>
+auto distance(R &&Range)
+    -> decltype(std::distance(Range.begin(), Range.end())) {
+  return std::distance(Range.begin(), Range.end());
+}
+
 //===----------------------------------------------------------------------===//
 //     Extra additions to <memory>
 //===----------------------------------------------------------------------===//

llvm/include/llvm/IR/CFG.h

@@ -107,6 +107,9 @@ inline const_pred_iterator pred_end(const BasicBlock *BB) {
 inline bool pred_empty(const BasicBlock *BB) {
   return pred_begin(BB) == pred_end(BB);
 }
+inline unsigned pred_size(const BasicBlock *BB) {
+  return std::distance(pred_begin(BB), pred_end(BB));
+}
 inline pred_range predecessors(BasicBlock *BB) {
   return pred_range(pred_begin(BB), pred_end(BB));
 }
@@ -140,6 +143,9 @@ inline succ_const_iterator succ_end(const BasicBlock *BB) {
 inline bool succ_empty(const BasicBlock *BB) {
   return succ_begin(BB) == succ_end(BB);
 }
+inline unsigned succ_size(const BasicBlock *BB) {
+  return std::distance(succ_begin(BB), succ_end(BB));
+}
 inline succ_range successors(BasicBlock *BB) {
   return succ_range(succ_begin(BB), succ_end(BB));
 }

llvm/lib/Analysis/BranchProbabilityInfo.cpp

@@ -873,8 +873,7 @@ BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
   if (I != Probs.end())
     return I->second;
 
-  return {1,
-          static_cast<uint32_t>(std::distance(succ_begin(Src), succ_end(Src)))};
+  return {1, static_cast<uint32_t>(succ_size(Src))};
 }
 
 BranchProbability

llvm/lib/Analysis/LazyCallGraph.cpp

@@ -1272,8 +1272,7 @@ LazyCallGraph::RefSCC::removeInternalRefEdge(Node &SourceN,
       // the removal hasn't changed the structure at all. This is an important
       // special case and we can directly exit the entire routine more
       // efficiently as soon as we discover it.
-      if (std::distance(RefSCCNodes.begin(), RefSCCNodes.end()) ==
-          NumRefSCCNodes) {
+      if (distance(RefSCCNodes) == NumRefSCCNodes) {
         // Clear out the low link field as we won't need it.
         for (Node *N : RefSCCNodes)
           N->LowLink = -1;
@@ -1739,7 +1738,7 @@ static void printNode(raw_ostream &OS, LazyCallGraph::Node &N) {
 }
 
 static void printSCC(raw_ostream &OS, LazyCallGraph::SCC &C) {
-  ptrdiff_t Size = std::distance(C.begin(), C.end());
+  ptrdiff_t Size = distance(C);
   OS << "    SCC with " << Size << " functions:\n";
 
   for (LazyCallGraph::Node &N : C)
@@ -1747,7 +1746,7 @@ static void printSCC(raw_ostream &OS, LazyCallGraph::SCC &C) {
 }
 
 static void printRefSCC(raw_ostream &OS, LazyCallGraph::RefSCC &C) {
-  ptrdiff_t Size = std::distance(C.begin(), C.end());
+  ptrdiff_t Size = distance(C);
   OS << "  RefSCC with " << Size << " call SCCs:\n";
 
   for (LazyCallGraph::SCC &InnerC : C)

llvm/lib/Analysis/LoopInfo.cpp

@@ -417,7 +417,7 @@ void Loop::getUniqueExitBlocks(
       // If a terminator has more then two successors, for example SwitchInst,
       // then it is possible that there are multiple edges from current block
       // to one exit block.
-      if (std::distance(succ_begin(BB), succ_end(BB)) <= 2) {
+      if (succ_size(BB) <= 2) {
         ExitBlocks.push_back(Successor);
         continue;
       }

llvm/lib/AsmParser/LLParser.cpp

@@ -6756,8 +6756,8 @@ bool LLParser::sortUseListOrder(Value *V, ArrayRef<unsigned> Indexes,
   if (NumUses < 2)
     return Error(Loc, "value only has one use");
   if (Order.size() != Indexes.size() || NumUses > Indexes.size())
-    return Error(Loc, "wrong number of indexes, expected " +
-                          Twine(std::distance(V->use_begin(), V->use_end())));
+    return Error(Loc,
+                 "wrong number of indexes, expected " + Twine(V->getNumUses()));
 
   V->sortUseList([&](const Use &L, const Use &R) {
     return Order.lookup(&L) < Order.lookup(&R);

llvm/lib/Bitcode/Writer/ValueEnumerator.cpp

@@ -489,7 +489,7 @@ void ValueEnumerator::print(raw_ostream &OS, const ValueMapType &Map,
     V->print(errs());
     errs() << '\n';
 
-    OS << " Uses(" << std::distance(V->use_begin(),V->use_end()) << "):";
+    OS << " Uses(" << V->getNumUses() << "):";
     for (const Use &U : V->uses()) {
       if (&U != &*V->use_begin())
         OS << ",";

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -3094,8 +3094,7 @@ private:
       Instruction *CurrentI = cast<Instruction>(CurrentValue);
       bool IsDefinedInThisBB = CurrentI->getParent() == CurrentBlock;
 
-      unsigned PredCount =
-          std::distance(pred_begin(CurrentBlock), pred_end(CurrentBlock));
+      unsigned PredCount = pred_size(CurrentBlock);
       // if Current Value is not defined in this basic block we are interested
       // in values in predecessors.
       if (!IsDefinedInThisBB) {

llvm/lib/CodeGen/ImplicitNullChecks.cpp

@@ -597,8 +597,7 @@ MachineInstr *ImplicitNullChecks::insertFaultingInstr(
   unsigned DefReg = NoRegister;
   if (NumDefs != 0) {
     DefReg = MI->defs().begin()->getReg();
-    assert(std::distance(MI->defs().begin(), MI->defs().end()) == 1 &&
-           "expected exactly one def!");
+    assert(distance(MI->defs()) == 1 && "expected exactly one def!");
   }
 
   FaultMaps::FaultKind FK;

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -1700,8 +1700,7 @@ SelectionDAGBuilder::getEdgeProbability(const MachineBasicBlock *Src,
   if (!BPI) {
     // If BPI is not available, set the default probability as 1 / N, where N is
     // the number of successors.
-    auto SuccSize = std::max<uint32_t>(
-        std::distance(succ_begin(SrcBB), succ_end(SrcBB)), 1);
+    auto SuccSize = std::max<uint32_t>(succ_size(SrcBB), 1);
     return BranchProbability(1, SuccSize);
   }
   return BPI->getEdgeProbability(SrcBB, DstBB);

llvm/lib/IR/Value.cpp

@@ -167,9 +167,7 @@ bool Value::isUsedInBasicBlock(const BasicBlock *BB) const {
   return false;
 }
 
-unsigned Value::getNumUses() const {
-  return (unsigned)std::distance(use_begin(), use_end());
-}
+unsigned Value::getNumUses() const { return (unsigned)distance(uses()); }
 
 static bool getSymTab(Value *V, ValueSymbolTable *&ST) {
   ST = nullptr;

llvm/lib/Target/PowerPC/PPCLoopPreIncPrep.cpp

@@ -254,8 +254,7 @@ bool PPCLoopPreIncPrep::runOnLoop(Loop *L) {
   const PPCSubtarget *ST =
     TM ? TM->getSubtargetImpl(*Header->getParent()) : nullptr;
 
-  unsigned HeaderLoopPredCount =
-    std::distance(pred_begin(Header), pred_end(Header));
+  unsigned HeaderLoopPredCount = pred_size(Header);
 
   // Collect buckets of comparable addresses used by loads and stores.
   SmallVector<Bucket, 16> Buckets;

llvm/lib/Transforms/IPO/PartialInlining.cpp

@@ -403,8 +403,7 @@ PartialInlinerImpl::computeOutliningColdRegionsInfo(Function *F,
 
   auto IsSingleEntry = [](SmallVectorImpl<BasicBlock *> &BlockList) {
     BasicBlock *Dom = BlockList.front();
-    return BlockList.size() > 1 &&
-           std::distance(pred_begin(Dom), pred_end(Dom)) == 1;
+    return BlockList.size() > 1 && pred_size(Dom) == 1;
   };
 
   auto IsSingleExit =
@@ -556,10 +555,6 @@ PartialInlinerImpl::computeOutliningInfo(Function *F) {
     return is_contained(successors(BB), Succ);
   };
 
-  auto SuccSize = [](BasicBlock *BB) {
-    return std::distance(succ_begin(BB), succ_end(BB));
-  };
-
   auto IsReturnBlock = [](BasicBlock *BB) {
     TerminatorInst *TI = BB->getTerminator();
     return isa<ReturnInst>(TI);
@@ -596,7 +591,7 @@ PartialInlinerImpl::computeOutliningInfo(Function *F) {
     if (OutliningInfo->GetNumInlinedBlocks() >= MaxNumInlineBlocks)
       break;
 
-    if (SuccSize(CurrEntry) != 2)
+    if (succ_size(CurrEntry) != 2)
       break;
 
     BasicBlock *Succ1 = *succ_begin(CurrEntry);
@@ -670,7 +665,7 @@ PartialInlinerImpl::computeOutliningInfo(Function *F) {
   // peeling off dominating blocks from the outlining region:
   while (OutliningInfo->GetNumInlinedBlocks() < MaxNumInlineBlocks) {
     BasicBlock *Cand = OutliningInfo->NonReturnBlock;
-    if (SuccSize(Cand) != 2)
+    if (succ_size(Cand) != 2)
       break;
 
     if (HasNonEntryPred(Cand))

llvm/lib/Transforms/Scalar/GVNHoist.cpp

@@ -578,7 +578,7 @@ private:
 
   // Returns true when the values are flowing out to each edge.
   bool valueAnticipable(CHIArgs C, TerminatorInst *TI) const {
-    if (TI->getNumSuccessors() > (unsigned)std::distance(C.begin(), C.end()))
+    if (TI->getNumSuccessors() > (unsigned)distance(C))
       return false; // Not enough args in this CHI.
 
     for (auto CHI : C) {

llvm/lib/Transforms/Scalar/JumpThreading.cpp

@@ -945,10 +945,10 @@ static unsigned GetBestDestForJumpOnUndef(BasicBlock *BB) {
   unsigned MinSucc = 0;
   BasicBlock *TestBB = BBTerm->getSuccessor(MinSucc);
   // Compute the successor with the minimum number of predecessors.
-  unsigned MinNumPreds = std::distance(pred_begin(TestBB), pred_end(TestBB));
+  unsigned MinNumPreds = pred_size(TestBB);
   for (unsigned i = 1, e = BBTerm->getNumSuccessors(); i != e; ++i) {
     TestBB = BBTerm->getSuccessor(i);
-    unsigned NumPreds = std::distance(pred_begin(TestBB), pred_end(TestBB));
+    unsigned NumPreds = pred_size(TestBB);
     if (NumPreds < MinNumPreds) {
       MinSucc = i;
       MinNumPreds = NumPreds;
@@ -1648,8 +1648,7 @@ bool JumpThreadingPass::ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
   // not thread. By doing so, we do not need to duplicate the current block and
   // also miss potential opportunities in case we dont/cant duplicate.
   if (OnlyDest && OnlyDest != MultipleDestSentinel) {
-    if (PredWithKnownDest ==
-        (size_t)std::distance(pred_begin(BB), pred_end(BB))) {
+    if (PredWithKnownDest == (size_t)pred_size(BB)) {
       bool SeenFirstBranchToOnlyDest = false;
       std::vector <DominatorTree::UpdateType> Updates;
       Updates.reserve(BB->getTerminator()->getNumSuccessors() - 1);

llvm/lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

@@ -285,8 +285,7 @@ bool MergedLoadStoreMotion::mergeStores(BasicBlock *T) {
     return false; // No. More than 2 predecessors.
 
   // #Instructions in Succ1 for Compile Time Control
-  int Size1 = std::distance(Pred1->instructionsWithoutDebug().begin(),
-                            Pred1->instructionsWithoutDebug().end());
+  int Size1 = distance(Pred1->instructionsWithoutDebug());
   int NStores = 0;
 
   for (BasicBlock::reverse_iterator RBI = Pred0->rbegin(), RBE = Pred0->rend();

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

@@ -966,7 +966,7 @@ static Value *findBasePointer(Value *I, DefiningValueMapTy &Cache) {
     auto MakeBaseInstPlaceholder = [](Instruction *I) -> Instruction* {
       if (isa<PHINode>(I)) {
         BasicBlock *BB = I->getParent();
-        int NumPreds = std::distance(pred_begin(BB), pred_end(BB));
+        int NumPreds = pred_size(BB);
         assert(NumPreds > 0 && "how did we reach here");
         std::string Name = suffixed_name_or(I, ".base", "base_phi");
         return PHINode::Create(I->getType(), NumPreds, Name, I);
@@ -1811,7 +1811,7 @@ static void relocationViaAlloca(
 
     SmallVector<Instruction *, 20> Uses;
     // PERF: trade a linear scan for repeated reallocation
-    Uses.reserve(std::distance(Def->user_begin(), Def->user_end()));
+    Uses.reserve(Def->getNumUses());
     for (User *U : Def->users()) {
       if (!isa<ConstantExpr>(U)) {
         // If the def has a ConstantExpr use, then the def is either a

llvm/lib/Transforms/Utils/CloneFunction.cpp

@@ -538,7 +538,7 @@ void llvm::CloneAndPruneIntoFromInst(Function *NewFunc, const Function *OldFunc,
     // phi nodes will have invalid entries.  Update the PHI nodes in this
     // case.
     PHINode *PN = cast<PHINode>(NewBB->begin());
-    NumPreds = std::distance(pred_begin(NewBB), pred_end(NewBB));
+    NumPreds = pred_size(NewBB);
     if (NumPreds != PN->getNumIncomingValues()) {
       assert(NumPreds < PN->getNumIncomingValues());
       // Count how many times each predecessor comes to this block.

llvm/lib/Transforms/Utils/Local.cpp

@@ -669,8 +669,7 @@ void llvm::MergeBasicBlockIntoOnlyPred(BasicBlock *DestBB, DominatorTree *DT,
   // dominator edges will be redirected to DestBB.
   std::vector <DominatorTree::UpdateType> Updates;
   if (DDT && !ReplaceEntryBB) {
-    Updates.reserve(1 +
-                    (2 * std::distance(pred_begin(PredBB), pred_end(PredBB))));
+    Updates.reserve(1 + (2 * pred_size(PredBB)));
     Updates.push_back({DominatorTree::Delete, PredBB, DestBB});
     for (auto I = pred_begin(PredBB), E = pred_end(PredBB); I != E; ++I) {
       Updates.push_back({DominatorTree::Delete, *I, PredBB});
@@ -975,7 +974,7 @@ bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB,
 
   std::vector<DominatorTree::UpdateType> Updates;
   if (DDT) {
-    Updates.reserve(1 + (2 * std::distance(pred_begin(BB), pred_end(BB))));
+    Updates.reserve(1 + (2 * pred_size(BB)));
     Updates.push_back({DominatorTree::Delete, BB, Succ});
     // All predecessors of BB will be moved to Succ.
     for (auto I = pred_begin(BB), E = pred_end(BB); I != E; ++I) {

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

@@ -295,7 +295,7 @@ private:
   unsigned getNumPreds(const BasicBlock *BB) {
     unsigned &NP = BBNumPreds[BB];
     if (NP == 0)
-      NP = std::distance(pred_begin(BB), pred_end(BB)) + 1;
+      NP = pred_size(BB) + 1;
     return NP - 1;
   }
 

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -688,9 +688,7 @@ Value *SimplifyCFGOpt::isValueEqualityComparison(TerminatorInst *TI) {
   if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
     // Do not permit merging of large switch instructions into their
     // predecessors unless there is only one predecessor.
-    if (SI->getNumSuccessors() * std::distance(pred_begin(SI->getParent()),
-                                               pred_end(SI->getParent())) <=
-        128)
+    if (SI->getNumSuccessors() * pred_size(SI->getParent()) <= 128)
       CV = SI->getCondition();
   } else if (BranchInst *BI = dyn_cast<BranchInst>(TI))
     if (BI->isConditional() && BI->getCondition()->hasOneUse())
@@ -2871,7 +2869,7 @@ static Value *ensureValueAvailableInSuccessor(Value *V, BasicBlock *BB,
       if (!AlternativeV)
         break;
 
-      assert(std::distance(pred_begin(Succ), pred_end(Succ)) == 2);
+      assert(pred_size(Succ) == 2);
       auto PredI = pred_begin(Succ);
       BasicBlock *OtherPredBB = *PredI == BB ? *++PredI : *PredI;
       if (PHI->getIncomingValueForBlock(OtherPredBB) == AlternativeV)
@@ -5752,7 +5750,7 @@ bool SimplifyCFGOpt::SimplifyUncondBranch(BranchInst *BI,
   // backedge, so we can eliminate BB.
   bool NeedCanonicalLoop =
       Options.NeedCanonicalLoop &&
-      (LoopHeaders && std::distance(pred_begin(BB), pred_end(BB)) > 1 &&
+      (LoopHeaders && pred_size(BB) > 1 &&
        (LoopHeaders->count(BB) || LoopHeaders->count(Succ)));
   BasicBlock::iterator I = BB->getFirstNonPHIOrDbg()->getIterator();
   if (I->isTerminator() && BB != &BB->getParent()->getEntryBlock() &&

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -356,7 +356,7 @@ void VPlan::updateDominatorTree(DominatorTree *DT, BasicBlock *LoopPreHeaderBB,
            "One successor of a basic block does not lead to the other.");
     assert(InterimSucc->getSinglePredecessor() &&
            "Interim successor has more than one predecessor.");
-    assert(std::distance(pred_begin(PostDomSucc), pred_end(PostDomSucc)) == 2 &&
+    assert(pred_size(PostDomSucc) == 2 &&
            "PostDom successor has more than two predecessors.");
     DT->addNewBlock(InterimSucc, BB);
     DT->addNewBlock(PostDomSucc, BB);

llvm/unittests/ADT/IteratorTest.cpp

@@ -365,4 +365,12 @@ TEST(ZipIteratorTest, Reverse) {
   EXPECT_TRUE(all_of(ascending, [](unsigned n) { return (n & 0x01) == 0; }));
 }
 
+TEST(RangeTest, Distance) {
+  std::vector<int> v1;
+  std::vector<int> v2{1, 2, 3};
+
+  EXPECT_EQ(std::distance(v1.begin(), v1.end()), distance(v1));
+  EXPECT_EQ(std::distance(v2.begin(), v2.end()), distance(v2));
+}
+
 } // anonymous namespace

llvm/unittests/IR/BasicBlockTest.cpp

@@ -73,9 +73,9 @@ TEST(BasicBlockTest, PhiRange) {
   auto isPhi = [](Instruction &I) { return isa<PHINode>(&I); };
   auto Phis = make_filter_range(*BB, isPhi);
   auto ReversedPhis = reverse(make_filter_range(*BB, isPhi));
-  EXPECT_EQ(std::distance(Phis.begin(), Phis.end()), 3);
+  EXPECT_EQ(distance(Phis), 3);
   EXPECT_EQ(&*Phis.begin(), P1);
-  EXPECT_EQ(std::distance(ReversedPhis.begin(), ReversedPhis.end()), 3);
+  EXPECT_EQ(distance(ReversedPhis), 3);
   EXPECT_EQ(&*ReversedPhis.begin(), P3);
 
   // And iterate a const range.
@@ -87,8 +87,7 @@ TEST(BasicBlockTest, PhiRange) {
 }
 
 #define CHECK_ITERATORS(Range1, Range2)                                        \
-  EXPECT_EQ(std::distance(Range1.begin(), Range1.end()),                       \
-            std::distance(Range2.begin(), Range2.end()));                      \
+  EXPECT_EQ(distance(Range1), distance(Range2));                               \
   for (auto Pair : zip(Range1, Range2))                                        \
     EXPECT_EQ(&std::get<0>(Pair), std::get<1>(Pair));