[llvm] r372813 - [ADT] Add StringMap::insert_or_assign

[Fangrui Song via llvm-commits]

Author: maskray
Date: Tue Sep 24 21:58:02 2019
New Revision: 372813

URL: http://llvm.org/viewvc/llvm-project?rev=372813&view=rev
Log:
[ADT] Add StringMap::insert_or_assign

Summary: Similar to std::unordered_map::insert_or_assign

Reviewers: alexshap, bkramer, dblaikie, lhames

Subscribers: jkorous, dexonsmith, kristina, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D67668

Modified:
    llvm/trunk/include/llvm/ADT/StringMap.h
    llvm/trunk/unittests/ADT/StringMapTest.cpp

Modified: llvm/trunk/include/llvm/ADT/StringMap.h
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/include/llvm/ADT/StringMap.h?rev=372813&r1=372812&r2=372813&view=diff
==============================================================================
--- llvm/trunk/include/llvm/ADT/StringMap.h (original)
+++ llvm/trunk/include/llvm/ADT/StringMap.h Tue Sep 24 21:58:02 2019
@@ -391,6 +391,16 @@ public:
     return try_emplace(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(StringRef Key, V &&Val) {
+    auto Ret = try_emplace(Key, std::forward<V>(Val));
+    if (!Ret.second)
+      Ret.first->second = std::forward<V>(Val);
+    return Ret;
+  }
+
   /// Emplace a new element for the specified key into the map if the key isn't
   /// already in the map. The bool component of the returned pair is true
   /// if and only if the insertion takes place, and the iterator component of

Modified: llvm/trunk/unittests/ADT/StringMapTest.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/unittests/ADT/StringMapTest.cpp?rev=372813&r1=372812&r2=372813&view=diff
==============================================================================
--- llvm/trunk/unittests/ADT/StringMapTest.cpp (original)
+++ llvm/trunk/unittests/ADT/StringMapTest.cpp Tue Sep 24 21:58:02 2019
@@ -75,6 +75,16 @@ const char* StringMapTest::testKeyFirst
 size_t StringMapTest::testKeyLength = sizeof(testKey) - 1;
 const std::string StringMapTest::testKeyStr(testKey);
 
+struct CountCopyAndMove {
+  CountCopyAndMove() = default;
+  CountCopyAndMove(const CountCopyAndMove &) { copy = 1; }
+  CountCopyAndMove(CountCopyAndMove &&) { move = 1; }
+  void operator=(const CountCopyAndMove &) { ++copy; }
+  void operator=(CountCopyAndMove &&) { ++move; }
+  int copy = 0;
+  int move = 0;
+};
+
 // Empty map tests.
 TEST_F(StringMapTest, EmptyMapTest) {
   assertEmptyMap();
@@ -269,6 +279,27 @@ TEST_F(StringMapTest, InsertRehashingPai
   EXPECT_EQ(42u, It->second);
 }
 
+TEST_F(StringMapTest, InsertOrAssignTest) {
+  struct A : CountCopyAndMove {
+    A(int v) : v(v) {}
+    int v;
+  };
+  StringMap<A> t(0);
+
+  auto try1 = t.insert_or_assign("A", A(1));
+  EXPECT_TRUE(try1.second);
+  EXPECT_EQ(1, try1.first->second.v);
+  EXPECT_EQ(1, try1.first->second.move);
+
+  auto try2 = t.insert_or_assign("A", A(2));
+  EXPECT_FALSE(try2.second);
+  EXPECT_EQ(2, try2.first->second.v);
+  EXPECT_EQ(2, try1.first->second.move);
+
+  EXPECT_EQ(try1.first, try2.first);
+  EXPECT_EQ(0, try1.first->second.copy);
+}
+
 TEST_F(StringMapTest, IterMapKeys) {
   StringMap<int> Map;
   Map["A"] = 1;