[IR] Allow Value::replaceUsesWithIf() to process constants

[rampitec]

Summary

The change is currently NFC, but exploited by the depending D102954.
Code to handle constants is borrowed from the general implementation
of Value::doRAUW().

[lebedev.ri]

lg

[efriedma]

llvm/lib/IR/Value.cpp

543

I'm suspicious of the way you're iterating over the use list here; incrementing early isn't enough to avoid a use-after-free, I think. The incremented UI might refer to the same constant.

doRAUW repeatedly accesses *use_begin() to avoid issues.

[rampitec]

llvm/lib/IR/Value.cpp

543

I do not think this is currently exploitable as the only caller with constants convertConstantExprsToInstructions() guards it, but this may be true in general. Since not all uses are going to be replaced I cannot do the same thing as doRAUW. Do you think this will resolve the potential issue?

@@ -527,22 +527,26 @@ void Value::replaceNonMetadataUsesWith(Value *New) {

 void Value::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!");

+  SmallPtrSet<Constant *, 8> Visited;
   for (use_iterator UI = use_begin(), E = use_end(); UI != E;) {
     Use &U = *UI;
     ++UI;
+    auto *C = dyn_cast<Constant>(U.getUser());
+    if (C && !Visited.insert(C).second)
+      continue;
     if (!ShouldReplace(U))
       continue;
     // Must handle Constants specially, we cannot call replaceUsesOfWith on a
     // constant because they are uniqued.
-    if (auto *C = dyn_cast<Constant>(U.getUser())) {
+    if (C) {
       if (!isa<GlobalValue>(C)) {
         C->handleOperandChange(this, New);
         continue;
       }
     }
     U.set(New);
   }

[efriedma]

llvm/lib/IR/Value.cpp

543

I don't think that's quite right. The problem is ensuring that UI points to something valid when we execute the ++UI. Adding a "continue" after that doesn't really help.

The most straightforward solution I can think of is to dodge the issue by not calling handleOperandChange with an active use_iterator. Instead, we collect all the constant users into a vector of TrackingVH, then iterate over that.

Also, minor related issue: handleOperandChange updates all the uses in a given Constant, not just the one passed to ShouldReplace.

[rampitec]

llvm/lib/IR/Value.cpp

543

Thanks! A good idea, here is the patch: D105061
For the handleOperandChange() I have left a FIXME for now in the same review. This is a very good point, although not essential for the current use.

llvm/include/llvm/IR/Value.h

#endif
   ///
   /// 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.
+  /// values.
   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!");
-
-    for (use_iterator UI = use_begin(), E = use_end(); UI != E;) {
-      Use &U = *UI;
-      ++UI;
-      if (!ShouldReplace(U))
-        continue;
-      U.set(New);
-    }
-  }
+                         llvm::function_ref<bool(Use &U)> ShouldReplace);
 
   /// 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
-  /// values or constant users.
+  /// values.
   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:
   // The methods that start with materialized_ check the uses that are

llvm/lib/IR/Value.cpp

 void Value::replaceAllUsesWith(Value *New) {
   doRAUW(New, ReplaceMetadataUses::Yes);
 }
 
 void Value::replaceNonMetadataUsesWith(Value *New) {
   doRAUW(New, ReplaceMetadataUses::No);
 }
 
+void Value::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!");
+
+  for (use_iterator UI = use_begin(), E = use_end(); UI != E;) {
+    Use &U = *UI;
+    ++UI;
+    if (!ShouldReplace(U))
+      continue;
+    // Must handle Constants specially, we cannot call replaceUsesOfWith on a
+    // constant because they are uniqued.
+    if (auto *C = dyn_cast<Constant>(U.getUser())) {
+      if (!isa<GlobalValue>(C)) {
+        C->handleOperandChange(this, New);
+        continue;
+      }
+    }
+    U.set(New);
+  }
+}
+
 /// Replace llvm.dbg.* uses of MetadataAsValue(ValueAsMetadata(V)) outside BB
 /// with New.
 static void replaceDbgUsesOutsideBlock(Value *V, Value *New, BasicBlock *BB) {
   SmallVector<DbgVariableIntrinsic *> DbgUsers;
   findDbgUsers(DbgUsers, V);
   for (auto *DVI : DbgUsers) {
     if (DVI->getParent() != BB)
       DVI->replaceVariableLocationOp(V, New);