[IR] Add PHINode::removeIncomingValueIf() (NFC)

[nikic]

Summary

Add an API that allows removing multiple incoming phi values based on a predicate callback, as suggested on D157621.

This makes sure that the removal is linear time rather than quadratic, and avoids subtleties around iterator invalidation.

I have replaced some of the more straightforward users with the new API, though there's a couple more places that should be able to use it.

[efriedma]

LGTM

llvm/lib/IR/Instructions.cpp

150

This isn't the most efficient implementation... instead of using a set, you could iterate over the value operands and the blocks at the same time. But I guess we can improve it later if it matters.

llvm/include/llvm/IR/Instructions.h

public:
   Value *removeIncomingValue(unsigned Idx, bool DeletePHIIfEmpty = true);
 
   Value *removeIncomingValue(const BasicBlock *BB, bool DeletePHIIfEmpty=true) {
     int Idx = getBasicBlockIndex(BB);
     assert(Idx >= 0 && "Invalid basic block argument to remove!");
     return removeIncomingValue(Idx, DeletePHIIfEmpty);
   }
 
+  /// Remove all incoming values for which the predicate returns true.
+  /// The predicate accepts the incoming value index.
+  void removeIncomingValueIf(function_ref<bool(unsigned)> Predicate,
+                             bool DeletePHIIfEmpty = true);
+
   /// Return the first index of the specified basic
   /// block in the value list for this PHI.  Returns -1 if no instance.
   ///
   int getBasicBlockIndex(const BasicBlock *BB) const {
     for (unsigned i = 0, e = getNumOperands(); i != e; ++i)
       if (block_begin()[i] == BB)
         return i;
     return -1;

llvm/lib/IR/Instructions.cpp

Value *PHINode::removeIncomingValue(unsigned Idx, bool DeletePHIIfEmpty) {
   if (getNumOperands() == 0 && DeletePHIIfEmpty) {
     // If anyone is using this PHI, make them use a dummy value instead...
     replaceAllUsesWith(PoisonValue::get(getType()));
     eraseFromParent();
   }
   return Removed;
 }
 
+void PHINode::removeIncomingValueIf(function_ref<bool(unsigned)> Predicate,
+                                    bool DeletePHIIfEmpty) {
+  SmallDenseSet<unsigned> RemoveIndices;
+  for (unsigned Idx = 0; Idx < getNumIncomingValues(); ++Idx)
+    if (Predicate(Idx))
+      RemoveIndices.insert(Idx);
+
+  if (RemoveIndices.empty())
+    return;
+
+  // Remove operands.
+  auto NewOpEnd = remove_if(operands(), [&](Use &U) {
+    return RemoveIndices.contains(U.getOperandNo());
+  });
+  for (Use &U : make_range(NewOpEnd, op_end()))
+    U.set(nullptr);
+
+  // Remove incoming blocks.
+  std::remove_if(const_cast<block_iterator>(block_begin()),
+                 const_cast<block_iterator>(block_end()), [&](BasicBlock *&BB) {
+                   return RemoveIndices.contains(&BB - block_begin());
+                 });
+
+  setNumHungOffUseOperands(getNumOperands() - RemoveIndices.size());
+
+  // If the PHI node is dead, because it has zero entries, nuke it now.
+  if (getNumOperands() == 0 && DeletePHIIfEmpty) {
+    // If anyone is using this PHI, make them use a dummy value instead...
+    replaceAllUsesWith(PoisonValue::get(getType()));
+    eraseFromParent();
+  }
+}
+
 /// growOperands - grow operands - This grows the operand list in response
 /// to a push_back style of operation.  This grows the number of ops by 1.5
 /// times.
 ///
 void PHINode::growOperands() {
   unsigned e = getNumOperands();
   unsigned NumOps = e + e / 2;
   if (NumOps < 2) NumOps = 2;      // 2 op PHI nodes are VERY common.

llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp

for (BasicBlock *BB : RI.R->blocks()) { // This includes the blocks in the
                                             // sub-Scopes.
       assert(BB != PreEntryBlock && "Don't copy the preetntry block");
       BasicBlock *NewBB = CloneBasicBlock(BB, VMap, ".nonchr", &F);
       NewBlocks.push_back(NewBB);
       VMap[BB] = NewBB;
 
       // Unreachable predecessors will not be cloned and will not have an edge
       // to the cloned block. As such, also remove them from any phi nodes.
-      // To avoid iterator invalidation, first collect the dead predecessors
-      // from the first phi node, and then perform the actual removal.
-      if (auto *FirstPN = dyn_cast<PHINode>(NewBB->begin())) {
-        SmallVector<BasicBlock *> DeadPreds;
-        for (BasicBlock *Pred : FirstPN->blocks())
-          if (!DT.isReachableFromEntry(Pred))
-            DeadPreds.push_back(Pred);
-        for (PHINode &PN : make_early_inc_range(NewBB->phis()))
-          for (BasicBlock *Pred : DeadPreds)
-            PN.removeIncomingValue(Pred);
-      }
+        PN.removeIncomingValueIf([&](unsigned Idx) {
+          return !DT.isReachableFromEntry(PN.getIncomingBlock(Idx));
+        });
     }
 
   // Place the cloned blocks right after the original blocks (right before the
   // exit block of.)
   if (ExitBlock)
     F.splice(ExitBlock->getIterator(), &F, NewBlocks[0]->getIterator(),
              F.end());
 

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

if (!HasLoopExit) {
         }
       }
     }
 
     if (InVal) {
       // If all incoming values for the new PHI would be the same, just don't
       // make a new PHI.  Instead, just remove the incoming values from the old
       // PHI.
-
-      // NOTE! This loop walks backwards for a reason! First off, this minimizes
-      // the cost of removal if we end up removing a large number of values, and
-      // second off, this ensures that the indices for the incoming values
-      // aren't invalidated when we remove one.
-      for (int64_t i = PN->getNumIncomingValues() - 1; i >= 0; --i)
-        if (PredSet.count(PN->getIncomingBlock(i)))
-          PN->removeIncomingValue(i, false);
+      PN->removeIncomingValueIf(
+          [&](unsigned Idx) {
+            return PredSet.contains(PN->getIncomingBlock(Idx));
+          },
+          /* DeletePHIIfEmpty */ false);
 
       // Add an incoming value to the PHI node in the loop for the preheader
       // edge.
       PN->addIncoming(InVal, NewBB);
       continue;
     }
 
     // If the values coming into the block are not the same, we need a new

llvm/lib/Transforms/Utils/LoopSimplify.cpp

for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
 
     // Delete all of the incoming values from the old PN except the preheader's
     assert(PreheaderIdx != ~0U && "PHI has no preheader entry??");
     if (PreheaderIdx != 0) {
       PN->setIncomingValue(0, PN->getIncomingValue(PreheaderIdx));
       PN->setIncomingBlock(0, PN->getIncomingBlock(PreheaderIdx));
     }
     // Nuke all entries except the zero'th.
-    for (unsigned i = 0, e = PN->getNumIncomingValues()-1; i != e; ++i)
-      PN->removeIncomingValue(e-i, false);
+    PN->removeIncomingValueIf([](unsigned Idx) { return Idx != 0; },
+                              /* DeletePHIIfEmpty */ false);
 
     // Finally, add the newly constructed PHI node as the entry for the BEBlock.
     PN->addIncoming(NewPN, BEBlock);
 
     // As an optimization, if all incoming values in the new PhiNode (which is a
     // subset of the incoming values of the old PHI node) have the same value,
     // eliminate the PHI Node.
     if (HasUniqueIncomingValue) {

llvm/lib/Transforms/Utils/LoopUtils.cpp

for (PHINode &P : ExitBlock->phis()) {
       // Set the zero'th element of Phi to be from the preheader and remove all
       // other incoming values. Given the loop has dedicated exits, all other
       // incoming values must be from the exiting blocks.
       int PredIndex = 0;
       P.setIncomingBlock(PredIndex, Preheader);
       // Removes all incoming values from all other exiting blocks (including
       // duplicate values from an exiting block).
       // Nuke all entries except the zero'th entry which is the preheader entry.
-      // NOTE! We need to remove Incoming Values in the reverse order as done
-      // below, to keep the indices valid for deletion (removeIncomingValues
-      // updates getNumIncomingValues and shifts all values down into the
-      // operand being deleted).
-      for (unsigned i = 0, e = P.getNumIncomingValues() - 1; i != e; ++i)
-        P.removeIncomingValue(e - i, false);
+      P.removeIncomingValueIf([](unsigned Idx) { return Idx != 0; },
+                              /* DeletePHIIfEmpty */ false);
 
       assert((P.getNumIncomingValues() == 1 &&
               P.getIncomingBlock(PredIndex) == Preheader) &&
              "Should have exactly one value and that's from the preheader!");
     }
 
     if (DT) {
       DTU.applyUpdates({{DominatorTree::Insert, Preheader, ExitBlock}});

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

static void removeSwitchAfterSelectFold(SwitchInst *SI, PHINode *PHI,
   BasicBlock *DestBB = PHI->getParent();
 
   if (DTU && !is_contained(predecessors(DestBB), SelectBB))
     Updates.push_back({DominatorTree::Insert, SelectBB, DestBB});
   Builder.CreateBr(DestBB);
 
   // Remove the switch.
 
-  while (PHI->getBasicBlockIndex(SelectBB) >= 0)
-    PHI->removeIncomingValue(SelectBB);
+  PHI->removeIncomingValueIf(
+      [&](unsigned Idx) { return PHI->getIncomingBlock(Idx) == SelectBB; });
   PHI->addIncoming(SelectValue, SelectBB);
 
   SmallPtrSet<BasicBlock *, 4> RemovedSuccessors;
   for (unsigned i = 0, e = SI->getNumSuccessors(); i < e; ++i) {
     BasicBlock *Succ = SI->getSuccessor(i);
 
     if (Succ == DestBB)
       continue;