[ADT] Implement {DenseMap,MapVector,StringMap}::contains

[kazu]

Summary

This patch implements the C++20-style contains() for DenseMap,
MapVector, and StringMap.

With this patch, every set and map container type that has count()
also has contains().

[kuhar]

LGTM.

Separately, it would be cool if we could make llvm::is_contained use the contains member when available, instead of doing a linear scan. It seems like a bit of a foot gun to me to have 2 functions with almost identical names but vastly different performance characteristics.

[kazu]

In D145895#4189370, @kuhar wrote:

LGTM.

Separately, it would be cool if we could make llvm::is_contained use the contains member when available, instead of doing a linear scan. It seems like a bit of a foot gun to me to have 2 functions with almost identical names but vastly different performance characteristics.

Thanks for the reivew! I might just migrate those to $set.contains() or $map.contains(). In theory, it's possible to write a template that calls llvm::is_contained so that it can work with either an array or a set, but I don't know how common that is.

With some template tricks, we could implement a check for the contains() method. If so, we could do static_assert(!sizeof(T*), "Use contains() instead.") or something, at least in our working source tree, maybe not in production.

llvm/include/llvm/ADT/DenseMap.h

if (std::is_trivially_destructible<ValueT>::value) {
       }
       assert(NumEntries == 0 && "Node count imbalance!");
       (void)NumEntries;
     }
     setNumEntries(0);
     setNumTombstones(0);
   }
 
+  /// Return true if the specified key is in the map, false otherwise.
+  bool contains(const_arg_type_t<KeyT> Val) const {
+    const BucketT *TheBucket;
+    return LookupBucketFor(Val, TheBucket);
+  }
+
   /// Return 1 if the specified key is in the map, 0 otherwise.
   size_type count(const_arg_type_t<KeyT> Val) const {
-    const BucketT *TheBucket;
-    return LookupBucketFor(Val, TheBucket) ? 1 : 0;
+    return contains(Val) ? 1 : 0;
   }
 
   iterator find(const_arg_type_t<KeyT> Val) {
     BucketT *TheBucket;
     if (LookupBucketFor(Val, TheBucket))
       return makeIterator(TheBucket,
                           shouldReverseIterate<KeyT>() ? getBuckets()
                                                        : getBucketsEnd(),

llvm/include/llvm/ADT/MapVector.h

std::pair<iterator, bool> insert(std::pair<KeyT, ValueT> &&KV) {
     if (Result.second) {
       Vector.push_back(std::move(KV));
       I = Vector.size() - 1;
       return std::make_pair(std::prev(end()), true);
     }
     return std::make_pair(begin() + I, false);
   }
 
-  size_type count(const KeyT &Key) const {
-    typename MapType::const_iterator Pos = Map.find(Key);
-    return Pos == Map.end()? 0 : 1;
-  }
+  bool contains(const KeyT &Key) const { return Map.find(Key) != Map.end(); }
+
+  size_type count(const KeyT &Key) const { return contains(Key) ? 1 : 0; }
 
   iterator find(const KeyT &Key) {
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end()? Vector.end() :
                             (Vector.begin() + Pos->second);
   }
 
   const_iterator find(const KeyT &Key) const {

llvm/include/llvm/ADT/StringMap.h

const ValueTy &at(StringRef Val) const {
     assert(Iter != this->end() && "StringMap::at failed due to a missing key");
     return Iter->second;
   }
 
   /// Lookup the ValueTy for the \p Key, or create a default constructed value
   /// if the key is not in the map.
   ValueTy &operator[](StringRef Key) { return try_emplace(Key).first->second; }
 
+  /// contains - Return true if the element is in the map, false otherwise.
+  bool contains(StringRef Key) const { return find(Key) != end(); }
+
   /// count - Return 1 if the element is in the map, 0 otherwise.
-  size_type count(StringRef Key) const { return find(Key) == end() ? 0 : 1; }
+  size_type count(StringRef Key) const { return contains(Key) ? 1 : 0; }
 
   template <typename InputTy>
   size_type count(const StringMapEntry<InputTy> &MapEntry) const {
     return count(MapEntry.getKey());
   }
 
   /// equal - check whether both of the containers are equal.
   bool operator==(const StringMap &RHS) const {

llvm/unittests/ADT/DenseMapTest.cpp

TYPED_TEST(DenseMapTest, EmptyIntMapTest) {
   EXPECT_EQ(0u, this->Map.size());
   EXPECT_TRUE(this->Map.empty());
 
   // Iterator tests
   EXPECT_TRUE(this->Map.begin() == this->Map.end());
 
   // Lookup tests
   EXPECT_FALSE(this->Map.count(this->getKey()));
+  EXPECT_FALSE(this->Map.contains(this->getKey()));
   EXPECT_TRUE(this->Map.find(this->getKey()) == this->Map.end());
   EXPECT_EQ(typename TypeParam::mapped_type(),
             this->Map.lookup(this->getKey()));
 }
 
 // Constant map tests
 TYPED_TEST(DenseMapTest, ConstEmptyMapTest) {
   const TypeParam &ConstMap = this->Map;
TYPED_TEST(DenseMapTest, SingleEntryMapTest) {
   typename TypeParam::iterator it = this->Map.begin();
   EXPECT_EQ(this->getKey(), it->first);
   EXPECT_EQ(this->getValue(), it->second);
   ++it;
   EXPECT_TRUE(it == this->Map.end());
 
   // Lookup tests
   EXPECT_TRUE(this->Map.count(this->getKey()));
+  EXPECT_TRUE(this->Map.contains(this->getKey()));
   EXPECT_TRUE(this->Map.find(this->getKey()) == this->Map.begin());
   EXPECT_EQ(this->getValue(), this->Map.lookup(this->getKey()));
   EXPECT_EQ(this->getValue(), this->Map[this->getKey()]);
 }
 
 TYPED_TEST(DenseMapTest, AtTest) {
   this->Map[this->getKey(0)] = this->getValue(0);
   this->Map[this->getKey(1)] = this->getValue(1);

llvm/unittests/ADT/MapVectorTest.cpp

 TEST(MapVectorTest, erase) {
   MapVector<int, int> MV;
 
   MV.insert(std::make_pair(1, 2));
   MV.insert(std::make_pair(3, 4));
   MV.insert(std::make_pair(5, 6));
   ASSERT_EQ(MV.size(), 3u);
 
+  ASSERT_TRUE(MV.contains(1));
   MV.erase(MV.find(1));
   ASSERT_EQ(MV.size(), 2u);
+  ASSERT_FALSE(MV.contains(1));
   ASSERT_EQ(MV.find(1), MV.end());
   ASSERT_EQ(MV[3], 4);
   ASSERT_EQ(MV[5], 6);
 
   ASSERT_EQ(MV.erase(3), 1u);
   ASSERT_EQ(MV.size(), 1u);
   ASSERT_EQ(MV.find(3), MV.end());
   ASSERT_EQ(MV[5], 6);

llvm/unittests/ADT/StringMapTest.cpp

void assertEmptyMap() {
     // Size tests
     EXPECT_EQ(0u, testMap.size());
     EXPECT_TRUE(testMap.empty());
 
     // Iterator tests
     EXPECT_TRUE(testMap.begin() == testMap.end());
 
     // Lookup tests
+    EXPECT_FALSE(testMap.contains(testKey));
     EXPECT_EQ(0u, testMap.count(testKey));
     EXPECT_EQ(0u, testMap.count(StringRef(testKeyFirst, testKeyLength)));
     EXPECT_EQ(0u, testMap.count(testKeyStr));
     EXPECT_TRUE(testMap.find(testKey) == testMap.end());
     EXPECT_TRUE(testMap.find(StringRef(testKeyFirst, testKeyLength)) ==
                 testMap.end());
     EXPECT_TRUE(testMap.find(testKeyStr) == testMap.end());
   }
 
   void assertSingleItemMap() {
     // Size tests
     EXPECT_EQ(1u, testMap.size());
     EXPECT_FALSE(testMap.begin() == testMap.end());
     EXPECT_FALSE(testMap.empty());
 
     // Iterator tests
     StringMap<uint32_t>::iterator it = testMap.begin();
     EXPECT_STREQ(testKey, it->first().data());
     EXPECT_EQ(testValue, it->second);
     ++it;
     EXPECT_TRUE(it == testMap.end());
 
     // Lookup tests
+    EXPECT_TRUE(testMap.contains(testKey));
     EXPECT_EQ(1u, testMap.count(testKey));
     EXPECT_EQ(1u, testMap.count(StringRef(testKeyFirst, testKeyLength)));
     EXPECT_EQ(1u, testMap.count(testKeyStr));
     EXPECT_TRUE(testMap.find(testKey) == testMap.begin());
     EXPECT_TRUE(testMap.find(StringRef(testKeyFirst, testKeyLength)) ==
                 testMap.begin());
     EXPECT_TRUE(testMap.find(testKeyStr) == testMap.begin());
   }