[ADT] Add StringMap::insert_or_assign

Summary: Similar to std::unordered_map::insert_or_assign

Reviewers: alexshap, bkramer, dblaikie, lhames

Subscribers: jkorous, dexonsmith, kristina, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D67668

llvm-svn: 372813
This commit is contained in:
Fangrui Song 2019-09-25 04:58:02 +00:00
parent b2c6e9bf99
commit 0ab311bbff
2 changed files with 41 additions and 0 deletions

View File

@ -391,6 +391,16 @@ public:
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

View File

@ -75,6 +75,16 @@ 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();
@ -269,6 +279,27 @@ TEST_F(StringMapTest, InsertRehashingPairTest) {
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;