[ADT] Add 2 DenseMap::insert_or_assign overloads

[MaskRay]

Summary

I plan to use one overload in an llvm-objcopy patch.

[dexonsmith]

include/llvm/ADT/DenseMap.h

219–222

This looks accidentally recursive to me, and I don't see test coverage for it.

[MaskRay]

Fix std::pair<iterator, bool> insert_or_assign(const KeyT &Key, V &&Val)
and add tests

[dblaikie]

include/llvm/ADT/DenseMap.h

206–231

Is this significantly cheaper than - say, try_emplace + assignment? (the latter seems like it'd keep the implementation simpler/less risky)

auto {Iter, Inserted} = m.try_emplace(key, value);
if (!Inserted)
  Iter->second = value

& at that point, should it perhaps be a utility, rather than reimplementing it in every container?

226–227

This doesn't look like assignment - seems like it's inconsistent with the name of the function and the functionality of the standard version of this function?

unittests/ADT/DenseMapTest.cpp

195

I'd probably test this with something simpler than std::unique_ptr - to make the test more targeted/clear as to what it's testing.

198–199

Does this leak memory?

[MaskRay]

Address review comments

include/llvm/ADT/DenseMap.h

206–231

Thanks! I can do:

BucketT *B;
bool Added = !LookupBucketFor(Key, B);
if (Added)
  B = InsertIntoBucket(B, std::move(Key), std::forward<V>(Val));
else
  B->getSecond() = std::forward<V>(Val);
return std::make_pair(makeIterator(B, getBucketsEnd(), *this, true), Added);

but that seems no better than:

auto Ret = try_emplace(std::move(Key), std::forward<V>(Val));
if (!Ret.second)
  Ret.first->second = std::forward<V>(Val);
return Ret;

So I'll use the simple version.

unittests/ADT/DenseMapTest.cpp

195

I agree it is less clear but it has a nice property. std::unique_ptr<int> as the value type can be used to detect memory leaks in a -DLLVM_USE_SANITIZER=Address build.

[MaskRay]

Add a comment

[MaskRay]

unique_ptr<int> -> NonDefaultConstructible

[MaskRay]

improve test

include/llvm/ADT/DenseMap.h

public:
 
   // 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
   // value.
   std::pair<iterator, bool> insert(std::pair<KeyT, ValueT> &&KV) {
     return try_emplace(std::move(KV.first), std::move(KV.second));
   }
 
+  template <typename V>
+  std::pair<iterator, bool> insert_or_assign(KeyT &&Key, V &&Val) {
+    auto Ret = try_emplace(std::move(Key), std::forward<V>(Val));
+    if (!Ret.second)
+      Ret.first->second = std::forward<V>(Val);
+    return Ret;
+  }
+
+  template <typename V>
+  std::pair<iterator, bool> insert_or_assign(const KeyT &Key, V &&Val) {
+    auto Ret = try_emplace(Key, std::forward<V>(Val));
+    if (!Ret.second)
+      Ret.first->second = std::forward<V>(Val);
+    return Ret;
+  }
+
   // Inserts key,value pair into the map if the key isn't already in the map.
   // The value is constructed in-place if the key is not in the map, otherwise
   // it is not moved.
   template <typename... Ts>
   std::pair<iterator, bool> try_emplace(KeyT &&Key, Ts &&... Args) {
     BucketT *TheBucket;
     if (LookupBucketFor(Key, TheBucket))
       return std::make_pair(
                makeIterator(TheBucket, getBucketsEnd(), *this, true),
                false); // Already in map.
 
     // Otherwise, insert the new element.
     TheBucket =
         InsertIntoBucket(TheBucket, std::move(Key), std::forward<Ts>(Args)...);
     return std::make_pair(
              makeIterator(TheBucket, getBucketsEnd(), *this, true),
              true);
   }

unittests/ADT/DenseMapTest.cpp

 
 // Test insert() method
 TYPED_TEST(DenseMapTest, InsertTest) {
   this->Map.insert(std::make_pair(this->getKey(), this->getValue()));
   EXPECT_EQ(1u, this->Map.size());
   EXPECT_EQ(this->getValue(), this->Map[this->getKey()]);
 }
 
+TEST(DenseMapCustomTest, InsertOrAssignTest) {
+  DenseMap<int, std::unique_ptr<int>> Map;
+  auto Try1 = Map.insert_or_assign(0, std::make_unique<int>(1));
+  EXPECT_TRUE(Try1.second);
+  auto Try2 = Map.insert_or_assign(0, std::make_unique<int>(2));
+  EXPECT_FALSE(Try2.second);
+  EXPECT_EQ(2, *Try2.first->getSecond());
+  EXPECT_EQ(Try1.first, Try2.first);
+
+  int Key = 1;
+  auto Try3 = Map.insert_or_assign(Key, std::make_unique<int>(3));
+  EXPECT_TRUE(Try3.second);
+  auto Try4 = Map.insert_or_assign(Key, std::make_unique<int>(4));
+  EXPECT_FALSE(Try4.second);
+  EXPECT_EQ(4, *Try4.first->getSecond());
+  EXPECT_EQ(Try3.first, Try4.first);
+}
+
 // Test copy constructor method
 TYPED_TEST(DenseMapTest, CopyConstructorTest) {
   this->Map[this->getKey()] = this->getValue();
   TypeParam copyMap(this->Map);
 
   EXPECT_EQ(1u, copyMap.size());
   EXPECT_EQ(this->getValue(), copyMap[this->getKey()]);
 }