[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

200

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.

203–204

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

200

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));
   }
 
+  /// Inserts an element or assigns to the current element if the key already
+  /// exists. The return type is the same as try_emplace.
+  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),

unittests/ADT/DenseMapTest.cpp

uint32_t *getTestKey(int i, uint32_t **) {
   return &dummy_arr1[i];
 }
 uint32_t *getTestValue(int i, uint32_t **) {
   static uint32_t dummy_arr1[8192];
   assert(i < 8192 && "Only support 8192 dummy keys.");
   return &dummy_arr1[i];
 }
 
+struct NotDefaultConstructible {
+  NotDefaultConstructible() = delete;
+  int v;
+};
+
 /// A test class that tries to check that construction and destruction
 /// occur correctly.
 class CtorTester {
   static std::set<CtorTester *> Constructed;
   int Value;
 
 public:
   explicit CtorTester(int Value = 0) : Value(Value) {

 
 // 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, NotDefaultConstructible> Map;
+  auto Try1 = Map.insert_or_assign(0, NotDefaultConstructible{1});
+  EXPECT_TRUE(Try1.second);
+  auto Try2 = Map.insert_or_assign(0, NotDefaultConstructible{2});
+  EXPECT_FALSE(Try2.second);
+  EXPECT_EQ(2, Try2.first->second.v);
+  EXPECT_EQ(Try1.first, Try2.first);
+
+  int Key = 1;
+  auto Try3 = Map.insert_or_assign(Key, NotDefaultConstructible{3});
+  EXPECT_TRUE(Try3.second);
+  auto Try4 = Map.insert_or_assign(Key, NotDefaultConstructible{4});
+  EXPECT_FALSE(Try4.second);
+  EXPECT_EQ(4, Try4.first->second.v);
+  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()]);
 }