[libc++] Fix UB in bitwise logic of {std, ranges}::{fill, fill_n} algorithms (#122410)

This PR addresses an undefined behavior that arises when using the
`std::fill` and `std::fill_n` algorithms, as well as their ranges
counterparts `ranges::fill` and `ranges::fill_n`, with `vector<bool, Alloc>`
that utilizes a custom-sized allocator with small integral types.

Author: winner245

libcxx/include/__algorithm/fill_n.h

@@ -41,11 +41,7 @@ __fill_n_bool(__bit_iterator<_Cp, false> __first, typename __size_difference_typ
   if (__first.__ctz_ != 0) {
     __storage_type __clz_f = static_cast<__storage_type>(__bits_per_word - __first.__ctz_);
     __storage_type __dn    = std::min(__clz_f, __n);
-    __storage_type __m     = (~__storage_type(0) << __first.__ctz_) & (~__storage_type(0) >> (__clz_f - __dn));
-    if (_FillVal)
-      *__first.__seg_ |= __m;
-    else
-      *__first.__seg_ &= ~__m;
+    std::__fill_masked_range(std::__to_address(__first.__seg_), __clz_f - __dn, __first.__ctz_, _FillVal);
     __n -= __dn;
     ++__first.__seg_;
   }
@@ -56,11 +52,7 @@ __fill_n_bool(__bit_iterator<_Cp, false> __first, typename __size_difference_typ
   // do last partial word
   if (__n > 0) {
     __first.__seg_ += __nw;
-    __storage_type __m = ~__storage_type(0) >> (__bits_per_word - __n);
-    if (_FillVal)
-      *__first.__seg_ |= __m;
-    else
-      *__first.__seg_ &= ~__m;
+    std::__fill_masked_range(std::__to_address(__first.__seg_), __bits_per_word - __n, 0u, _FillVal);
   }
 }
 

libcxx/include/__bit_reference

@@ -14,6 +14,7 @@
 #include <__algorithm/copy_backward.h>
 #include <__algorithm/copy_n.h>
 #include <__algorithm/min.h>
+#include <__assert>
 #include <__bit/countr.h>
 #include <__compare/ordering.h>
 #include <__config>
@@ -24,10 +25,13 @@
 #include <__memory/construct_at.h>
 #include <__memory/pointer_traits.h>
 #include <__type_traits/conditional.h>
+#include <__type_traits/enable_if.h>
 #include <__type_traits/is_constant_evaluated.h>
+#include <__type_traits/is_unsigned.h>
 #include <__type_traits/void_t.h>
 #include <__utility/pair.h>
 #include <__utility/swap.h>
+#include <climits>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -58,6 +62,26 @@ struct __size_difference_type_traits<_Cp, __void_t<typename _Cp::difference_type
   using size_type       = typename _Cp::size_type;
 };
 
+// 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.
+// See https://github.com/llvm/llvm-project/pull/122410.
+template <class _StoragePointer>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
+__fill_masked_range(_StoragePointer __word, unsigned __clz, unsigned __ctz, bool __fill_val) {
+  static_assert(is_unsigned<typename pointer_traits<_StoragePointer>::element_type>::value,
+                "__fill_masked_range must be called with unsigned type");
+  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);
+  if (__fill_val)
+    *__word |= __m;
+  else
+    *__word &= static_cast<_StorageType>(~__m);
+}
+
 template <class _Cp, bool = __has_storage_type<_Cp>::value>
 class __bit_reference {
   using __storage_type _LIBCPP_NODEBUG    = typename _Cp::__storage_type;

libcxx/include/__fwd/bit_reference.h

@@ -10,6 +10,9 @@
 #define _LIBCPP___FWD_BIT_REFERENCE_H
 
 #include <__config>
+#include <__memory/pointer_traits.h>
+#include <__type_traits/enable_if.h>
+#include <__type_traits/is_unsigned.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -23,6 +26,10 @@ class __bit_iterator;
 template <class, class = void>
 struct __size_difference_type_traits;
 
+template <class _StoragePointer>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
+__fill_masked_range(_StoragePointer __word, unsigned __ctz, unsigned __clz, bool __fill_val);
+
 _LIBCPP_END_NAMESPACE_STD
 
 #endif // _LIBCPP___FWD_BIT_REFERENCE_H

libcxx/test/std/algorithms/alg.modifying.operations/alg.fill/fill.pass.cpp

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <vector>
 
+#include "sized_allocator.h"
 #include "test_macros.h"
 #include "test_iterators.h"
 
@@ -46,6 +47,39 @@ struct Test {
   }
 };
 
+// Make sure std::fill behaves properly with std::vector<bool> iterators with custom size types.
+// See https://github.com/llvm/llvm-project/pull/122410.
+TEST_CONSTEXPR_CXX20 void test_bititer_with_custom_sized_types() {
+  {
+    using Alloc = sized_allocator<bool, std::uint8_t, std::int8_t>;
+    std::vector<bool, Alloc> in(100, false, Alloc(1));
+    std::vector<bool, Alloc> expected(100, true, Alloc(1));
+    std::fill(in.begin(), in.end(), true);
+    assert(in == expected);
+  }
+  {
+    using Alloc = sized_allocator<bool, std::uint16_t, std::int16_t>;
+    std::vector<bool, Alloc> in(200, false, Alloc(1));
+    std::vector<bool, Alloc> expected(200, true, Alloc(1));
+    std::fill(in.begin(), in.end(), true);
+    assert(in == expected);
+  }
+  {
+    using Alloc = sized_allocator<bool, std::uint32_t, std::int32_t>;
+    std::vector<bool, Alloc> in(200, false, Alloc(1));
+    std::vector<bool, Alloc> expected(200, true, Alloc(1));
+    std::fill(in.begin(), in.end(), true);
+    assert(in == expected);
+  }
+  {
+    using Alloc = sized_allocator<bool, std::uint64_t, std::int64_t>;
+    std::vector<bool, Alloc> in(200, false, Alloc(1));
+    std::vector<bool, Alloc> expected(200, true, Alloc(1));
+    std::fill(in.begin(), in.end(), true);
+    assert(in == expected);
+  }
+}
+
 TEST_CONSTEXPR_CXX20 bool test() {
   types::for_each(types::forward_iterator_list<char*>(), Test<char>());
   types::for_each(types::forward_iterator_list<int*>(), Test<int>());
@@ -93,6 +127,9 @@ TEST_CONSTEXPR_CXX20 bool test() {
       assert(in == expected);
     }
   }
+
+  test_bititer_with_custom_sized_types();
+
   return true;
 }