[DenseMap] Add a C++17-style try_emplace method.

This provides an elegant pattern to solve the "construct if not in map
already" problem we have many times in LLVM. Without try_emplace we
either have to rely on a sentinel value (nullptr) or do two lookups.

llvm-svn: 276277

2 files changed, 45 insertions, 37 deletions

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index 917c086..aaaf508e 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -169,30 +169,45 @@ public:
   // If the key is already in the map, it returns false and doesn't update the
   // value.
   std::pair<iterator, bool> insert(const std::pair<KeyT, ValueT> &KV) {
+    return try_emplace(KV.first, KV.second);
+  }
+
+  // Inserts key,value pair into the map if the key isn't already in the map.
+  // If the key is already in the map, it returns false and doesn't update the
+  // value.
+  std::pair<iterator, bool> insert(std::pair<KeyT, ValueT> &&KV) {
+    return try_emplace(std::move(KV.first), std::move(KV.second));
+  }
+
+  // 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.
+  template <typename... Ts>
+  std::pair<iterator, bool> try_emplace(KeyT &&Key, Ts &&... Args) {
     BucketT *TheBucket;
-    if (LookupBucketFor(KV.first, TheBucket))
+    if (LookupBucketFor(Key, TheBucket))
       return std::make_pair(iterator(TheBucket, getBucketsEnd(), *this, true),
                             false); // Already in map.
 
     // Otherwise, insert the new element.
-    TheBucket = InsertIntoBucket(KV.first, KV.second, TheBucket);
+    TheBucket =
+        InsertIntoBucket(TheBucket, std::move(Key), std::forward<Ts>(Args)...);
     return std::make_pair(iterator(TheBucket, getBucketsEnd(), *this, true),
                           true);
   }
 
   // Inserts key,value pair into the map if the key isn't already in the map.
-  // If the key is already in the map, it returns false and doesn't update the
-  // value.
-  std::pair<iterator, bool> insert(std::pair<KeyT, ValueT> &&KV) {
+  // The value is constructed in-place if the key is not in the map, otherwise
+  // it is not moved.
+  template <typename... Ts>
+  std::pair<iterator, bool> try_emplace(const KeyT &Key, Ts &&... Args) {
     BucketT *TheBucket;
-    if (LookupBucketFor(KV.first, TheBucket))
+    if (LookupBucketFor(Key, TheBucket))
       return std::make_pair(iterator(TheBucket, getBucketsEnd(), *this, true),
                             false); // Already in map.
 
     // Otherwise, insert the new element.
-    TheBucket = InsertIntoBucket(std::move(KV.first),
-                                 std::move(KV.second),
-                                 TheBucket);
+    TheBucket = InsertIntoBucket(TheBucket, Key, std::forward<Ts>(Args)...);
     return std::make_pair(iterator(TheBucket, getBucketsEnd(), *this, true),
                           true);
   }
@@ -211,8 +226,8 @@ public:
                             false); // Already in map.
 
     // Otherwise, insert the new element.
-    TheBucket = InsertIntoBucket(std::move(KV.first), std::move(KV.second), Val,
-                                 TheBucket);
+    TheBucket = InsertIntoBucketWithLookup(TheBucket, std::move(KV.first),
+                                           std::move(KV.second), Val);
     return std::make_pair(iterator(TheBucket, getBucketsEnd(), *this, true),
                           true);
   }
@@ -249,7 +264,7 @@ public:
     if (LookupBucketFor(Key, TheBucket))
       return *TheBucket;
 
-    return *InsertIntoBucket(Key, ValueT(), TheBucket);
+    return *InsertIntoBucket(TheBucket, Key);
   }
 
   ValueT &operator[](const KeyT &Key) {
@@ -261,7 +276,7 @@ public:
     if (LookupBucketFor(Key, TheBucket))
       return *TheBucket;
 
-    return *InsertIntoBucket(std::move(Key), ValueT(), TheBucket);
+    return *InsertIntoBucket(TheBucket, std::move(Key));
   }
 
   ValueT &operator[](KeyT &&Key) {
@@ -429,36 +444,19 @@ private:
     static_cast<DerivedT *>(this)->shrink_and_clear();
   }
 
-
-  BucketT *InsertIntoBucket(const KeyT &Key, const ValueT &Value,
-                            BucketT *TheBucket) {
-    TheBucket = InsertIntoBucketImpl(Key, Key, TheBucket);
-
-    TheBucket->getFirst() = Key;
-    ::new (&TheBucket->getSecond()) ValueT(Value);
-    return TheBucket;
-  }
-
-  BucketT *InsertIntoBucket(const KeyT &Key, ValueT &&Value,
-                            BucketT *TheBucket) {
-    TheBucket = InsertIntoBucketImpl(Key, Key, TheBucket);
-
-    TheBucket->getFirst() = Key;
-    ::new (&TheBucket->getSecond()) ValueT(std::move(Value));
-    return TheBucket;
-  }
-
-  BucketT *InsertIntoBucket(KeyT &&Key, ValueT &&Value, BucketT *TheBucket) {
+  template <typename KeyArg, typename... ValueArgs>
+  BucketT *InsertIntoBucket(BucketT *TheBucket, KeyArg &&Key,
+                            ValueArgs &&... Values) {
     TheBucket = InsertIntoBucketImpl(Key, Key, TheBucket);
 
-    TheBucket->getFirst() = std::move(Key);
-    ::new (&TheBucket->getSecond()) ValueT(std::move(Value));
+    TheBucket->getFirst() = std::forward<KeyArg>(Key);
+    ::new (&TheBucket->getSecond()) ValueT(std::forward<ValueArgs>(Values)...);
     return TheBucket;
   }
 
   template <typename LookupKeyT>
-  BucketT *InsertIntoBucket(KeyT &&Key, ValueT &&Value, LookupKeyT &Lookup,
-                            BucketT *TheBucket) {
+  BucketT *InsertIntoBucketWithLookup(BucketT *TheBucket, KeyT &&Key,
+                                      ValueT &&Value, LookupKeyT &Lookup) {
     TheBucket = InsertIntoBucketImpl(Key, Lookup, TheBucket);
 
     TheBucket->getFirst() = std::move(Key);
diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp
index db00f8c..cbf1a44 100644
--- a/llvm/unittests/ADT/DenseMapTest.cpp
+++ b/llvm/unittests/ADT/DenseMapTest.cpp
@@ -619,4 +619,14 @@ TEST(DenseMapCustomTest, SmallDenseMapGrowTest) {
   EXPECT_TRUE(map.find(32) == map.end());
 }
 
+TEST(DenseMapCustomTest, TryEmplaceTest) {
+  DenseMap<int, std::unique_ptr<int>> Map;
+  std::unique_ptr<int> P(new int(2));
+  auto Try1 = Map.try_emplace(0, new int(1));
+  EXPECT_TRUE(Try1.second);
+  auto Try2 = Map.try_emplace(0, std::move(P));
+  EXPECT_FALSE(Try2.second);
+  EXPECT_EQ(Try1.first, Try2.first);
+  EXPECT_NE(nullptr, P);
+}
 }