[PATCH] D67668: [ADT] Add StringMap::insert_or_assign

[Fangrui Song via Phabricator via llvm-commits]

MaskRay created this revision.
MaskRay added reviewers: alexshap, bkramer, dblaikie, lhames.
Herald added subscribers: llvm-commits, kristina, dexonsmith.
Herald added a project: LLVM.

Similar to std::unordered_map::insert_or_assign


Repository:
  rL LLVM

https://reviews.llvm.org/D67668

Files:
  include/llvm/ADT/StringMap.h
  unittests/ADT/StringMapTest.cpp


Index: unittests/ADT/StringMapTest.cpp
===================================================================
--- unittests/ADT/StringMapTest.cpp
+++ unittests/ADT/StringMapTest.cpp
@@ -269,6 +269,17 @@
   EXPECT_EQ(42u, It->second);
 }
 
+TEST_F(StringMapTest, InsertOrAssignTest) {
+  StringMap<int> t(0);
+  auto try1 = t.insert_or_assign("a", 1);
+  EXPECT_TRUE(try1.second);
+  EXPECT_EQ(1, try1.first->second);
+  auto try2 = t.insert_or_assign("a", 2);
+  EXPECT_FALSE(try2.second);
+  EXPECT_EQ(2, try2.first->second);
+  EXPECT_EQ(try1.first, try2.first);
+}
+
 TEST_F(StringMapTest, IterMapKeys) {
   StringMap<int> Map;
   Map["A"] = 1;
Index: include/llvm/ADT/StringMap.h
===================================================================
--- include/llvm/ADT/StringMap.h
+++ include/llvm/ADT/StringMap.h
@@ -391,6 +391,27 @@
     return try_emplace(KV.first, std::move(KV.second));
   }
 
+  template <typename V>
+  std::pair<iterator, bool> insert_or_assign(StringRef Key, V &&Val) {
+    unsigned BucketNo = LookupBucketFor(Key);
+    StringMapEntryBase *&Bucket = TheTable[BucketNo];
+    if (Bucket && Bucket != getTombstoneVal()) {
+      auto &Value = static_cast<MapEntryTy*>(Bucket)->getValue();
+      Value.~ValueTy();
+      ::new (&Value) ValueTy(std::forward<V>(Val));
+      return std::make_pair(iterator(TheTable + BucketNo, false), false);
+    }
+
+    if (Bucket == getTombstoneVal())
+      --NumTombstones;
+    Bucket = MapEntryTy::Create(Key, Allocator, std::forward<V>(Val));
+    ++NumItems;
+    assert(NumItems + NumTombstones <= NumBuckets);
+
+    BucketNo = RehashTable(BucketNo);
+    return std::make_pair(iterator(TheTable + BucketNo, false), true);
+  }
+
   /// 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


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D67668.220528.patch
Type: text/x-patch
Size: 1937 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20190917/df9e8083/attachment.bin>