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 + std::pair insert_or_assign(KeyT &&Key, V &&Val) { + auto Ret = try_emplace(std::move(Key), std::forward(Val)); + if (!Ret.second) + Ret.first->second = std::forward(Val); + return Ret; + } + + template + std::pair insert_or_assign(const KeyT &Key, V &&Val) { + auto Ret = try_emplace(Key, std::forward(Val)); + if (!Ret.second) + Ret.first->second = std::forward(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. Index: unittests/ADT/DenseMapTest.cpp =================================================================== --- unittests/ADT/DenseMapTest.cpp +++ unittests/ADT/DenseMapTest.cpp @@ -29,6 +29,11 @@ return &dummy_arr1[i]; } +struct NotDefaultConstructible { + NotDefaultConstructible() = delete; + int v; +}; + /// A test class that tries to check that construction and destruction /// occur correctly. class CtorTester { @@ -191,6 +196,24 @@ EXPECT_EQ(this->getValue(), this->Map[this->getKey()]); } +TEST(DenseMapCustomTest, InsertOrAssignTest) { + DenseMap Map; + auto Try1 = Map.insert_or_assign(0, NotDefaultConstructible{1}); + EXPECT_TRUE(Try1.second); + auto Try2 = Map.insert_or_assign(0, NotDefaultConstructible{2}); + EXPECT_FALSE(Try2.second); + EXPECT_EQ(2, Try2.first->second.v); + EXPECT_EQ(Try1.first, Try2.first); + + int Key = 1; + auto Try3 = Map.insert_or_assign(Key, NotDefaultConstructible{3}); + EXPECT_TRUE(Try3.second); + auto Try4 = Map.insert_or_assign(Key, NotDefaultConstructible{4}); + EXPECT_FALSE(Try4.second); + EXPECT_EQ(4, Try4.first->second.v); + EXPECT_EQ(Try3.first, Try4.first); +} + // Test copy constructor method TYPED_TEST(DenseMapTest, CopyConstructorTest) { this->Map[this->getKey()] = this->getValue();