[PATCH] D67665: [ADT] Add 2 DenseMap::insert_or_assign overloads

[Fangrui Song via Phabricator via llvm-commits]

MaskRay updated this revision to Diff 220531.
MaskRay added a reviewer: dexonsmith.
MaskRay removed a subscriber: dexonsmith.
MaskRay added a comment.

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


Repository:
  rL LLVM

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D67665/new/

https://reviews.llvm.org/D67665

Files:
  include/llvm/ADT/DenseMap.h
  unittests/ADT/DenseMapTest.cpp


Index: unittests/ADT/DenseMapTest.cpp
===================================================================
--- unittests/ADT/DenseMapTest.cpp
+++ unittests/ADT/DenseMapTest.cpp
@@ -191,6 +191,24 @@
   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, new int(1));
+  EXPECT_TRUE(Try1.second);
+  auto Try2 = Map.insert_or_assign(0, new 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, new int(3));
+  EXPECT_TRUE(Try3.second);
+  auto Try4 = Map.insert_or_assign(Key, new 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();
Index: include/llvm/ADT/DenseMap.h
===================================================================
--- include/llvm/ADT/DenseMap.h
+++ include/llvm/ADT/DenseMap.h
@@ -203,6 +203,32 @@
     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) {
+    BucketT *B;
+    bool Added = !LookupBucketFor(Key, B);
+    if (Added) {
+      B = InsertIntoBucket(B, std::move(Key), std::forward<V>(Val));
+    } else {
+      B->getSecond().~ValueT();
+      ::new (&B->getSecond()) ValueT(std::forward<V>(Val));
+    }
+    return std::make_pair(makeIterator(B, getBucketsEnd(), *this, true), Added);
+  }
+
+  template <typename V>
+  std::pair<iterator, bool> insert_or_assign(const KeyT &Key, V &&Val) {
+    BucketT *B;
+    bool Added = !LookupBucketFor(Key, B);
+    if (Added) {
+      B = InsertIntoBucket(B, Key, std::forward<V>(Val));
+    } else {
+      B->getSecond().~ValueT();
+      ::new (&B->getSecond()) ValueT(std::forward<V>(Val));
+    }
+    return std::make_pair(makeIterator(B, getBucketsEnd(), *this, true), Added);
+  }
+
   // 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.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D67665.220531.patch
Type: text/x-patch
Size: 2332 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20190917/102f08c5/attachment.bin>