diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h index 3ef6a7cd1b4b5..7ccc9445c0a7b 100644 --- a/llvm/include/llvm/ADT/DenseMap.h +++ b/llvm/include/llvm/ADT/DenseMap.h @@ -312,6 +312,22 @@ class DenseMapBase : public DebugEpochBase { insert(*I); } + 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; + } + + 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; + } + /// Returns the value associated to the key in the map if it exists. If it /// does not exist, emplace a default value for the key and returns a /// reference to the newly created value. diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp index 09f9a57647709..cb45a6e681ce4 100644 --- a/llvm/unittests/ADT/DenseMapTest.cpp +++ b/llvm/unittests/ADT/DenseMapTest.cpp @@ -475,6 +475,42 @@ TEST(DenseMapCustomTest, ReserveTest) { } } +TEST(DenseMapCustomTest, InsertOrAssignTest) { + DenseMap Map; + + CountCopyAndMove val1(1); + CountCopyAndMove::ResetCounts(); + auto try0 = Map.insert_or_assign(0, val1); + EXPECT_TRUE(try0.second); + EXPECT_EQ(0, CountCopyAndMove::TotalMoves()); + EXPECT_EQ(1, CountCopyAndMove::CopyConstructions); + EXPECT_EQ(0, CountCopyAndMove::CopyAssignments); + + CountCopyAndMove::ResetCounts(); + auto try1 = Map.insert_or_assign(0, val1); + EXPECT_FALSE(try1.second); + EXPECT_EQ(0, CountCopyAndMove::TotalMoves()); + EXPECT_EQ(0, CountCopyAndMove::CopyConstructions); + EXPECT_EQ(1, CountCopyAndMove::CopyAssignments); + + int key2 = 2; + CountCopyAndMove val2(2); + CountCopyAndMove::ResetCounts(); + auto try2 = Map.insert_or_assign(key2, std::move(val2)); + EXPECT_TRUE(try2.second); + EXPECT_EQ(0, CountCopyAndMove::TotalCopies()); + EXPECT_EQ(1, CountCopyAndMove::MoveConstructions); + EXPECT_EQ(0, CountCopyAndMove::MoveAssignments); + + CountCopyAndMove val3(3); + CountCopyAndMove::ResetCounts(); + auto try3 = Map.insert_or_assign(key2, std::move(val3)); + EXPECT_FALSE(try3.second); + EXPECT_EQ(0, CountCopyAndMove::TotalCopies()); + EXPECT_EQ(0, CountCopyAndMove::MoveConstructions); + EXPECT_EQ(1, CountCopyAndMove::MoveAssignments); +} + // Make sure DenseMap works with StringRef keys. TEST(DenseMapCustomTest, StringRefTest) { DenseMap M;