1 files changed, 87 insertions, 17 deletions

diff --git a/Source/WTF/wtf/HashSet.h b/Source/WTF/wtf/HashSet.h
index 46c3fe5d1..cf1f94214 100644
--- a/Source/WTF/wtf/HashSet.h
+++ b/Source/WTF/wtf/HashSet.h
@@ -23,6 +23,7 @@
 
 #include <initializer_list>
 #include <wtf/FastMalloc.h>
+#include <wtf/GetPtr.h>
 #include <wtf/HashTable.h>
 
 namespace WTF {
@@ -32,11 +33,12 @@ namespace WTF {
     template<typename Value, typename HashFunctions, typename Traits> class HashSet;
 
     template<typename ValueArg, typename HashArg = typename DefaultHash<ValueArg>::Hash,
-        typename TraitsArg = HashTraits<ValueArg>> class HashSet {
+        typename TraitsArg = HashTraits<ValueArg>> class HashSet final {
         WTF_MAKE_FAST_ALLOCATED;
     private:
         typedef HashArg HashFunctions;
         typedef TraitsArg ValueTraits;
+        typedef typename ValueTraits::TakeType TakeType;
 
     public:
         typedef typename ValueTraits::TraitType ValueType;
@@ -62,8 +64,8 @@ namespace WTF {
 
         void swap(HashSet&);
 
-        int size() const;
-        int capacity() const;
+        unsigned size() const;
+        unsigned capacity() const;
         bool isEmpty() const;
 
         iterator begin() const;
@@ -100,13 +102,24 @@ namespace WTF {
 
         bool remove(const ValueType&);
         bool remove(iterator);
+        template<typename Functor>
+        void removeIf(const Functor&);
         void clear();
 
-        ValueType take(const ValueType&);
+        TakeType take(const ValueType&);
+        TakeType take(iterator);
+        TakeType takeAny();
+
+        // Overloads for smart pointer values that take the raw pointer type as the parameter.
+        template<typename V = ValueType> typename std::enable_if<IsSmartPtr<V>::value, iterator>::type find(typename GetPtrHelper<V>::PtrType) const;
+        template<typename V = ValueType> typename std::enable_if<IsSmartPtr<V>::value, bool>::type contains(typename GetPtrHelper<V>::PtrType) const;
+        template<typename V = ValueType> typename std::enable_if<IsSmartPtr<V>::value, bool>::type remove(typename GetPtrHelper<V>::PtrType);
+        template<typename V = ValueType> typename std::enable_if<IsSmartPtr<V>::value, TakeType>::type take(typename GetPtrHelper<V>::PtrType);
 
         static bool isValidValue(const ValueType&);
-        
-        bool operator==(const HashSet&) const;
+
+        template<typename OtherCollection>
+        bool operator==(const OtherCollection&) const;
 
     private:
         HashTableType m_impl;
@@ -116,6 +129,16 @@ namespace WTF {
         template<typename T> static const T& extract(const T& t) { return t; }
     };
 
+    template<typename ValueTraits, typename HashFunctions>
+    struct HashSetTranslator {
+        template<typename T> static unsigned hash(const T& key) { return HashFunctions::hash(key); }
+        template<typename T, typename U> static bool equal(const T& a, const U& b) { return HashFunctions::equal(a, b); }
+        template<typename T, typename U, typename V> static void translate(T& location, U&&, V&& value)
+        { 
+            ValueTraits::assignToEmpty(location, std::forward<V>(value));
+        }
+    };
+
     template<typename Translator>
     struct HashSetTranslatorAdapter {
         template<typename T> static unsigned hash(const T& key) { return Translator::hash(key); }
@@ -133,13 +156,13 @@ namespace WTF {
     }
 
     template<typename T, typename U, typename V>
-    inline int HashSet<T, U, V>::size() const
+    inline unsigned HashSet<T, U, V>::size() const
     {
         return m_impl.size(); 
     }
 
     template<typename T, typename U, typename V>
-    inline int HashSet<T, U, V>::capacity() const
+    inline unsigned HashSet<T, U, V>::capacity() const
     {
         return m_impl.capacity(); 
     }
@@ -197,7 +220,7 @@ namespace WTF {
     template<typename T, typename U, typename V>
     inline auto HashSet<T, U, V>::add(ValueType&& value) -> AddResult
     {
-        return m_impl.add(std::move(value));
+        return m_impl.add(WTFMove(value));
     }
 
     template<typename Value, typename HashFunctions, typename Traits>
@@ -234,24 +257,70 @@ namespace WTF {
     }
 
     template<typename T, typename U, typename V>
+    template<typename Functor>
+    inline void HashSet<T, U, V>::removeIf(const Functor& functor)
+    {
+        m_impl.removeIf(functor);
+    }
+
+    template<typename T, typename U, typename V>
     inline void HashSet<T, U, V>::clear()
     {
         m_impl.clear(); 
     }
 
     template<typename T, typename U, typename V>
-    auto HashSet<T, U, V>::take(const ValueType& value) -> ValueType
+    inline auto HashSet<T, U, V>::take(iterator it) -> TakeType
     {
-        auto it = find(value);
         if (it == end())
-            return ValueTraits::emptyValue();
+            return ValueTraits::take(ValueTraits::emptyValue());
 
-        ValueType result = std::move(const_cast<ValueType&>(*it));
+        auto result = ValueTraits::take(WTFMove(const_cast<ValueType&>(*it)));
         remove(it);
         return result;
     }
 
     template<typename T, typename U, typename V>
+    inline auto HashSet<T, U, V>::take(const ValueType& value) -> TakeType
+    {
+        return take(find(value));
+    }
+
+    template<typename T, typename U, typename V>
+    inline auto HashSet<T, U, V>::takeAny() -> TakeType
+    {
+        return take(begin());
+    }
+
+    template<typename Value, typename HashFunctions, typename Traits>
+    template<typename V>
+    inline auto HashSet<Value, HashFunctions, Traits>::find(typename GetPtrHelper<V>::PtrType value) const -> typename std::enable_if<IsSmartPtr<V>::value, iterator>::type
+    {
+        return m_impl.template find<HashSetTranslator<Traits, HashFunctions>>(value);
+    }
+
+    template<typename Value, typename HashFunctions, typename Traits>
+    template<typename V>
+    inline auto HashSet<Value, HashFunctions, Traits>::contains(typename GetPtrHelper<V>::PtrType value) const -> typename std::enable_if<IsSmartPtr<V>::value, bool>::type
+    {
+        return m_impl.template contains<HashSetTranslator<Traits, HashFunctions>>(value);
+    }
+
+    template<typename Value, typename HashFunctions, typename Traits>
+    template<typename V>
+    inline auto HashSet<Value, HashFunctions, Traits>::remove(typename GetPtrHelper<V>::PtrType value) -> typename std::enable_if<IsSmartPtr<V>::value, bool>::type
+    {
+        return remove(find(value));
+    }
+
+    template<typename Value, typename HashFunctions, typename Traits>
+    template<typename V>
+    inline auto HashSet<Value, HashFunctions, Traits>::take(typename GetPtrHelper<V>::PtrType value) -> typename std::enable_if<IsSmartPtr<V>::value, TakeType>::type
+    {
+        return take(find(value));
+    }
+
+    template<typename T, typename U, typename V>
     inline bool HashSet<T, U, V>::isValidValue(const ValueType& value)
     {
         if (ValueTraits::isDeletedValue(value))
@@ -282,12 +351,13 @@ namespace WTF {
     }  
 
     template<typename T, typename U, typename V>
-    inline bool HashSet<T, U, V>::operator==(const HashSet& other) const
+    template<typename OtherCollection>
+    inline bool HashSet<T, U, V>::operator==(const OtherCollection& otherCollection) const
     {
-        if (size() != other.size())
+        if (size() != otherCollection.size())
             return false;
-        for (const_iterator iter = begin(); iter != end(); ++iter) {
-            if (!other.contains(*iter))
+        for (const auto& other : otherCollection) {
+            if (!contains(other))
                 return false;
         }
         return true;