Index: llvm/trunk/include/llvm/ADT/StringMap.h =================================================================== --- llvm/trunk/include/llvm/ADT/StringMap.h +++ llvm/trunk/include/llvm/ADT/StringMap.h @@ -391,6 +391,16 @@ 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 + std::pair insert_or_assign(StringRef Key, V &&Val) { + auto Ret = try_emplace(Key, std::forward(Val)); + if (!Ret.second) + Ret.first->second = std::forward(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 Index: llvm/trunk/unittests/ADT/StringMapTest.cpp =================================================================== --- llvm/trunk/unittests/ADT/StringMapTest.cpp +++ llvm/trunk/unittests/ADT/StringMapTest.cpp @@ -75,6 +75,16 @@ 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 @@ EXPECT_EQ(42u, It->second); } +TEST_F(StringMapTest, InsertOrAssignTest) { + struct A : CountCopyAndMove { + A(int v) : v(v) {} + int v; + }; + StringMap 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 Map; Map["A"] = 1;