Fix ranges::equal for vector<bool> with small storage types

Author: winner245

libcxx/include/__algorithm/equal.h

@@ -54,24 +54,27 @@ __equal_unaligned(__bit_iterator<_Cp, _IsConst1> __first1,
       unsigned __clz_f     = __bits_per_word - __first1.__ctz_;
       difference_type __dn = std::min(static_cast<difference_type>(__clz_f), __n);
       __n -= __dn;
-      __storage_type __m   = (~__storage_type(0) << __first1.__ctz_) & (~__storage_type(0) >> (__clz_f - __dn));
+      __storage_type __m   = std::__middle_mask<__storage_type>(__clz_f - __dn, __first1.__ctz_);
       __storage_type __b   = *__first1.__seg_ & __m;
       unsigned __clz_r     = __bits_per_word - __first2.__ctz_;
       __storage_type __ddn = std::min<__storage_type>(__dn, __clz_r);
-      __m                  = (~__storage_type(0) << __first2.__ctz_) & (~__storage_type(0) >> (__clz_r - __ddn));
+      __m                  = std::__middle_mask<__storage_type>(__clz_r - __ddn, __first2.__ctz_);
       if (__first2.__ctz_ > __first1.__ctz_) {
-        if ((*__first2.__seg_ & __m) != (__b << (__first2.__ctz_ - __first1.__ctz_)))
+        if (static_cast<__storage_type>(*__first2.__seg_ & __m) !=
+            static_cast<__storage_type>(__b << (__first2.__ctz_ - __first1.__ctz_)))
           return false;
       } else {
-        if ((*__first2.__seg_ & __m) != (__b >> (__first1.__ctz_ - __first2.__ctz_)))
+        if (static_cast<__storage_type>(*__first2.__seg_ & __m) !=
+            static_cast<__storage_type>(__b >> (__first1.__ctz_ - __first2.__ctz_)))
           return false;
       }
       __first2.__seg_ += (__ddn + __first2.__ctz_) / __bits_per_word;
       __first2.__ctz_ = static_cast<unsigned>((__ddn + __first2.__ctz_) % __bits_per_word);
       __dn -= __ddn;
       if (__dn > 0) {
-        __m = ~__storage_type(0) >> (__bits_per_word - __dn);
-        if ((*__first2.__seg_ & __m) != (__b >> (__first1.__ctz_ + __ddn)))
+        __m = std::__trailing_mask<__storage_type>(__bits_per_word - __n);
+        if (static_cast<__storage_type>(*__first2.__seg_ & __m) !=
+            static_cast<__storage_type>(__b >> (__first1.__ctz_ + __ddn)))
           return false;
         __first2.__ctz_ = static_cast<unsigned>(__dn);
       }
@@ -81,29 +84,30 @@ __equal_unaligned(__bit_iterator<_Cp, _IsConst1> __first1,
     // __first1.__ctz_ == 0;
     // do middle words
     unsigned __clz_r   = __bits_per_word - __first2.__ctz_;
-    __storage_type __m = ~__storage_type(0) << __first2.__ctz_;
+    __storage_type __m = std::__leading_mask<__storage_type>(__first2.__ctz_);
     for (; __n >= __bits_per_word; __n -= __bits_per_word, ++__first1.__seg_) {
       __storage_type __b = *__first1.__seg_;
-      if ((*__first2.__seg_ & __m) != (__b << __first2.__ctz_))
+      if (static_cast<__storage_type>(*__first2.__seg_ & __m) != static_cast<__storage_type>(__b << __first2.__ctz_))
         return false;
       ++__first2.__seg_;
-      if ((*__first2.__seg_ & ~__m) != (__b >> __clz_r))
+      if (static_cast<__storage_type>(*__first2.__seg_ & static_cast<__storage_type>(~__m)) !=
+          static_cast<__storage_type>(__b >> __clz_r))
         return false;
     }
     // do last word
     if (__n > 0) {
-      __m                 = ~__storage_type(0) >> (__bits_per_word - __n);
+      __m                 = std::__trailing_mask<__storage_type>(__bits_per_word - __n);
       __storage_type __b  = *__first1.__seg_ & __m;
       __storage_type __dn = std::min(__n, static_cast<difference_type>(__clz_r));
-      __m                 = (~__storage_type(0) << __first2.__ctz_) & (~__storage_type(0) >> (__clz_r - __dn));
-      if ((*__first2.__seg_ & __m) != (__b << __first2.__ctz_))
+      __m                 = std::__middle_mask<__storage_type>(__clz_r - __dn, __first2.__ctz_);
+      if (static_cast<__storage_type>(*__first2.__seg_ & __m) != static_cast<__storage_type>(__b << __first2.__ctz_))
         return false;
       __first2.__seg_ += (__dn + __first2.__ctz_) / __bits_per_word;
       __first2.__ctz_ = static_cast<unsigned>((__dn + __first2.__ctz_) % __bits_per_word);
       __n -= __dn;
       if (__n > 0) {
-        __m = ~__storage_type(0) >> (__bits_per_word - __n);
-        if ((*__first2.__seg_ & __m) != (__b >> __dn))
+        __m = std::__trailing_mask<__storage_type>(__bits_per_word - __n);
+        if (static_cast<__storage_type>(*__first2.__seg_ & __m) != static_cast<__storage_type>(__b >> __dn))
           return false;
       }
     }
@@ -128,7 +132,7 @@ __equal_aligned(__bit_iterator<_Cp, _IsConst1> __first1,
       unsigned __clz       = __bits_per_word - __first1.__ctz_;
       difference_type __dn = std::min(static_cast<difference_type>(__clz), __n);
       __n -= __dn;
-      __storage_type __m = (~__storage_type(0) << __first1.__ctz_) & (~__storage_type(0) >> (__clz - __dn));
+      __storage_type __m = std::__middle_mask<__storage_type>(__clz - __dn, __first1.__ctz_);
       if ((*__first2.__seg_ & __m) != (*__first1.__seg_ & __m))
         return false;
       ++__first2.__seg_;
@@ -144,7 +148,7 @@ __equal_aligned(__bit_iterator<_Cp, _IsConst1> __first1,
         return false;
     // do last word
     if (__n > 0) {
-      __storage_type __m = ~__storage_type(0) >> (__bits_per_word - __n);
+      __storage_type __m = std::__trailing_mask<__storage_type>(__bits_per_word - __n);
       if ((*__first2.__seg_ & __m) != (*__first1.__seg_ & __m))
         return false;
     }

libcxx/include/__bit_reference

@@ -68,6 +68,30 @@ struct __size_difference_type_traits<_Cp, __void_t<typename _Cp::difference_type
   using size_type       = typename _Cp::size_type;
 };
 
+// Creates a mask of type `_StorageType` with a specified number of leading zeros (__clz) and sets all remaining
+// bits to one.
+template <class _StorageType>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __trailing_mask(unsigned __clz) {
+  static_assert(is_unsigned<_StorageType>::value, "__trailing_mask only works with unsigned types");
+  return static_cast<_StorageType>(~static_cast<_StorageType>(0)) >> __clz;
+}
+
+// Creates a mask of type `_StorageType` with a specified number of trailing zeros (__ctz) and sets all remaining
+// bits to one.
+template <class _StorageType>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __leading_mask(unsigned __ctz) {
+  static_assert(is_unsigned<_StorageType>::value, "__leading_mask only works with unsigned types");
+  return static_cast<_StorageType>(~static_cast<_StorageType>(0)) << __ctz;
+}
+
+// Creates a mask of type `_StorageType` with a specified number of leading zeros (__clz), a specified number of
+// trailing zeros (__ctz), and sets all bits in between to one.
+template <class _StorageType>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __middle_mask(unsigned __clz, unsigned __ctz) {
+  static_assert(is_unsigned<_StorageType>::value, "__middle_mask only works with unsigned types");
+  return std::__leading_mask<_StorageType>(__ctz) & std::__trailing_mask<_StorageType>(__clz);
+}
+
 // This function is designed to operate correctly even for smaller integral types like `uint8_t`, `uint16_t`,
 // or `unsigned short`. Casting back to _StorageType is crucial to prevent undefined behavior that can arise
 // from integral promotions.
@@ -80,8 +104,7 @@ __fill_masked_range(_StoragePointer __word, unsigned __clz, unsigned __ctz, bool
   using _StorageType = typename pointer_traits<_StoragePointer>::element_type;
   _LIBCPP_ASSERT_VALID_INPUT_RANGE(
       __ctz + __clz < sizeof(_StorageType) * CHAR_BIT, "__fill_masked_range called with invalid range");
-  _StorageType __m = static_cast<_StorageType>(static_cast<_StorageType>(~static_cast<_StorageType>(0)) >> __clz) &
-                     static_cast<_StorageType>(static_cast<_StorageType>(~static_cast<_StorageType>(0)) << __ctz);
+  _StorageType __m = std::__middle_mask<_StorageType>(__clz, __ctz);
   if (__fill_val)
     *__word |= __m;
   else

libcxx/include/__fwd/bit_reference.h

@@ -30,6 +30,15 @@ template <class _StoragePointer>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
 __fill_masked_range(_StoragePointer __word, unsigned __ctz, unsigned __clz, bool __fill_val);
 
+template <class _StorageType>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __trailing_mask(unsigned __clz);
+
+template <class _StorageType>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __leading_mask(unsigned __ctz);
+
+template <class _StorageType>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __middle_mask(unsigned __clz, unsigned __ctz);
+
 _LIBCPP_END_NAMESPACE_STD
 
 #endif // _LIBCPP___FWD_BIT_REFERENCE_H

libcxx/include/__vector/comparison.h

@@ -29,20 +29,6 @@ operator==(const vector<_Tp, _Allocator>& __x, const vector<_Tp, _Allocator>& __
   return __sz == __y.size() && std::equal(__x.begin(), __x.end(), __y.begin());
 }
 
-// FIXME: Remove this `vector<bool>` overload once #126369 is resolved, reverting to the generic `operator==`
-// with `std::equal` for better performance.
-template <class _Allocator>
-_LIBCPP_CONSTEXPR_SINCE_CXX20 inline _LIBCPP_HIDE_FROM_ABI bool
-operator==(const vector<bool, _Allocator>& __x, const vector<bool, _Allocator>& __y) {
-  const typename vector<bool, _Allocator>::size_type __sz = __x.size();
-  if (__sz != __y.size())
-    return false;
-  for (typename vector<bool, _Allocator>::size_type __i = 0; __i < __sz; ++__i)
-    if (__x[__i] != __y[__i])
-      return false;
-  return true;
-}
-
 #if _LIBCPP_STD_VER <= 17
 
 template <class _Tp, class _Allocator>

libcxx/test/std/algorithms/alg.nonmodifying/alg.equal/equal.pass.cpp

@@ -24,12 +24,14 @@
 // MSVC warning C4244: 'argument': conversion from 'wchar_t' to 'const _Ty', possible loss of data
 // MSVC warning C4389: '==': signed/unsigned mismatch
 // ADDITIONAL_COMPILE_FLAGS(cl-style-warnings): /wd4242 /wd4244 /wd4389
+// XFAIL: FROZEN-CXX03-HEADERS-FIXME
 
 #include <algorithm>
 #include <cassert>
 #include <functional>
 #include <vector>
 
+#include "sized_allocator.h"
 #include "test_iterators.h"
 #include "test_macros.h"
 #include "type_algorithms.h"
@@ -173,6 +175,90 @@ TEST_CONSTEXPR_CXX20 bool test() {
     test_vector_bool<256>();
   }
 
+  // Make sure std::equal behaves properly with std::vector<bool> iterators with custom size types.
+  // See issue: https://github.com/llvm/llvm-project/issues/126369.
+  {
+    //// Tests for std::equal with aligned bits
+
+    { // Test the first (partial) word for uint8_t
+      using Alloc = sized_allocator<bool, std::uint8_t, std::int8_t>;
+      std::vector<bool, Alloc> in(6, true, Alloc(1));
+      std::vector<bool, Alloc> expected(8, true, Alloc(1));
+      assert(std::equal(in.begin() + 4, in.end(), expected.begin() + 4));
+    }
+    { // Test the last word for uint8_t
+      using Alloc = sized_allocator<bool, std::uint8_t, std::int8_t>;
+      std::vector<bool, Alloc> in(12, true, Alloc(1));
+      std::vector<bool, Alloc> expected(16, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin()));
+    }
+    { // Test middle words for uint8_t
+      using Alloc = sized_allocator<bool, std::uint8_t, std::int8_t>;
+      std::vector<bool, Alloc> in(24, true, Alloc(1));
+      std::vector<bool, Alloc> expected(29, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin()));
+    }
+
+    { // Test the first (partial) word for uint16_t
+      using Alloc = sized_allocator<bool, std::uint16_t, std::int16_t>;
+      std::vector<bool, Alloc> in(12, true, Alloc(1));
+      std::vector<bool, Alloc> expected(16, true, Alloc(1));
+      assert(std::equal(in.begin() + 4, in.end(), expected.begin() + 4));
+    }
+    { // Test the last word for uint16_t
+      using Alloc = sized_allocator<bool, std::uint16_t, std::int16_t>;
+      std::vector<bool, Alloc> in(24, true, Alloc(1));
+      std::vector<bool, Alloc> expected(32, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin()));
+    }
+    { // Test middle words for uint16_t
+      using Alloc = sized_allocator<bool, std::uint16_t, std::int16_t>;
+      std::vector<bool, Alloc> in(48, true, Alloc(1));
+      std::vector<bool, Alloc> expected(55, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin()));
+    }
+
+    //// Tests for std::equal with unaligned bits
+
+    { // Test the first (partial) word for uint8_t
+      using Alloc = sized_allocator<bool, std::uint8_t, std::int8_t>;
+      std::vector<bool, Alloc> in(6, true, Alloc(1));
+      std::vector<bool, Alloc> expected(8, true, Alloc(1));
+      assert(std::equal(in.begin() + 4, in.end(), expected.begin()));
+    }
+    { // Test the last word for uint8_t
+      using Alloc = sized_allocator<bool, std::uint8_t, std::int8_t>;
+      std::vector<bool, Alloc> in(4, true, Alloc(1));
+      std::vector<bool, Alloc> expected(8, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin() + 3));
+    }
+    { // Test middle words for uint8_t
+      using Alloc = sized_allocator<bool, std::uint8_t, std::int8_t>;
+      std::vector<bool, Alloc> in(16, true, Alloc(1));
+      std::vector<bool, Alloc> expected(24, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin() + 4));
+    }
+
+    { // Test the first (partial) word for uint16_t
+      using Alloc = sized_allocator<bool, std::uint16_t, std::int16_t>;
+      std::vector<bool, Alloc> in(12, true, Alloc(1));
+      std::vector<bool, Alloc> expected(16, true, Alloc(1));
+      assert(std::equal(in.begin() + 4, in.end(), expected.begin()));
+    }
+    { // Test the last word for uint16_t
+      using Alloc = sized_allocator<bool, std::uint16_t, std::int16_t>;
+      std::vector<bool, Alloc> in(12, true, Alloc(1));
+      std::vector<bool, Alloc> expected(16, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin() + 3));
+    }
+    { // Test the middle words for uint16_t
+      using Alloc = sized_allocator<bool, std::uint16_t, std::int16_t>;
+      std::vector<bool, Alloc> in(32, true, Alloc(1));
+      std::vector<bool, Alloc> expected(64, true, Alloc(1));
+      assert(std::equal(in.begin(), in.end(), expected.begin() + 4));
+    }
+  }
+
   return true;
 }