[llvm] 07b8990 - [ADT] Add C++17-style insert_or_assign for DenseMap (#94151)

[via llvm-commits]

Author: c8ef
Date: 2024-06-05T07:37:40-07:00
New Revision: 07b8990d38ebaf42b0ed1c60d49b484a6e2497e8

URL: https://github.com/llvm/llvm-project/commit/07b8990d38ebaf42b0ed1c60d49b484a6e2497e8
DIFF: https://github.com/llvm/llvm-project/commit/07b8990d38ebaf42b0ed1c60d49b484a6e2497e8.diff

LOG: [ADT] Add C++17-style insert_or_assign for DenseMap (#94151)

add C++17-style insert_or_assign for DenseMap

close: #94115

Added: 
    

Modified: 
    llvm/include/llvm/ADT/DenseMap.h
    llvm/unittests/ADT/DenseMapTest.cpp

Removed: 
    


################################################################################
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 <typename V>
+  std::pair<iterator, bool> insert_or_assign(const KeyT &Key, V &&Val) {
+    auto Ret = try_emplace(Key, std::forward<V>(Val));
+    if (!Ret.second)
+      Ret.first->second = std::forward<V>(Val);
+    return Ret;
+  }
+
+  template <typename V>
+  std::pair<iterator, bool> insert_or_assign(KeyT &&Key, V &&Val) {
+    auto Ret = try_emplace(std::move(Key), std::forward<V>(Val));
+    if (!Ret.second)
+      Ret.first->second = std::forward<V>(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<int, CountCopyAndMove> 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<StringRef, int> M;