[vector] Add specialization of count

Author: miscco

stl/inc/vector

@@ -3149,6 +3149,49 @@ _NODISCARD _CONSTEXPR20 _InIt _Find_vbool(_InIt _First, const _InIt _Last, const
 
     return _Last;
 }
+
+template <class _InIt, class _Ty>
+_NODISCARD _CONSTEXPR20 _Iter_diff_t<_InIt> _Count_vbool(_InIt _First, const _InIt _Last, const _Ty& _Val) {
+    _Iter_diff_t<_InIt> _Count = 0;
+
+#ifdef __cpp_lib_bitops // TRANSITION, VSO-1020212
+    _Vbase* _UFirst            = const_cast<_Vbase*>(_First._Myptr);
+    const _Vbase* const _ULast = _Last._Myptr;
+
+    // Fast path for less than _VBITS
+    if (_UFirst == _ULast) {
+        const auto _Mask      = (_Vbase(-1) << _First._Myoff) & (_Vbase(-1) >> (_VBITS - _Last._Myoff));
+        const auto _SelectVal = _Val ? *_UFirst : static_cast<_Vbase>(~*_UFirst);
+        return __builtin_popcount(_SelectVal & _Mask);
+    }
+
+    if (_First._Myoff != 0) {
+        const auto _FirstMask = _Vbase(-1) << _First._Myoff;
+        const auto _SelectVal = _Val ? *_UFirst : static_cast<_Vbase>(~*_UFirst);
+        _Count += __builtin_popcount(_SelectVal & _FirstMask);
+        ++_UFirst;
+    }
+
+    for (; _UFirst != _ULast; ++_UFirst) {
+        const auto _SelectVal = _Val ? *_UFirst : static_cast<_Vbase>(~*_UFirst);
+        _Count += __builtin_popcount(_SelectVal);
+    }
+
+    if (_Last._Myoff != 0) {
+        const auto _LastMask  = _Vbase(-1) >> (_VBITS - _Last._Myoff);
+        const auto _SelectVal = _Val ? *_UFirst : static_cast<_Vbase>(~*_UFirst);
+        _Count += __builtin_popcount(_SelectVal & _LastMask);
+    }
+#else // ^^^ __cpp_lib_bitops / !__cpp_lib_bitops vvv
+    for (; _First != _Last; ++_First) {
+        if (*_First == _Val) {
+            ++_Count;
+        }
+    }
+#endif // !__cpp_lib_bitops
+
+    return _Count;
+}
 #endif // _HAS_IF_CONSTEXPR
 _STD_END
 

stl/inc/xutility

@@ -5012,17 +5012,25 @@ template <class _InIt, class _Ty>
 _NODISCARD _CONSTEXPR20 _Iter_diff_t<_InIt> count(const _InIt _First, const _InIt _Last, const _Ty& _Val) {
     // count elements that match _Val
     _Adl_verify_range(_First, _Last);
-    auto _UFirst               = _Get_unwrapped(_First);
-    const auto _ULast          = _Get_unwrapped(_Last);
-    _Iter_diff_t<_InIt> _Count = 0;
+#if _HAS_IF_CONSTEXPR
+    if constexpr (_Is_vb_iterator<_InIt> && is_same_v<_Ty, bool>) {
+        return _Count_vbool(_First, _Last, _Val);
+    } else {
+#endif // _HAS_IF_CONSTEXPR
+        auto _UFirst               = _Get_unwrapped(_First);
+        const auto _ULast          = _Get_unwrapped(_Last);
+        _Iter_diff_t<_InIt> _Count = 0;
 
-    for (; _UFirst != _ULast; ++_UFirst) {
-        if (*_UFirst == _Val) {
-            ++_Count;
+        for (; _UFirst != _ULast; ++_UFirst) {
+            if (*_UFirst == _Val) {
+                ++_Count;
+            }
         }
-    }
 
-    return _Count;
+        return _Count;
+#if _HAS_IF_CONSTEXPR
+    }
+#endif // _HAS_IF_CONSTEXPR
 }
 
 #if _HAS_CXX17

tests/std/tests/GH_000625_vector_bool_optimization/test.cpp

@@ -887,9 +887,59 @@ bool test_find() {
     return true;
 }
 
+bool test_count() {
+    // Less than blockSize
+    {
+        const auto result_true = count(source.begin(), next(source.begin(), 8), true);
+        assert(result_true == 5);
+
+        const auto result_false = count(source.begin(), next(source.begin(), 8), false);
+        assert(result_false == 3);
+    }
+
+    // Exactly blockSize
+    {
+        const auto result_true = count(source.begin(), next(source.begin(), blockSize), true);
+        assert(result_true == 20);
+
+        const auto result_false = count(source.begin(), next(source.begin(), blockSize), false);
+        assert(result_false == 12);
+    }
+
+    // More than blockSize
+    {
+        const auto result_true = count(source.begin(), next(source.begin(), blockSize + 3), true);
+        assert(result_true == 22);
+
+        const auto result_false = count(source.begin(), next(source.begin(), blockSize + 3), false);
+        assert(result_false == 13);
+    }
+
+    // Multiple blockSize ends at boundary
+    {
+        const auto result_true = count(source.begin(), next(source.begin(), 2 * blockSize), true);
+        assert(result_true == 40);
+
+        const auto result_false = count(source.begin(), next(source.begin(), 2 * blockSize), false);
+        assert(result_false == 24);
+    }
+
+    // Multiple blockSize ends with offset
+    {
+        const auto result_true = count(source.begin(), next(source.begin(), 2 * blockSize + 5), true);
+        assert(result_true == 43);
+
+        const auto result_false = count(source.begin(), next(source.begin(), 2 * blockSize + 5), false);
+        assert(result_false == 26);
+    }
+
+    return true;
+}
+
 int main() {
     test_copy();
     test_fill();
     test_equal();
     test_find();
+    test_count();
 }