Index: llvm/trunk/include/llvm/ADT/StringMap.h =================================================================== --- llvm/trunk/include/llvm/ADT/StringMap.h +++ llvm/trunk/include/llvm/ADT/StringMap.h @@ -359,6 +359,11 @@ return find(Key) == end() ? 0 : 1; } + template + size_type count(const StringMapEntry &MapEntry) const { + return count(MapEntry.getKey()); + } + /// insert - Insert the specified key/value pair into the map. If the key /// already exists in the map, return false and ignore the request, otherwise /// insert it and return true. Index: llvm/trunk/include/llvm/ADT/StringSet.h =================================================================== --- llvm/trunk/include/llvm/ADT/StringSet.h +++ llvm/trunk/include/llvm/ADT/StringSet.h @@ -45,6 +45,12 @@ for (auto It = Begin; It != End; ++It) base::insert(std::make_pair(*It, '\0')); } + + template + std::pair + insert(const StringMapEntry &MapEntry) { + return insert(MapEntry.getKey()); + } }; } // end namespace llvm Index: llvm/trunk/unittests/ADT/CMakeLists.txt =================================================================== --- llvm/trunk/unittests/ADT/CMakeLists.txt +++ llvm/trunk/unittests/ADT/CMakeLists.txt @@ -65,6 +65,7 @@ StringExtrasTest.cpp StringMapTest.cpp StringRefTest.cpp + StringSetTest.cpp StringSwitchTest.cpp TinyPtrVectorTest.cpp TripleTest.cpp Index: llvm/trunk/unittests/ADT/StringMapTest.cpp =================================================================== --- llvm/trunk/unittests/ADT/StringMapTest.cpp +++ llvm/trunk/unittests/ADT/StringMapTest.cpp @@ -7,7 +7,6 @@ //===----------------------------------------------------------------------===// #include "llvm/ADT/StringMap.h" -#include "llvm/ADT/StringSet.h" #include "llvm/ADT/Twine.h" #include "llvm/Support/DataTypes.h" #include "gtest/gtest.h" @@ -284,20 +283,6 @@ EXPECT_EQ(Expected, Keys); } -TEST_F(StringMapTest, IterSetKeys) { - StringSet<> Set; - Set.insert("A"); - Set.insert("B"); - Set.insert("C"); - Set.insert("D"); - - auto Keys = to_vector<4>(Set.keys()); - llvm::sort(Keys); - - SmallVector Expected = {"A", "B", "C", "D"}; - EXPECT_EQ(Expected, Keys); -} - // Create a non-default constructable value struct StringMapTestStruct { StringMapTestStruct(int i) : i(i) {} Index: llvm/trunk/unittests/ADT/StringSetTest.cpp =================================================================== --- llvm/trunk/unittests/ADT/StringSetTest.cpp +++ llvm/trunk/unittests/ADT/StringSetTest.cpp @@ -0,0 +1,44 @@ +//===- llvm/unittest/ADT/StringSetTest.cpp - StringSet unit tests ----------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "llvm/ADT/StringSet.h" +#include "gtest/gtest.h" +using namespace llvm; + +namespace { + +// Test fixture +class StringSetTest : public testing::Test {}; + +TEST_F(StringSetTest, IterSetKeys) { + StringSet<> Set; + Set.insert("A"); + Set.insert("B"); + Set.insert("C"); + Set.insert("D"); + + auto Keys = to_vector<4>(Set.keys()); + llvm::sort(Keys); + + SmallVector Expected = {"A", "B", "C", "D"}; + EXPECT_EQ(Expected, Keys); +} + +TEST_F(StringSetTest, InsertAndCountStringMapEntry) { + // Test insert(StringMapEntry) and count(StringMapEntry) + // which are required for set_difference(StringSet, StringSet). + StringSet<> Set; + StringMapEntry *Element = StringMapEntry::Create("A"); + Set.insert(*Element); + size_t Count = Set.count(*Element); + size_t Expected = 1; + EXPECT_EQ(Expected, Count); + Element->Destroy(); +} + +} // end anonymous namespace