2 files changed, 72 insertions, 0 deletions

diff --git a/llvm/include/llvm/ADT/DenseSet.h b/llvm/include/llvm/ADT/DenseSet.h
index 3724a09623f..d13f269f530 100644
--- a/llvm/include/llvm/ADT/DenseSet.h
+++ b/llvm/include/llvm/ADT/DenseSet.h
@@ -58,6 +58,10 @@ public:
   /// the Size of the set.
   void resize(size_t Size) { TheMap.resize(Size); }
 
+  /// Grow the DenseSet so that it can contain at least \p NumEntries items
+  /// before resizing again.
+  void reserve(size_t Size) { TheMap.reserve(Size); }
+
   void clear() {
     TheMap.clear();
   }
@@ -154,6 +158,11 @@ public:
     return TheMap.insert(std::make_pair(V, Empty));
   }
 
+  std::pair<iterator, bool> insert(ValueT &&V) {
+    detail::DenseSetEmpty Empty;
+    return TheMap.insert(std::make_pair(std::move(V), Empty));
+  }
+
   /// Alternative version of insert that uses a different (and possibly less
   /// expensive) key type.
   template <typename LookupKeyT>
diff --git a/llvm/unittests/ADT/DenseSetTest.cpp b/llvm/unittests/ADT/DenseSetTest.cpp
index 5952353034f..f85025fd5ac 100644
--- a/llvm/unittests/ADT/DenseSetTest.cpp
+++ b/llvm/unittests/ADT/DenseSetTest.cpp
@@ -65,4 +65,67 @@ TEST(DenseSetCustomTest, FindAsTest) {
   EXPECT_TRUE(set.find_as("d") == set.end());
 }
 
+// Simple class that counts how many moves and copy happens when growing a map
+struct CountCopyAndMove {
+  static int Move;
+  static int Copy;
+  int Value;
+  CountCopyAndMove(int Value) : Value(Value) {}
+
+  CountCopyAndMove(const CountCopyAndMove &) { Copy++; }
+  CountCopyAndMove &operator=(const CountCopyAndMove &) {
+    Copy++;
+    return *this;
+  }
+  CountCopyAndMove(CountCopyAndMove &&) { Move++; }
+  CountCopyAndMove &operator=(const CountCopyAndMove &&) {
+    Move++;
+    return *this;
+  }
+};
+int CountCopyAndMove::Copy = 0;
+int CountCopyAndMove::Move = 0;
+} // anonymous namespace
+
+namespace llvm {
+// Specialization required to insert a CountCopyAndMove into a DenseSet.
+template <> struct DenseMapInfo<CountCopyAndMove> {
+  static inline CountCopyAndMove getEmptyKey() { return CountCopyAndMove(-1); };
+  static inline CountCopyAndMove getTombstoneKey() {
+    return CountCopyAndMove(-2);
+  };
+  static unsigned getHashValue(const CountCopyAndMove &Val) {
+    return Val.Value;
+  }
+  static bool isEqual(const CountCopyAndMove &LHS,
+                      const CountCopyAndMove &RHS) {
+    return LHS.Value == RHS.Value;
+  }
+};
+}
+
+namespace {
+// Make sure reserve actually gives us enough buckets to insert N items
+// without increasing allocation size.
+TEST(DenseSetCustomTest, ReserveTest) {
+  // Test a few different size, 48 is *not* a random choice: we need a value
+  // that is 2/3 of a power of two to stress the grow() condition, and the power
+  // of two has to be at least 64 because of minimum size allocation in the
+  // DenseMa. 66 is a value just above the 64 default init.
+  for (auto Size : {1, 2, 48, 66}) {
+    DenseSet<CountCopyAndMove> Set;
+    Set.reserve(Size);
+    unsigned MemorySize = Set.getMemorySize();
+    CountCopyAndMove::Copy = 0;
+    CountCopyAndMove::Move = 0;
+    for (int i = 0; i < Size; ++i)
+      Set.insert(CountCopyAndMove(i));
+    // Check that we didn't grow
+    EXPECT_EQ(MemorySize, Set.getMemorySize());
+    // Check that move was called the expected number of times
+    EXPECT_EQ(Size, CountCopyAndMove::Move);
+    // Check that no copy occured
+    EXPECT_EQ(0, CountCopyAndMove::Copy);
+  }
+}
 }