[IR] Value: add replaceUsesWithIf() utility

[lebedev.ri]

Summary

While there is always a Value::replaceAllUsesWith(),
sometimes the replacement needs to be conditional.

I have only cleaned a few cases where replaceUsesWithIf()
could be used, to both add test coverage,
and show that it is actually useful.

[lebedev.ri]

Test plan: $ ninja check-llvm

[asbirlea]

llvm/lib/Analysis/MemorySSAUpdater.cpp

280

Could you move this comment inside the lambda?

[lebedev.ri]

Move comment in MemorySSAUpdater.cpp

[jdoerfert]

for (; UI != E;) aka. while (UI != E)

LGTM.

llvm/include/llvm/IR/Value.h

 //===- llvm/Value.h - Definition of the Value class -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the Value class.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_IR_VALUE_H
 #define LLVM_IR_VALUE_H
 
 #include "llvm-c/Types.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/Use.h"
 #include "llvm/Support/CBindingWrapping.h"
 #include "llvm/Support/Casting.h"
 #include <cassert>
 #include <iterator>
 #include <memory>
 
public:
   void replaceAllUsesWith(Value *V);
 
   /// Change non-metadata uses of this to point to a new Value.
   ///
   /// Go through the uses list for this definition and make each use point to
   /// "V" instead of "this". This function skips metadata entries in the list.
   void replaceNonMetadataUsesWith(Value *V);
 
+  /// Go through the uses list for this definition and make each use point
+  /// to "V" if the callback ShouldReplace returns true for the given Use.
+  /// Unlike replaceAllUsesWith() this function does not support basic block
+  /// values or constant users.
+  void replaceUsesWithIf(Value *New,
+                         llvm::function_ref<bool(Use &U)> ShouldReplace) {
+    assert(New && "Value::replaceUsesWithIf(<null>) is invalid!");
+    assert(New->getType() == getType() &&
+           "replaceUses of value with new value of different type!");
+
+    use_iterator UI = use_begin(), E = use_end();
+    for (; UI != E;) {
+      Use &U = *UI;
+      ++UI;
+      if (!ShouldReplace(U))
+        continue;
+      U.set(New);
+    }
+  }
+
   /// replaceUsesOutsideBlock - Go through the uses list for this definition and
   /// make each use point to "V" instead of "this" when the use is outside the
   /// block. 'This's use list is expected to have at least one element.
-  /// Unlike replaceAllUsesWith this function does not support basic block
+  /// Unlike replaceAllUsesWith() this function does not support basic block
   /// values or constant users.
   void replaceUsesOutsideBlock(Value *V, BasicBlock *BB);
 
   //----------------------------------------------------------------------
   // Methods for handling the chain of uses of this Value.
   //
   // Materializing a function can introduce new uses, so these methods come in
   // two variants:

llvm/lib/Analysis/MemorySSAUpdater.cpp

void MemorySSAUpdater::insertDef(MemoryDef *MD, bool RenameUses) {
   MemoryAccess *DefBefore = getPreviousDef(MD);
   bool DefBeforeSameBlock = DefBefore->getBlock() == MD->getBlock();
 
   // There is a def before us, which means we can replace any store/phi uses
   // of that thing with us, since we are in the way of whatever was there
   // before.
   // We now define that def's memorydefs and memoryphis
   if (DefBeforeSameBlock) {
-    for (auto UI = DefBefore->use_begin(), UE = DefBefore->use_end();
-         UI != UE;) {
-      Use &U = *UI++;
+    DefBefore->replaceUsesWithIf(MD, [MD](Use &U) {
       // Leave the MemoryUses alone.
       // Also make sure we skip ourselves to avoid self references.
-      if (isa<MemoryUse>(U.getUser()) || U.getUser() == MD)
-        continue;
+      User *Usr = U.getUser();
+      return !isa<MemoryUse>(Usr) && Usr != MD;
       // Defs are automatically unoptimized when the user is set to MD below,
       // because the isOptimized() call will fail to find the same ID.
-      U.set(MD);
-    }
+    });
   }
 
   // and that def is now our defining access.
   MD->setDefiningAccess(DefBefore);
 
   // Remember the index where we may insert new phis below.
   unsigned NewPhiIndex = InsertedPHIs.size();
 
   SmallVector<WeakVH, 8> FixupList(InsertedPHIs.begin(), InsertedPHIs.end());
   if (!DefBeforeSameBlock) {
     // If there was a local def before us, we must have the same effect it

llvm/lib/IR/Value.cpp

 void Value::replaceUsesOutsideBlock(Value *New, BasicBlock *BB) {
   assert(New && "Value::replaceUsesOutsideBlock(<null>, BB) is invalid!");
   assert(!contains(New, this) &&
          "this->replaceUsesOutsideBlock(expr(this), BB) is NOT valid!");
   assert(New->getType() == getType() &&
          "replaceUses of value with new value of different type!");
   assert(BB && "Basic block that may contain a use of 'New' must be defined\n");
 
-  use_iterator UI = use_begin(), E = use_end();
-  for (; UI != E;) {
-    Use &U = *UI;
-    ++UI;
-    auto *Usr = dyn_cast<Instruction>(U.getUser());
-    if (Usr && Usr->getParent() == BB)
-      continue;
-    U.set(New);
-  }
+  replaceUsesWithIf(New, [BB](Use &U) {
+    auto *I = dyn_cast<Instruction>(U.getUser());
+    // Don't replace if it's an instruction in the BB basic block.
+    return !I || I->getParent() != BB;
+  });
 }
 
 namespace {
 // Various metrics for how much to strip off of pointers.
 enum PointerStripKind {
   PSK_ZeroIndices,
   PSK_ZeroIndicesAndAliases,
   PSK_ZeroIndicesAndAliasesSameRepresentation,

llvm/lib/Target/WebAssembly/WebAssemblyOptimizeReturned.cpp

void OptimizeReturned::visitCallSite(CallSite CS) {
   for (unsigned I = 0, E = CS.getNumArgOperands(); I < E; ++I)
     if (CS.paramHasAttr(I, Attribute::Returned)) {
       Instruction *Inst = CS.getInstruction();
       Value *Arg = CS.getArgOperand(I);
       // Ignore constants, globals, undef, etc.
       if (isa<Constant>(Arg))
         continue;
       // Like replaceDominatedUsesWith but using Instruction/Use dominance.
-      for (auto UI = Arg->use_begin(), UE = Arg->use_end(); UI != UE;) {
-        Use &U = *UI++;
-        if (DT->dominates(Inst, U))
-          U.set(Inst);
-      }
+      Arg->replaceUsesWithIf(Inst,
+                             [&](Use &U) { return DT->dominates(Inst, U); });
     }
 }
 
 bool OptimizeReturned::runOnFunction(Function &F) {
   LLVM_DEBUG(dbgs() << "********** Optimize returned Attributes **********\n"
                        "********** Function: "
                     << F.getName() << '\n');
 
   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   visit(F);
   return true;
 }

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

void LowerTypeTestsModule::replaceCfiUses(Function *Old, Value *New, bool IsDefinition) {
   }
 
   // Process operand replacement of saved constants.
   for (auto *C : Constants)
     C->handleOperandChange(Old, New);
 }
 
 void LowerTypeTestsModule::replaceDirectCalls(Value *Old, Value *New) {
-  auto UI = Old->use_begin(), E = Old->use_end();
-  for (; UI != E;) {
-    Use &U = *UI;
-    ++UI;
-
-    if (!isDirectCall(U))
-      continue;
-
-    U.set(New);
-  }
+  Old->replaceUsesWithIf(New, [](Use &U) { return isDirectCall(U); });
 }
 
 bool LowerTypeTestsModule::lower() {
   // If only some of the modules were split, we cannot correctly perform
   // this transformation. We already checked for the presense of type tests
   // with partially split modules during the thin link, and would have emitted
   // an error if any were found, so here we can simply return.
   if ((ExportSummary && ExportSummary->partiallySplitLTOUnits()) ||

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

for (unsigned N = 0; N < Allocas.size(); ++N) {
     // Replace uses of the alloca with tagged address.
     Value *Tag = getAllocaTag(IRB, StackTag, AI, N);
     Value *AILong = IRB.CreatePointerCast(AI, IntptrTy);
     Value *Replacement = tagPointer(IRB, AI->getType(), AILong, Tag);
     std::string Name =
         AI->hasName() ? AI->getName().str() : "alloca." + itostr(N);
     Replacement->setName(Name + ".hwasan");
 
-    for (auto UI = AI->use_begin(), UE = AI->use_end(); UI != UE;) {
-      Use &U = *UI++;
-      if (U.getUser() != AILong)
-        U.set(Replacement);
-    }
+    AI->replaceUsesWithIf(Replacement,
+                          [AILong](Use &U) { return U.getUser() != AILong; });
 
     for (auto *DDI : AllocaDeclareMap.lookup(AI)) {
       DIExpression *OldExpr = DDI->getExpression();
       DIExpression *NewExpr = DIExpression::append(
           OldExpr, {dwarf::DW_OP_LLVM_tag_offset, RetagMask(N)});
       DDI->setArgOperand(2, MetadataAsValue::get(*C, NewExpr));
     }
 

llvm/lib/Transforms/Scalar/LoopSink.cpp

for (BasicBlock *N : makeArrayRef(SortedBBsToSinkInto).drop_front(1)) {
     assert(LoopBlockNumber.find(N)->second >
                LoopBlockNumber.find(MoveBB)->second &&
            "BBs not sorted!");
     // Clone I and replace its uses.
     Instruction *IC = I.clone();
     IC->setName(I.getName());
     IC->insertBefore(&*N->getFirstInsertionPt());
     // Replaces uses of I with IC in N
-    for (Value::use_iterator UI = I.use_begin(), UE = I.use_end(); UI != UE;) {
-      Use &U = *UI++;
-      auto *I = cast<Instruction>(U.getUser());
-      if (I->getParent() == N)
-        U.set(IC);
-    }
+    I.replaceUsesWithIf(IC, [N](Use &U) {
+      return cast<Instruction>(U.getUser())->getParent() == N;
+    });
     // Replaces uses of I with IC in blocks dominated by N
     replaceDominatedUsesWith(&I, IC, DT, N);
     LLVM_DEBUG(dbgs() << "Sinking a clone of " << I << " To: " << N->getName()
                       << '\n');
     NumLoopSunkCloned++;
   }
   LLVM_DEBUG(dbgs() << "Sinking " << I << " To: " << MoveBB->getName() << '\n');
   NumLoopSunk++;