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 + std::pair insert_or_assign(StringRef Key, V &&Val) { + unsigned BucketNo = LookupBucketFor(Key); + StringMapEntryBase *&Bucket = TheTable[BucketNo]; + if (Bucket && Bucket != getTombstoneVal()) { + auto &Value = static_cast(Bucket)->getValue(); + Value.~ValueTy(); + ::new (&Value) ValueTy(std::forward(Val)); + return std::make_pair(iterator(TheTable + BucketNo, false), false); + } + + if (Bucket == getTombstoneVal()) + --NumTombstones; + Bucket = MapEntryTy::Create(Key, Allocator, std::forward(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 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 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 Map; Map["A"] = 1;