Index: include/llvm/ADT/DenseMap.h =================================================================== --- include/llvm/ADT/DenseMap.h +++ include/llvm/ADT/DenseMap.h @@ -203,6 +203,22 @@ return try_emplace(std::move(KV.first), std::move(KV.second)); } + 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 @@ -191,6 +191,24 @@ EXPECT_EQ(this->getValue(), this->Map[this->getKey()]); } +TEST(DenseMapCustomTest, InsertOrAssignTest) { + DenseMap> Map; + auto Try1 = Map.insert_or_assign(0, std::make_unique(1)); + EXPECT_TRUE(Try1.second); + auto Try2 = Map.insert_or_assign(0, std::make_unique(2)); + EXPECT_FALSE(Try2.second); + EXPECT_EQ(2, *Try2.first->getSecond()); + EXPECT_EQ(Try1.first, Try2.first); + + int Key = 1; + auto Try3 = Map.insert_or_assign(Key, std::make_unique(3)); + EXPECT_TRUE(Try3.second); + auto Try4 = Map.insert_or_assign(Key, std::make_unique(4)); + EXPECT_FALSE(Try4.second); + EXPECT_EQ(4, *Try4.first->getSecond()); + EXPECT_EQ(Try3.first, Try4.first); +} + // Test copy constructor method TYPED_TEST(DenseMapTest, CopyConstructorTest) { this->Map[this->getKey()] = this->getValue();