Index: llvm/trunk/include/llvm/ADT/SmallPtrSet.h =================================================================== --- llvm/trunk/include/llvm/ADT/SmallPtrSet.h +++ llvm/trunk/include/llvm/ADT/SmallPtrSet.h @@ -162,22 +162,38 @@ /// 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()) + return false; + + const void ** Loc = const_cast(P); + assert(*Loc == Ptr && "broken find!"); + *Loc = getTombstoneMarker(); + NumTombstones++; + return true; + } - 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; } @@ -362,7 +378,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: llvm/trunk/lib/Support/SmallPtrSet.cpp =================================================================== --- llvm/trunk/lib/Support/SmallPtrSet.cpp +++ llvm/trunk/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;