Add support for DenseMap/DenseSet count and find using const pointers

[• dberlin]

Summary

Similar to SmallPtrSet, this makes find and count work with both const
referneces and const pointers.

[dblaikie]

Probably nice to generalize this more fully to support heterogenous lookup (added to standard C++ containers in C++14), but that's probably sufficiently more work not to hold this common case up behind it.

llvm/trunk/include/llvm/ADT/DenseMap.h

 template <
     typename KeyT, typename ValueT, typename KeyInfoT = DenseMapInfo<KeyT>,
     typename Bucket = detail::DenseMapPair<KeyT, ValueT>, bool IsConst = false>
 class DenseMapIterator;
 
 template <typename DerivedT, typename KeyT, typename ValueT, typename KeyInfoT,
           typename BucketT>
 class DenseMapBase : public DebugEpochBase {
+  template <typename T>
+  using const_arg_type_t = typename const_pointer_or_const_ref<T>::type;
+
 public:
   typedef unsigned size_type;
   typedef KeyT key_type;
   typedef ValueT mapped_type;
   typedef BucketT value_type;
 
   typedef DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT> iterator;
   typedef DenseMapIterator<KeyT, ValueT, KeyInfoT, BucketT, true>
for (BucketT *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P) {
       }
     }
     assert(NumEntries == 0 && "Node count imbalance!");
     setNumEntries(0);
     setNumTombstones(0);
   }
 
   /// Return 1 if the specified key is in the map, 0 otherwise.
-  size_type count(const KeyT &Val) const {
+  size_type count(const_arg_type_t<KeyT> Val) const {
     const BucketT *TheBucket;
     return LookupBucketFor(Val, TheBucket) ? 1 : 0;
   }
 
-  iterator find(const KeyT &Val) {
+  iterator find(const_arg_type_t<KeyT> Val) {
     BucketT *TheBucket;
     if (LookupBucketFor(Val, TheBucket))
       return iterator(TheBucket, getBucketsEnd(), *this, true);
     return end();
   }
-  const_iterator find(const KeyT &Val) const {
+  const_iterator find(const_arg_type_t<KeyT> Val) const {
     const BucketT *TheBucket;
     if (LookupBucketFor(Val, TheBucket))
       return const_iterator(TheBucket, getBucketsEnd(), *this, true);
     return end();
   }
 
   /// Alternate version of find() which allows a different, and possibly
   /// less expensive, key type.
const_iterator find_as(const LookupKeyT &Val) const {
     const BucketT *TheBucket;
     if (LookupBucketFor(Val, TheBucket))
       return const_iterator(TheBucket, getBucketsEnd(), *this, true);
     return end();
   }
 
   /// lookup - Return the entry for the specified key, or a default
   /// constructed value if no such entry exists.
-  ValueT lookup(const KeyT &Val) const {
+  ValueT lookup(const_arg_type_t<KeyT> Val) const {
     const BucketT *TheBucket;
     if (LookupBucketFor(Val, TheBucket))
       return TheBucket->getSecond();
     return ValueT();
   }
 
   // Inserts key,value pair into the map if the key isn't already in the map.
   // If the key is already in the map, it returns false and doesn't update the

llvm/trunk/include/llvm/ADT/DenseSet.h

 ///
 /// or the equivalent SmallDenseMap type.  ValueInfoT must implement the
 /// DenseMapInfo "concept".
 template <typename ValueT, typename MapTy, typename ValueInfoT>
 class DenseSetImpl {
   static_assert(sizeof(typename MapTy::value_type) == sizeof(ValueT),
                 "DenseMap buckets unexpectedly large!");
   MapTy TheMap;
+  template <typename T>
+  using const_arg_type_t = typename const_pointer_or_const_ref<T>::type;
 
 public:
   typedef ValueT key_type;
   typedef ValueT value_type;
   typedef unsigned size_type;
 
   explicit DenseSetImpl(unsigned InitialReserve = 0) : TheMap(InitialReserve) {}
 
public:
   /// before resizing again.
   void reserve(size_t Size) { TheMap.reserve(Size); }
 
   void clear() {
     TheMap.clear();
   }
 
   /// Return 1 if the specified key is in the set, 0 otherwise.
-  size_type count(const ValueT &V) const {
+  size_type count(const_arg_type_t<ValueT> V) const {
     return TheMap.count(V);
   }
 
   bool erase(const ValueT &V) {
     return TheMap.erase(V);
   }
 
   void swap(DenseSetImpl &RHS) { TheMap.swap(RHS.TheMap); }
public:
   typedef ConstIterator const_iterator;
 
   iterator begin() { return Iterator(TheMap.begin()); }
   iterator end() { return Iterator(TheMap.end()); }
 
   const_iterator begin() const { return ConstIterator(TheMap.begin()); }
   const_iterator end() const { return ConstIterator(TheMap.end()); }
 
-  iterator find(const ValueT &V) { return Iterator(TheMap.find(V)); }
-  const_iterator find(const ValueT &V) const {
+  iterator find(const_arg_type_t<ValueT> V) { return Iterator(TheMap.find(V)); }
+  const_iterator find(const_arg_type_t<ValueT> V) const {
     return ConstIterator(TheMap.find(V));
   }
 
   /// Alternative version of find() which allows a different, and possibly less
   /// expensive, key type.
   /// The DenseMapInfo is responsible for supplying methods
   /// getHashValue(LookupKeyT) and isEqual(LookupKeyT, KeyT) for each key type
   /// used.

llvm/trunk/include/llvm/Support/type_traits.h

 struct add_const_past_pointer { typedef const T type; };
 
 template <typename T>
 struct add_const_past_pointer<
     T, typename std::enable_if<std::is_pointer<T>::value>::type> {
   typedef const typename std::remove_pointer<T>::type *type;
 };
 
+template <typename T, typename Enable = void>
+struct const_pointer_or_const_ref {
+  using type = const T &;
+};
+template <typename T>
+struct const_pointer_or_const_ref<
+    T, typename std::enable_if<std::is_pointer<T>::value>::type> {
+  using type = typename add_const_past_pointer<T>::type;
+};
 }
 
 // If the compiler supports detecting whether a class is final, define
 // an LLVM_IS_FINAL macro. If it cannot be defined properly, this
 // macro will be left undefined.
 #if __cplusplus >= 201402L
 #define LLVM_IS_FINAL(Ty) std::is_final<Ty>()
 #elif __has_feature(is_final) || LLVM_GNUC_PREREQ(4, 7, 0)
 #define LLVM_IS_FINAL(Ty) __is_final(Ty)
 #endif
 
 #ifdef LLVM_DEFINED_HAS_FEATURE
 #undef __has_feature
 #endif
 
 #endif

llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp

for (auto &I : MissingVals) {
     // don't already have an incoming value for the middle block.
     if (PHI->getBasicBlockIndex(MiddleBlock) == -1)
       PHI->addIncoming(I.second, MiddleBlock);
   }
 }
 
 namespace {
 struct CSEDenseMapInfo {
-  static bool canHandle(Instruction *I) {
+  static bool canHandle(const Instruction *I) {
     return isa<InsertElementInst>(I) || isa<ExtractElementInst>(I) ||
            isa<ShuffleVectorInst>(I) || isa<GetElementPtrInst>(I);
   }
   static inline Instruction *getEmptyKey() {
     return DenseMapInfo<Instruction *>::getEmptyKey();
   }
   static inline Instruction *getTombstoneKey() {
     return DenseMapInfo<Instruction *>::getTombstoneKey();
   }
-  static unsigned getHashValue(Instruction *I) {
+  static unsigned getHashValue(const Instruction *I) {
     assert(canHandle(I) && "Unknown instruction!");
     return hash_combine(I->getOpcode(), hash_combine_range(I->value_op_begin(),
                                                            I->value_op_end()));
   }
-  static bool isEqual(Instruction *LHS, Instruction *RHS) {
+  static bool isEqual(const Instruction *LHS, const Instruction *RHS) {
     if (LHS == getEmptyKey() || RHS == getEmptyKey() ||
         LHS == getTombstoneKey() || RHS == getTombstoneKey())
       return LHS == RHS;
     return LHS->isIdenticalTo(RHS);
   }
 };
 }
 

llvm/trunk/unittests/ADT/DenseMapTest.cpp

TEST(DenseMapCustomTest, TryEmplaceTest) {
   std::unique_ptr<int> P(new int(2));
   auto Try1 = Map.try_emplace(0, new int(1));
   EXPECT_TRUE(Try1.second);
   auto Try2 = Map.try_emplace(0, std::move(P));
   EXPECT_FALSE(Try2.second);
   EXPECT_EQ(Try1.first, Try2.first);
   EXPECT_NE(nullptr, P);
 }
+
+TEST(DenseMapCustomTest, ConstTest) {
+  // Test that const pointers work okay for count and find, even when the
+  // underlying map is a non-const pointer.
+  DenseMap<int *, int> Map;
+  int A;
+  int *B = &A;
+  const int *C = &A;
+  Map.insert({B, 0});
+  EXPECT_EQ(Map.count(B), 1u);
+  EXPECT_EQ(Map.count(C), 1u);
+  EXPECT_NE(Map.find(B), Map.end());
+  EXPECT_NE(Map.find(C), Map.end());
+}
 }

llvm/trunk/unittests/ADT/DenseSetTest.cpp

for (auto Size : {1, 2, 48, 66}) {
     // Check that we didn't grow
     EXPECT_EQ(MemorySize, Set.getMemorySize());
     // Check that move was called the expected number of times
     EXPECT_EQ(Size, CountCopyAndMove::Move);
     // Check that no copy occured
     EXPECT_EQ(0, CountCopyAndMove::Copy);
   }
 }
+TEST(DenseSetCustomTest, ConstTest) {
+  // Test that const pointers work okay for count and find, even when the
+  // underlying map is a non-const pointer.
+  DenseSet<int *> Map;
+  int A;
+  int *B = &A;
+  const int *C = &A;
+  Map.insert(B);
+  EXPECT_EQ(Map.count(B), 1u);
+  EXPECT_EQ(Map.count(C), 1u);
+  EXPECT_NE(Map.find(B), Map.end());
+  EXPECT_NE(Map.find(C), Map.end());
+}
 }