Fix use-after-free bug in AffectedValueCallbackVH::allUsesReplacedWith

[hfinkel]

Summary

When transferring affected values in the cache from an old value, identified by the value of the current callback, to the specified new value we might need to insert a new entry into the DenseMap which constitutes the cache. Doing so might delete the current callback object. Move the copying logic into a new function, a member of the assumption cache itself, so that we don't run into UB should the callback handle itself be removed mid-copy.

[sanjoy]

The code lgtm, but

• Can you easily test this using a C++ unit test? I'll settle for something that normally runs "fine", but trips asan if the build is using asan.
• getAffectedValues should be renamed to getOrInsertAffectedValues.

lib/Analysis/AssumptionCache.cpp

121 ↗

(On Diff #84505)

Can you use llvm::find here (in a later change)?

132 ↗

(On Diff #84505)

Please add why it may dangle.

[hfinkel]

In D28749#646897, @sanjoy wrote:

The code lgtm, but

• Can you easily test this using a C++ unit test? I'll settle for something that normally runs "fine", but trips asan if the build is using asan.

I'll experiment with this today and commit one in follow-up. I might also be able to create an IR test case for instcombine if things work out just right.

• getAffectedValues should be renamed to getOrInsertAffectedValues.

Sure.

lib/Analysis/AssumptionCache.cpp

121 ↗

(On Diff #84505)

Yes, I think so. I'd forgotten about that.

llvm/trunk/include/llvm/Analysis/AssumptionCache.h

class AssumptionCache {
   /// \brief A map of values about which an assumption might be providing
   /// information to the relevant set of assumptions.
   using AffectedValuesMap =
     DenseMap<AffectedValueCallbackVH, SmallVector<WeakVH, 1>,
              AffectedValueCallbackVH::DMI>;
   AffectedValuesMap AffectedValues;
 
   /// Get the vector of assumptions which affect a value from the cache.
-  SmallVector<WeakVH, 1> &getAffectedValues(Value *V);
+  SmallVector<WeakVH, 1> &getOrInsertAffectedValues(Value *V);
+
+  /// Copy affected values in the cache for OV to be affected values for NV.
+  void copyAffectedValuesInCache(Value *OV, Value *NV);
 
   /// \brief Flag tracking whether we have scanned the function yet.
   ///
   /// We want to be as lazy about this as possible, and so we scan the function
   /// at the last moment.
   bool Scanned;
 
   /// \brief Scan the function for assumptions and add them to the cache.

llvm/trunk/lib/Analysis/AssumptionCache.cpp

 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/Debug.h"
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
-SmallVector<WeakVH, 1> &AssumptionCache::getAffectedValues(Value *V) {
+SmallVector<WeakVH, 1> &AssumptionCache::getOrInsertAffectedValues(Value *V) {
   // Try using find_as first to avoid creating extra value handles just for the
   // purpose of doing the lookup.
   auto AVI = AffectedValues.find_as(V);
   if (AVI != AffectedValues.end())
     return AVI->second;
 
   auto AVIP = AffectedValues.insert({
       AffectedValueCallbackVH(V, this), SmallVector<WeakVH, 1>()});
if (Pred == ICmpInst::ICMP_EQ) {
       };
 
       AddAffectedFromEq(A);
       AddAffectedFromEq(B);
     }
   }
 
   for (auto &AV : Affected) {
-    auto &AVV = getAffectedValues(AV);
+    auto &AVV = getOrInsertAffectedValues(AV);
     if (std::find(AVV.begin(), AVV.end(), CI) == AVV.end())
       AVV.push_back(CI);
   }
 }
 
 void AssumptionCache::AffectedValueCallbackVH::deleted() {
   auto AVI = AC->AffectedValues.find(getValPtr());
   if (AVI != AC->AffectedValues.end())
     AC->AffectedValues.erase(AVI);
   // 'this' now dangles!
 }
 
+void AssumptionCache::copyAffectedValuesInCache(Value *OV, Value *NV) {
+  auto &NAVV = getOrInsertAffectedValues(NV);
+  auto AVI = AffectedValues.find(OV);
+  if (AVI == AffectedValues.end())
+    return;
+
+  for (auto &A : AVI->second)
+    if (std::find(NAVV.begin(), NAVV.end(), A) == NAVV.end())
+      NAVV.push_back(A);
+}
+
 void AssumptionCache::AffectedValueCallbackVH::allUsesReplacedWith(Value *NV) {
   if (!isa<Instruction>(NV) && !isa<Argument>(NV))
     return;
 
   // Any assumptions that affected this value now affect the new value.
 
-  auto &NAVV = AC->getAffectedValues(NV);
-  auto AVI = AC->AffectedValues.find(getValPtr());
-  if (AVI == AC->AffectedValues.end())
-    return;
-
-  for (auto &A : AVI->second)
-    if (std::find(NAVV.begin(), NAVV.end(), A) == NAVV.end())
-      NAVV.push_back(A);
+  AC->copyAffectedValuesInCache(getValPtr(), NV);
+  // 'this' now might dangle! If the AffectedValues map was resized to add an
+  // entry for NV then this object might have been destroyed in favor of some
+  // copy in the grown map.
 }
 
 void AssumptionCache::scanFunction() {
   assert(!Scanned && "Tried to scan the function twice!");
   assert(AssumeHandles.empty() && "Already have assumes when scanning!");
 
   // Go through all instructions in all blocks, add all calls to @llvm.assume
   // to this cache.