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

[Fangrui Song via Phabricator via llvm-commits]

MaskRay updated this revision to Diff 221922.
MaskRay added a comment.

improve test


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
@@ -29,6 +29,16 @@
   return &dummy_arr1[i];
 }
 
+struct CountValueCopyAndMove {
+  CountValueCopyAndMove() = default;
+  CountValueCopyAndMove(const CountValueCopyAndMove &) { copy = 1; }
+  CountValueCopyAndMove(CountValueCopyAndMove &&) { move = 1; }
+  void operator=(const CountValueCopyAndMove &) { ++copy; }
+  void operator=(CountValueCopyAndMove &&) { ++move; }
+  int copy = 0;
+  int move = 0;
+};
+
 /// A test class that tries to check that construction and destruction
 /// occur correctly.
 class CtorTester {
@@ -191,6 +201,35 @@
   EXPECT_EQ(this->getValue(), this->Map[this->getKey()]);
 }
 
+TEST(DenseMapCustomTest, InsertOrAssignTest) {
+  struct A : CountValueCopyAndMove {
+    A(int v) : v(v) {}
+    int v;
+  };
+  DenseMap<int, A> map;
+  int key = 1;
+
+  auto try1 = map.insert_or_assign(0, A(1));
+  EXPECT_TRUE(try1.second);
+  EXPECT_EQ(1, try1.first->second.move);
+  auto try2 = map.insert_or_assign(0, A(2));
+  EXPECT_FALSE(try2.second);
+  EXPECT_EQ(2, try2.first->second.v);
+  EXPECT_EQ(2, try2.first->second.move);
+  EXPECT_EQ(try1.first, try2.first);
+  EXPECT_EQ(0, try1.first->second.copy);
+
+  auto try3 = map.insert_or_assign(key, A(3));
+  EXPECT_TRUE(try3.second);
+  EXPECT_EQ(1, try3.first->second.move);
+  auto try4 = map.insert_or_assign(key, A(4));
+  EXPECT_FALSE(try4.second);
+  EXPECT_EQ(4, try4.first->second.v);
+  EXPECT_EQ(2, try4.first->second.move);
+  EXPECT_EQ(try3.first, try4.first);
+  EXPECT_EQ(0, try3.first->second.copy);
+}
+
 // 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,24 @@
     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.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D67665.221922.patch
Type: text/x-patch
Size: 2897 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20190926/42e28a4f/attachment.bin>