[ADT] Add StringMap::insert_or_assign

[MaskRay]

Summary

Similar to std::unordered_map::insert_or_assign

[MaskRay]

It turns out that I need StringMap::insert_or_assign, instead of DenseMap, for D67656.

[jkorous]

include/llvm/ADT/StringMap.h

394

Please add some description.

[MaskRay]

simplify

[dblaikie]

Could we have this as a standalone helper function (STLExtras or the like) rather than implementing it in every container?

[MaskRay]

In D67668#1674393, @dblaikie wrote:

Could we have this as a standalone helper function (STLExtras or the like) rather than implementing it in every container?

We only need it in 2 unordered associative containers: DenseMap, and StringMap. Member methods have the advantage that they have better discoverability and can be easily be completed. This method is simple enough now. Placing it in STLExtras does not save much.

[MaskRay]

Add a comment

[dblaikie]

In D67668#1674807, @MaskRay wrote:

In D67668#1674393, @dblaikie wrote:

Could we have this as a standalone helper function (STLExtras or the like) rather than implementing it in every container?

We only need it in 2 unordered associative containers: DenseMap, and StringMap. Member methods have the advantage that they have better discoverability and can be easily be completed. This method is simple enough now. Placing it in STLExtras does not save much.

There's also MapVector, I guess, but fair enough - for consistency with the standard containers it's not too much.

unittests/ADT/StringMapTest.cpp

283–290

You could use a noisy type here to test that only the appropriate operations were executed (and/or that a type that's not default-constructible is usable here, whereas it wouldn't be if you replaced the implementation with 't[x] = y' for instance)

Countable, further down in this test file, could be used to help there.

[MaskRay]

You could use a noisy type here to test that only the appropriate operations were executed

Add struct NoDefaultConstructible {...}

[dblaikie]

include/llvm/ADT/StringMap.h

397

Is this API consistent with the C++ standard version? I would guess the c++ standard version takes a pair<K, V>? I'd sort of come back to "why not a free function" if this isn't going to match the standard API & make API interoperability completely transparent.

unittests/ADT/StringMapTest.cpp

283–290

I was thinking/hoping one of the existing test types could be used and might provide more robust checking? One that demonstrates the first insert_or_assign only calls the move constructor and the second only calls the move assignment operator?

[MaskRay]

include/llvm/ADT/StringMap.h

397

The standard version provides two overloads:

template <class M> pair<iterator, bool> insert_or_assign(const key_type& k, M&& obj);
template <class M> pair<iterator, bool> insert_or_assign(key_type&& k, M&& obj);

This is unnecessary for StringRef, which is cheap to copy (2 words). So I simply add a method that takes a StringRef instead of creating 2 overloads.

unittests/ADT/StringMapTest.cpp

283–290

There is no existing trait-based test types. I can add move constructor/assignment operator to ensure no copy assignment is called if you are strong about this.

[dblaikie]

include/llvm/ADT/StringMap.h

397

Ah, cool - thanks!

unittests/ADT/StringMapTest.cpp

283–290

I was thinking something like CountCtorCopyAndMove (we might even have a generalized one of these somewhere as a test utility for all of libSupport? Or certainly they've been repeated in a few tests) that distinguishes (move/copy) assignments from (move/copy) constructions.

[dblaikie]

Seems good with the test fleshed out more before committing - thanks for your patience!

[MaskRay]

Add CountCopyAndMove to flesh out the test

include/llvm/ADT/StringMap.h

public:
   /// insert - Inserts the specified key/value pair into the map if the key
   /// isn't already in the map. The bool component of the returned pair is true
   /// if and only if the insertion takes place, and the iterator component of
   /// the pair points to the element with key equivalent to the key of the pair.
   std::pair<iterator, bool> insert(std::pair<StringRef, ValueTy> KV) {
     return try_emplace(KV.first, std::move(KV.second));
   }
 
+  /// Inserts an element or assigns to the current element if the key already
+  /// exists. The return type is the same as try_emplace.
+  template <typename V>
+  std::pair<iterator, bool> insert_or_assign(StringRef Key, V &&Val) {
+    auto Ret = try_emplace(Key, std::forward<V>(Val));
+    if (!Ret.second)
+      Ret.first->second = std::forward<V>(Val);
+    return Ret;
+  }
+
   /// Emplace a new element for the specified key into the map if the key isn't
   /// already in the map. The bool component of the returned pair is true
   /// if and only if the insertion takes place, and the iterator component of
   /// the pair points to the element with key equivalent to the key of the pair.
   template <typename... ArgsTy>
   std::pair<iterator, bool> try_emplace(StringRef Key, ArgsTy &&... Args) {
     unsigned BucketNo = LookupBucketFor(Key);
     StringMapEntryBase *&Bucket = TheTable[BucketNo];

unittests/ADT/StringMapTest.cpp

 };
 
 const char StringMapTest::testKey[] = "key";
 const uint32_t StringMapTest::testValue = 1u;
 const char* StringMapTest::testKeyFirst = testKey;
 size_t StringMapTest::testKeyLength = sizeof(testKey) - 1;
 const std::string StringMapTest::testKeyStr(testKey);
 
+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;
+};
+
 // Empty map tests.
 TEST_F(StringMapTest, EmptyMapTest) {
   assertEmptyMap();
 }
 
 // Constant map tests.
 TEST_F(StringMapTest, ConstEmptyMapTest) {
   const StringMap<uint32_t>& constTestMap = testMap;
TEST_F(StringMapTest, InsertRehashingPairTest) {
 
   StringMap<uint32_t>::iterator It =
     t.insert(std::make_pair("abcdef", 42)).first;
   EXPECT_EQ(16u, t.getNumBuckets());
   EXPECT_EQ("abcdef", It->first());
   EXPECT_EQ(42u, It->second);
 }
 
+TEST_F(StringMapTest, InsertOrAssignTest) {
+  struct A : CountCopyAndMove {
+    A(int v) : v(v) {}
+    int v;
+  };
+  StringMap<A> t(0);
+
+  auto try1 = t.insert_or_assign("A", A(1));
+  EXPECT_TRUE(try1.second);
+  EXPECT_EQ(1, try1.first->second.v);
+  EXPECT_EQ(1, try1.first->second.move);
+
+  auto try2 = t.insert_or_assign("A", A(2));
+  EXPECT_FALSE(try2.second);
+  EXPECT_EQ(2, try2.first->second.v);
+  EXPECT_EQ(2, try1.first->second.move);
+
+  EXPECT_EQ(try1.first, try2.first);
+  EXPECT_EQ(0, try1.first->second.copy);
+}
+
 TEST_F(StringMapTest, IterMapKeys) {
   StringMap<int> Map;
   Map["A"] = 1;
   Map["B"] = 2;
   Map["C"] = 3;
   Map["D"] = 3;
 
   auto Keys = to_vector<4>(Map.keys());