[SmallPtrSet] Add remove_if() method (#96468)

Add remove_if() method, similar to the one already present on SetVector.
It is intended to replace the following pattern:

    for (Foo *Ptr : Set)
      if (Pred(Ptr))
        Set.erase(Ptr);

With:

    Set.remove_if(Pred);

This pattern is commonly used for set intersection, where `Pred` is
something like `!OtherSet.contains(Ptr)`.

The implementation provided here is a bit more efficient than the naive
loop, because it does not require looking up the bucket during the
erase() operation again.

However, my actual motivation for this is to have a way to perform this
operation without relying on the current `std::set`-style guarantee that
erase() does not invalidate iterators. I'd like to stop making use of
tombstones in the small regime, which will make insertion operations a
good bit more efficient. However, this will invalidate iterators during
erase().

llvm/include/llvm/ADT/SmallPtrSet.h

@@ -356,6 +356,22 @@ class SmallPtrSetImpl : public SmallPtrSetImplBase {
   bool erase(PtrType Ptr) {
     return erase_imp(PtrTraits::getAsVoidPointer(Ptr));
   }
+
+  /// Remove elements that match the given predicate.
+  template <typename UnaryPredicate>
+  void remove_if(UnaryPredicate P) {
+    for (const void **APtr = CurArray, **E = EndPointer(); APtr != E; ++APtr) {
+      const void *Value = *APtr;
+      if (Value == getTombstoneMarker() || Value == getEmptyMarker())
+        continue;
+      PtrType Ptr = PtrTraits::getFromVoidPointer(const_cast<void *>(Value));
+      if (P(Ptr)) {
+        *APtr = getTombstoneMarker();
+        ++NumTombstones;
+      }
+    }
+  }
+
   /// count - Return 1 if the specified pointer is in the set, 0 otherwise.
   size_type count(ConstPtrType Ptr) const {
     return find_imp(ConstPtrTraits::getAsVoidPointer(Ptr)) != EndPointer();

llvm/unittests/ADT/SmallPtrSetTest.cpp

@@ -409,3 +409,36 @@ TEST(SmallPtrSetTest, InsertIterator) {
   for (const auto *Ptr : Buf)
     EXPECT_TRUE(Set.contains(Ptr));
 }
+
+TEST(SmallPtrSetTest, RemoveIf) {
+  SmallPtrSet<int *, 5> Set;
+  int Vals[6] = {0, 1, 2, 3, 4, 5};
+
+  // Stay in small regime.
+  Set.insert(&Vals[0]);
+  Set.insert(&Vals[1]);
+  Set.insert(&Vals[2]);
+  Set.insert(&Vals[3]);
+  Set.erase(&Vals[0]); // Leave a tombstone.
+
+  // Remove odd elements.
+  Set.remove_if([](int *Ptr) { return *Ptr % 2 != 0; });
+  // We should only have element 2 left now.
+  EXPECT_EQ(Set.size(), 1u);
+  EXPECT_TRUE(Set.contains(&Vals[2]));
+
+  // Switch to big regime.
+  Set.insert(&Vals[0]);
+  Set.insert(&Vals[1]);
+  Set.insert(&Vals[3]);
+  Set.insert(&Vals[4]);
+  Set.insert(&Vals[5]);
+  Set.erase(&Vals[0]); // Leave a tombstone.
+
+  // Remove odd elements.
+  Set.remove_if([](int *Ptr) { return *Ptr % 2 != 0; });
+  // We should only have elements 2 and 4 left now.
+  EXPECT_EQ(Set.size(), 2u);
+  EXPECT_TRUE(Set.contains(&Vals[2]));
+  EXPECT_TRUE(Set.contains(&Vals[4]));
+}