From 8ab112dbf48c3b2dbf8ae61b50732dd5982ee683 Mon Sep 17 00:00:00 2001
From: David Blaikie <dblaikie@gmail.com>
Date: Fri, 17 Jul 2020 10:43:12 -0700
Subject: [PATCH] [llvm] Add contains(KeyType) -> bool methods to StringSet

Matches C++20 API addition.

Differential Revision: https://reviews.llvm.org/D83449
---
 include/llvm/ADT/StringSet.h    |  3 +++
 unittests/ADT/StringSetTest.cpp | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/include/llvm/ADT/StringSet.h b/include/llvm/ADT/StringSet.h
index 63d929399a4..c4245175544 100644
--- a/include/llvm/ADT/StringSet.h
+++ b/include/llvm/ADT/StringSet.h
@@ -45,6 +45,9 @@ public:
   insert(const StringMapEntry<ValueTy> &mapEntry) {
     return insert(mapEntry.getKey());
   }
+
+  /// Check if the set contains the given \c key.
+  bool contains(StringRef key) const { return Base::FindKey(key) != -1; }
 };
 
 } // end namespace llvm
diff --git a/unittests/ADT/StringSetTest.cpp b/unittests/ADT/StringSetTest.cpp
index e27306dd696..0071f5a3775 100644
--- a/unittests/ADT/StringSetTest.cpp
+++ b/unittests/ADT/StringSetTest.cpp
@@ -53,4 +53,23 @@ TEST_F(StringSetTest, EmptyString) {
   EXPECT_EQ(Count, 1UL);
 }
 
+TEST_F(StringSetTest, Contains) {
+  StringSet<> Set;
+  EXPECT_FALSE(Set.contains(""));
+  EXPECT_FALSE(Set.contains("test"));
+
+  Set.insert("");
+  Set.insert("test");
+  EXPECT_TRUE(Set.contains(""));
+  EXPECT_TRUE(Set.contains("test"));
+
+  Set.insert("test");
+  EXPECT_TRUE(Set.contains(""));
+  EXPECT_TRUE(Set.contains("test"));
+
+  Set.erase("test");
+  EXPECT_TRUE(Set.contains(""));
+  EXPECT_FALSE(Set.contains("test"));
+}
+
 } // end anonymous namespace