[SmallPtrSet] Introduce a find primitive and rewrite count/erase in terms of it

[reames]

Summary

This was originally motivated by a compile time problem I've since figured out how to solve differently, but the cleanup seemed useful. We had the same logic - which essentially implemented find - in several places. By commoning them out, I can implement find and allow erase to be inlined at the call sites if profitable.

[chandlerc]

LGTM, minor formatting tweaks.

include/llvm/ADT/SmallPtrSet.h

197–198

extraneous blank lines?

[sanjoy]

Please also add some tests.

include/llvm/ADT/SmallPtrSet.h

167

Early return here if P == EndPointer()?

169

I'd add an assert here that *Loc is Ptr.

389

What is endPtr? Is this patch based on some other patch?

[reames]

include/llvm/ADT/SmallPtrSet.h

389

endPtr is already in tree. This is almost "end()", but it can't be because of the breaking API check logic. Honestly, the breaking API check logic looks broken to me. :( It doesn't seem to cover anything other than begin/end, whereas we return other forms of iterators and might reasonable compare them against end.

[sanjoy]

include/llvm/ADT/SmallPtrSet.h

389

Yeah, this part looks fine. I was looking at a stale version of SmallPtrSet.

include/llvm/ADT/SmallPtrSet.h

std::pair<const void *const *, bool> insert_imp(const void *Ptr) {
     }
     return insert_imp_big(Ptr);
   }
 
   /// erase_imp - If the set contains the specified pointer, remove it and
   /// return true, otherwise return false.  This is hidden from the client so
   /// that the derived class can check that the right type of pointer is passed
   /// in.
-  bool erase_imp(const void * Ptr);
+  bool erase_imp(const void * Ptr) {
+    const void *const *P = find_imp(Ptr);
+    if (P != EndPointer()) {
+      const void ** Loc = const_cast<const void **>(P);
+      *Loc = getTombstoneMarker();
+      NumTombstones++;
+      return true;
+    }
+    return false;
+  }
+
 
-  bool count_imp(const void * Ptr) const {
+  /// Returns the raw pointer needed to construct an iterator.  If element not
+  /// found, this will be EndPointer.  Otherwise, it will be a pointer to the
+  /// slot which stores Ptr;
+  const void *const * find_imp(const void * Ptr) const {
     if (isSmall()) {
       // Linear search for the item.
       for (const void *const *APtr = SmallArray,
                       *const *E = SmallArray + NumNonEmpty; APtr != E; ++APtr)
         if (*APtr == Ptr)
-          return true;
-      return false;
+          return APtr;
+      return EndPointer();
     }
 
     // Big set case.
-    return *FindBucketFor(Ptr) == Ptr;
+    auto *Bucket = FindBucketFor(Ptr);
+    if (*Bucket == Ptr)
+      return Bucket;
+    return EndPointer();
   }
 
+  
+  
 private:
   bool isSmall() const { return CurArray == SmallArray; }
 
   std::pair<const void *const *, bool> insert_imp_big(const void *Ptr);
 
   const void * const *FindBucketFor(const void *Ptr) const;
   void shrink_and_clear();
 
public:
   /// erase - If the set contains the specified pointer, remove it and return
   /// true, otherwise return false.
   bool erase(PtrType Ptr) {
     return erase_imp(PtrTraits::getAsVoidPointer(Ptr));
   }
 
   /// count - Return 1 if the specified pointer is in the set, 0 otherwise.
   size_type count(PtrType Ptr) const {
-    return count_imp(PtrTraits::getAsVoidPointer(Ptr)) ? 1 : 0;
+    return find(Ptr) != endPtr() ? 1 : 0;
+  }
+  iterator find(PtrType Ptr) const {
+    auto *P = find_imp(PtrTraits::getAsVoidPointer(Ptr));
+    return iterator(P, EndPointer());
   }
 
   template <typename IterT>
   void insert(IterT I, IterT E) {
     for (; I != E; ++I)
       insert(*I);
   }
 

lib/Support/SmallPtrSet.cpp

SmallPtrSetImplBase::insert_imp_big(const void *Ptr) {
   if (*Bucket == getTombstoneMarker())
     --NumTombstones;
   else
     ++NumNonEmpty; // Track density.
   *Bucket = Ptr;
   return std::make_pair(Bucket, true);
 }
 
-bool SmallPtrSetImplBase::erase_imp(const void * Ptr) {
-  if (isSmall()) {
-    // Check to see if it is in the set.
-    for (const void **APtr = CurArray, **E = CurArray + NumNonEmpty; APtr != E;
-         ++APtr)
-      if (*APtr == Ptr) {
-        // If it is in the set, replace this element.
-        *APtr = getTombstoneMarker();
-        ++NumTombstones;
-        return true;
-      }
-
-    return false;
-  }
-
-  // Okay, we know we have space.  Find a hash bucket.
-  void **Bucket = const_cast<void**>(FindBucketFor(Ptr));
-  if (*Bucket != Ptr) return false;  // Not in the set?
-
-  // Set this as a tombstone.
-  *Bucket = getTombstoneMarker();
-  ++NumTombstones;
-  return true;
-}
-
 const void * const *SmallPtrSetImplBase::FindBucketFor(const void *Ptr) const {
   unsigned Bucket = DenseMapInfo<void *>::getHashValue(Ptr) & (CurArraySize-1);
   unsigned ArraySize = CurArraySize;
   unsigned ProbeAmt = 1;
   const void *const *Array = CurArray;
   const void *const *Tombstone = nullptr;
   while (true) {
     // If we found an empty bucket, the pointer doesn't exist in the set.