Index: include/llvm/ADT/SmallPtrSet.h =================================================================== --- include/llvm/ADT/SmallPtrSet.h +++ include/llvm/ADT/SmallPtrSet.h @@ -162,22 +162,40 @@ /// 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(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; } @@ -368,7 +386,11 @@ /// 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 Index: lib/Support/SmallPtrSet.cpp =================================================================== --- lib/Support/SmallPtrSet.cpp +++ lib/Support/SmallPtrSet.cpp @@ -61,31 +61,6 @@ 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(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::getHashValue(Ptr) & (CurArraySize-1); unsigned ArraySize = CurArraySize;