[IR] Support efficient lookup of AssertingVH/PoisoningVH

[nikic]

Summary

Currently, there doesn't seem to be any way to look up a Value* in a map/set indexed by AssertingVH/PoisoningVH, without creating a value handle, which is fairly expensive, because it involves adding the value handle to the use list and immediately removing it again. Using find_as(Value *) does not work (and is in fact worse than just using find(Value *)), because it will end up creating multiple value handles during the lookup.

For AssertingVH, address this by simply using DenseMapInfo<T *> instead of manually implementing something. The AssertingVH<T> will now get coerced to T*, rather than the other way around.

For PoisoningVH, add extra overloads of getHashValue() and isEqual() that accept a T* argument.

This allows using find_as(Value *) to perform efficient lookups in assertion-enabled builds.

[MaskRay]

Thanks! Seems reasonable. Can you please add a DenseMap test to llvm/unittests/IR/ValueHandleTest.cpp?

llvm/include/llvm/IR/ValueHandle.h

313

public can be omitted.

Reusing DenseMapInfo<T *> is clever and more efficient.

[nikic]

Remove unnecessary public, add some basic tests. These essentially just check that AssertingVH/PoisoningVH inside a DenseMap work, because I don't know how I could test that using find_as() does not construct a temporary value handle.

[MaskRay]

Thanks!

llvm/unittests/IR/ValueHandleTest.cpp

496

{ K, V } -> {K, V}

llvm/include/llvm/IR/ValueHandle.h

ValueTy *operator=(const AssertingVH<ValueTy> &RHS) {
     setValPtr(RHS.getValPtr());
     return getValPtr();
   }
 
   ValueTy *operator->() const { return getValPtr(); }
   ValueTy &operator*() const { return *getValPtr(); }
 };
 
-// Specialize DenseMapInfo to allow AssertingVH to participate in DenseMap.
+// Treat AssertingVH<T> like T* inside maps. This also allows using find_as()
+// to look up a value without constructing a value handle.
 template<typename T>
-struct DenseMapInfo<AssertingVH<T>> {
-  static inline AssertingVH<T> getEmptyKey() {
-    AssertingVH<T> Res;
-    Res.setRawValPtr(DenseMapInfo<Value *>::getEmptyKey());
-    return Res;
-  }
-
-  static inline AssertingVH<T> getTombstoneKey() {
-    AssertingVH<T> Res;
-    Res.setRawValPtr(DenseMapInfo<Value *>::getTombstoneKey());
-    return Res;
-  }
-
-  static unsigned getHashValue(const AssertingVH<T> &Val) {
-    return DenseMapInfo<Value *>::getHashValue(Val.getRawValPtr());
-  }
-
-  static bool isEqual(const AssertingVH<T> &LHS, const AssertingVH<T> &RHS) {
-    return DenseMapInfo<Value *>::isEqual(LHS.getRawValPtr(),
-                                          RHS.getRawValPtr());
-  }
-};
+struct DenseMapInfo<AssertingVH<T>> : DenseMapInfo<T *> {};
 
 /// Value handle that tracks a Value across RAUW.
 ///
 /// TrackingVH is designed for situations where a client needs to hold a handle
 /// to a Value (or subclass) across some operations which may move that value,
 /// but should never destroy it or replace it with some unacceptable type.
 ///
 /// It is an error to attempt to replace a value with one of a type which is
template <typename T> struct DenseMapInfo<PoisoningVH<T>> {
   static unsigned getHashValue(const PoisoningVH<T> &Val) {
     return DenseMapInfo<Value *>::getHashValue(Val.getRawValPtr());
   }
 
   static bool isEqual(const PoisoningVH<T> &LHS, const PoisoningVH<T> &RHS) {
     return DenseMapInfo<Value *>::isEqual(LHS.getRawValPtr(),
                                           RHS.getRawValPtr());
   }
+
+  // Allow lookup by T* via find_as(), without constructing a temporary
+  // value handle.
+
+  static unsigned getHashValue(const T *Val) {
+    return DenseMapInfo<Value *>::getHashValue(Val);
+  }
+
+  static bool isEqual(const T *LHS, const PoisoningVH<T> &RHS) {
+    return DenseMapInfo<Value *>::isEqual(LHS, RHS.getRawValPtr());
+  }
 };
 
 } // end namespace llvm
 
 #endif // LLVM_IR_VALUEHANDLE_H

llvm/unittests/IR/ValueHandleTest.cpp

 }
 
 TEST_F(ValueHandle, PoisoningVH_DoesNotFollowRAUW) {
   PoisoningVH<Value> VH(BitcastV.get());
   BitcastV->replaceAllUsesWith(ConstantV);
   EXPECT_TRUE(DenseMapInfo<PoisoningVH<Value>>::isEqual(VH, BitcastV.get()));
 }
 
+TEST_F(ValueHandle, AssertingVH_DenseMap) {
+  DenseMap<AssertingVH<Value>, int> Map;
+  Map.insert({ BitcastV.get(), 1 });
+  Map.insert({ ConstantV, 2 });
+  // These will create a temporary AssertingVH during lookup.
+  EXPECT_TRUE(Map.find(BitcastV.get()) != Map.end());
+  EXPECT_TRUE(Map.find(ConstantV) != Map.end());
+  // These will not create a temporary AssertingVH.
+  EXPECT_TRUE(Map.find_as(BitcastV.get()) != Map.end());
+  EXPECT_TRUE(Map.find_as(ConstantV) != Map.end());
+}
+
+TEST_F(ValueHandle, PoisoningVH_DenseMap) {
+  DenseMap<PoisoningVH<Value>, int> Map;
+  Map.insert({ BitcastV.get(), 1 });
+  Map.insert({ ConstantV, 2 });
+  // These will create a temporary PoisoningVH during lookup.
+  EXPECT_TRUE(Map.find(BitcastV.get()) != Map.end());
+  EXPECT_TRUE(Map.find(ConstantV) != Map.end());
+  // These will not create a temporary PoisoningVH.
+  EXPECT_TRUE(Map.find_as(BitcastV.get()) != Map.end());
+  EXPECT_TRUE(Map.find_as(ConstantV) != Map.end());
+}
+
 #ifdef NDEBUG
 
 TEST_F(ValueHandle, PoisoningVH_ReducesToPointer) {
   EXPECT_EQ(sizeof(CastInst *), sizeof(PoisoningVH<CastInst>));
 }
 
 #else // !NDEBUG