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,16 @@ return &dummy_arr1[i]; } +struct CountValueCopyAndMove { + CountValueCopyAndMove() = default; + CountValueCopyAndMove(const CountValueCopyAndMove &) { copy = 1; } + CountValueCopyAndMove(CountValueCopyAndMove &&) { move = 1; } + void operator=(const CountValueCopyAndMove &) { ++copy; } + void operator=(CountValueCopyAndMove &&) { ++move; } + int copy = 0; + int move = 0; +}; + /// A test class that tries to check that construction and destruction /// occur correctly. class CtorTester { @@ -191,6 +201,35 @@ EXPECT_EQ(this->getValue(), this->Map[this->getKey()]); } +TEST(DenseMapCustomTest, InsertOrAssignTest) { + struct A : CountValueCopyAndMove { + A(int v) : v(v) {} + int v; + }; + DenseMap map; + int key = 1; + + auto try1 = map.insert_or_assign(0, A(1)); + EXPECT_TRUE(try1.second); + EXPECT_EQ(1, try1.first->second.move); + auto try2 = map.insert_or_assign(0, A(2)); + EXPECT_FALSE(try2.second); + EXPECT_EQ(2, try2.first->second.v); + EXPECT_EQ(2, try2.first->second.move); + EXPECT_EQ(try1.first, try2.first); + EXPECT_EQ(0, try1.first->second.copy); + + auto try3 = map.insert_or_assign(key, A(3)); + EXPECT_TRUE(try3.second); + EXPECT_EQ(1, try3.first->second.move); + auto try4 = map.insert_or_assign(key, A(4)); + EXPECT_FALSE(try4.second); + EXPECT_EQ(4, try4.first->second.v); + EXPECT_EQ(2, try4.first->second.move); + EXPECT_EQ(try3.first, try4.first); + EXPECT_EQ(0, try3.first->second.copy); +} + // Test copy constructor method TYPED_TEST(DenseMapTest, CopyConstructorTest) { this->Map[this->getKey()] = this->getValue();