MapVector: add C++17-style try_emplace and insert_or_assign

[MaskRay]

Similar to https://wg21.link/n4279

For example, insert_or_assign can be used to simplify
CodeGenModule::AddDeferredUnusedCoverageMapping in
clang/lib/CodeGen/CodeGenModule.cpp

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Fangrui Song (MaskRay)

Changes

They can avoid a copy/move for many insert/operator[] uses.

For example, insert_or_assign can be used to simplify
CodeGenModule::AddDeferredUnusedCoverageMapping in
clang/lib/CodeGen/CodeGenModule.cpp

---

Full diff: https://github.com/llvm/llvm-project/pull/71969.diff

2 Files Affected:

• (modified) llvm/include/llvm/ADT/MapVector.h (+35-10)
• (modified) llvm/unittests/ADT/MapVectorTest.cpp (+107)

diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h
index c45779c0ce8e0cf..56d02d73bb0fc30 100644
--- a/llvm/include/llvm/ADT/MapVector.h
+++ b/llvm/include/llvm/ADT/MapVector.h
@@ -114,31 +114,56 @@ class MapVector {
     return Pos == Map.end()? ValueT() : Vector[Pos->second].second;
   }
 
-  std::pair<iterator, bool> insert(const std::pair<KeyT, ValueT> &KV) {
-    std::pair<KeyT, typename MapType::mapped_type> Pair = std::make_pair(KV.first, 0);
-    std::pair<typename MapType::iterator, bool> Result = Map.insert(Pair);
+  template <typename... Ts>
+  std::pair<iterator, bool> try_emplace(const KeyT &Key, Ts &&...Args) {
+    std::pair<typename MapType::iterator, bool> Result =
+        Map.insert(std::make_pair(Key, 0));
     auto &I = Result.first->second;
     if (Result.second) {
-      Vector.push_back(std::make_pair(KV.first, KV.second));
+      Vector.emplace_back(std::piecewise_construct, std::forward_as_tuple(Key),
+                          std::forward_as_tuple(std::forward<Ts>(Args)...));
       I = Vector.size() - 1;
       return std::make_pair(std::prev(end()), true);
     }
     return std::make_pair(begin() + I, false);
   }
-
-  std::pair<iterator, bool> insert(std::pair<KeyT, ValueT> &&KV) {
-    // Copy KV.first into the map, then move it into the vector.
-    std::pair<KeyT, typename MapType::mapped_type> Pair = std::make_pair(KV.first, 0);
-    std::pair<typename MapType::iterator, bool> Result = Map.insert(Pair);
+  template <typename... Ts>
+  std::pair<iterator, bool> try_emplace(KeyT &&Key, Ts &&...Args) {
+    std::pair<typename MapType::iterator, bool> Result =
+        Map.insert(std::make_pair(Key, 0));
     auto &I = Result.first->second;
     if (Result.second) {
-      Vector.push_back(std::move(KV));
+      Vector.emplace_back(std::piecewise_construct,
+                          std::forward_as_tuple(static_cast<KeyT &&>(Key)),
+                          std::forward_as_tuple(std::forward<Ts>(Args)...));
       I = Vector.size() - 1;
       return std::make_pair(std::prev(end()), true);
     }
     return std::make_pair(begin() + I, false);
   }
 
+  std::pair<iterator, bool> insert(const std::pair<KeyT, ValueT> &KV) {
+    return try_emplace(KV.first, KV.second);
+  }
+  std::pair<iterator, bool> insert(std::pair<KeyT, ValueT> &&KV) {
+    return try_emplace(std::move(KV.first), std::move(KV.second));
+  }
+
+  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;
+  }
+
   bool contains(const KeyT &Key) const { return Map.find(Key) != Map.end(); }
 
   size_type count(const KeyT &Key) const { return contains(Key) ? 1 : 0; }
diff --git a/llvm/unittests/ADT/MapVectorTest.cpp b/llvm/unittests/ADT/MapVectorTest.cpp
index 1a371cbfba81e8e..40645b9a4de07b7 100644
--- a/llvm/unittests/ADT/MapVectorTest.cpp
+++ b/llvm/unittests/ADT/MapVectorTest.cpp
@@ -9,10 +9,36 @@
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/iterator_range.h"
 #include "gtest/gtest.h"
+#include <memory>
 #include <utility>
 
 using namespace llvm;
 
+namespace {
+struct CountCopyAndMove {
+  CountCopyAndMove() = default;
+  CountCopyAndMove(const CountCopyAndMove &) { copy = 1; }
+  CountCopyAndMove(CountCopyAndMove &&) { move = 1; }
+  void operator=(const CountCopyAndMove &) { ++copy; }
+  void operator=(CountCopyAndMove &&) { ++move; }
+  int copy = 0;
+  int move = 0;
+};
+
+struct A : CountCopyAndMove {
+  A(int v) : v(v) {}
+  int v;
+};
+} // namespace
+
+template <> struct DenseMapInfo<A> {
+  static inline A getEmptyKey() { return 0x7fffffff; }
+  static inline A getTombstoneKey() { return -0x7fffffff - 1; }
+  static unsigned getHashValue(const A &Val) { return (unsigned)(Val.v * 37U); }
+  static bool isEqual(const A &LHS, const A &RHS) { return LHS.v == RHS.v; }
+};
+
+namespace {
 TEST(MapVectorTest, swap) {
   MapVector<int, int> MV1, MV2;
   std::pair<MapVector<int, int>::iterator, bool> R;
@@ -79,6 +105,86 @@ TEST(MapVectorTest, insert_pop) {
   EXPECT_EQ(MV[4], 7);
 }
 
+TEST(MapVectorTest, try_emplace) {
+  struct AAndU {
+    A a;
+    std::unique_ptr<int> b;
+    AAndU(A a, std::unique_ptr<int> b) : a(a), b(std::move(b)) {}
+  };
+  MapVector<A, AAndU> mv;
+
+  A zero(0);
+  auto try0 = mv.try_emplace(zero, zero, nullptr);
+  EXPECT_TRUE(try0.second);
+  EXPECT_EQ(0, try0.first->second.a.v);
+  EXPECT_EQ(1, try0.first->second.a.copy);
+  EXPECT_EQ(0, try0.first->second.a.move);
+
+  auto try1 = mv.try_emplace(zero, zero, nullptr);
+  EXPECT_FALSE(try1.second);
+  EXPECT_EQ(0, try1.first->second.a.v);
+  EXPECT_EQ(1, try1.first->second.a.copy);
+  EXPECT_EQ(0, try1.first->second.a.move);
+
+  EXPECT_EQ(try0.first, try1.first);
+  EXPECT_EQ(1, try1.first->first.copy);
+  EXPECT_EQ(0, try1.first->first.move);
+
+  A two(2);
+  auto try2 = mv.try_emplace(2, (A &&)two, std::make_unique<int>(2));
+  EXPECT_TRUE(try2.second);
+  EXPECT_EQ(2, try2.first->second.a.v);
+  EXPECT_EQ(0, try2.first->second.a.move);
+
+  std::unique_ptr<int> p(new int(3));
+  auto try3 = mv.try_emplace((A &&)two, 3, std::move(p));
+  EXPECT_FALSE(try3.second);
+  EXPECT_EQ(2, try3.first->second.a.v);
+  EXPECT_EQ(1, try3.first->second.a.copy);
+  EXPECT_EQ(0, try3.first->second.a.move);
+
+  EXPECT_EQ(try2.first, try3.first);
+  EXPECT_EQ(0, try3.first->first.copy);
+  EXPECT_EQ(1, try3.first->first.move);
+  EXPECT_NE(nullptr, p);
+}
+
+TEST(MapVectorTest, insert_or_assign) {
+  MapVector<A, A> mv;
+
+  A zero(0);
+  auto try0 = mv.insert_or_assign(zero, zero);
+  EXPECT_TRUE(try0.second);
+  EXPECT_EQ(0, try0.first->second.v);
+  EXPECT_EQ(1, try0.first->second.copy);
+  EXPECT_EQ(0, try0.first->second.move);
+
+  auto try1 = mv.insert_or_assign(zero, zero);
+  EXPECT_FALSE(try1.second);
+  EXPECT_EQ(0, try1.first->second.v);
+  EXPECT_EQ(2, try1.first->second.copy);
+  EXPECT_EQ(0, try1.first->second.move);
+
+  EXPECT_EQ(try0.first, try1.first);
+  EXPECT_EQ(1, try1.first->first.copy);
+  EXPECT_EQ(0, try1.first->first.move);
+
+  A two(2);
+  auto try2 = mv.try_emplace(2, (A &&)two);
+  EXPECT_TRUE(try2.second);
+  EXPECT_EQ(2, try2.first->second.v);
+  EXPECT_EQ(1, try2.first->second.move);
+
+  auto try3 = mv.insert_or_assign((A &&)two, 3);
+  EXPECT_FALSE(try3.second);
+  EXPECT_EQ(3, try3.first->second.v);
+  EXPECT_EQ(0, try3.first->second.copy);
+  EXPECT_EQ(2, try3.first->second.move);
+
+  EXPECT_EQ(try2.first, try3.first);
+  EXPECT_EQ(1, try3.first->first.move);
+}
+
 TEST(MapVectorTest, erase) {
   MapVector<int, int> MV;
 
@@ -423,3 +529,4 @@ TEST(SmallMapVectorLargeTest, iteration_test) {
     count--;
   }
 }
+} // namespace

[dwblaikie]

Could this static_cast be std::forward maybe?

[lichray]

It could be std::move, but it's okay to aim for less debug codegen using a cast.

[dwblaikie]

std::move is a builtin in clang - so there's no extra codegen/debug info for it & not sure why this one place would be a raw static_cast but everywhere else is std::move if that's the tradeoff we're making? This one doesn't seem different & so probably shouldn't be written differently?

[MaskRay]

Changed to std::move. Thanks!

https://reviews.llvm.org/D123345 made std::move builtin. Even without the feature, libc++ annotates its move impl with __attribute__((__no_debug__)).

[lichray]

It could be std::move, but it's okay to aim for less debug codegen using a cast.

[dwblaikie]

Structure binding might be more legible & less verbose here (& similarly elsewhere)