[llvm] dd4426b - [llvm] Add contains(KeyType) -> bool methods to SmallSet

[David Blaikie via llvm-commits]

Author: David Blaikie
Date: 2020-07-17T11:26:27-07:00
New Revision: dd4426b9a66e0363184adcfc3967e6b9f36f95dc

URL: https://github.com/llvm/llvm-project/commit/dd4426b9a66e0363184adcfc3967e6b9f36f95dc
DIFF: https://github.com/llvm/llvm-project/commit/dd4426b9a66e0363184adcfc3967e6b9f36f95dc.diff

LOG: [llvm] Add contains(KeyType) -> bool methods to SmallSet

Matches C++20 API addition.

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

Added: 
    

Modified: 
    llvm/include/llvm/ADT/SmallSet.h
    llvm/unittests/ADT/SmallSetTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/SmallSet.h b/llvm/include/llvm/ADT/SmallSet.h
index a03fa7dd8423..0600e528ee69 100644
--- a/llvm/include/llvm/ADT/SmallSet.h
+++ b/llvm/include/llvm/ADT/SmallSet.h
@@ -232,6 +232,13 @@ class SmallSet {
     return {Set.end()};
   }
 
+  /// Check if the SmallSet contains the given element.
+  bool contains(const T &V) const {
+    if (isSmall())
+      return vfind(V) != Vector.end();
+    return Set.find(V) != Set.end();
+  }
+
 private:
   bool isSmall() const { return Set.empty(); }
 

diff  --git a/llvm/unittests/ADT/SmallSetTest.cpp b/llvm/unittests/ADT/SmallSetTest.cpp
index 06682ce823dc..26cab828c784 100644
--- a/llvm/unittests/ADT/SmallSetTest.cpp
+++ b/llvm/unittests/ADT/SmallSetTest.cpp
@@ -167,3 +167,28 @@ TEST(SmallSetTest, EqualityComparisonTest) {
   EXPECT_NE(s1small, s4large);
   EXPECT_NE(s4large, s3large);
 }
+
+TEST(SmallSetTest, Contains) {
+  SmallSet<int, 2> Set;
+  EXPECT_FALSE(Set.contains(0));
+  EXPECT_FALSE(Set.contains(1));
+
+  Set.insert(0);
+  Set.insert(1);
+  EXPECT_TRUE(Set.contains(0));
+  EXPECT_TRUE(Set.contains(1));
+
+  Set.insert(1);
+  EXPECT_TRUE(Set.contains(0));
+  EXPECT_TRUE(Set.contains(1));
+
+  Set.erase(1);
+  EXPECT_TRUE(Set.contains(0));
+  EXPECT_FALSE(Set.contains(1));
+
+  Set.insert(1);
+  Set.insert(2);
+  EXPECT_TRUE(Set.contains(0));
+  EXPECT_TRUE(Set.contains(1));
+  EXPECT_TRUE(Set.contains(2));
+}