Fix UB in bitwise logic of std/ranges::fill/fill_n algorithms

Author: winner245

libcxx/include/__algorithm/fill_n.h

@@ -41,7 +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));
+    __storage_type __m     = std::__middle_mask<__storage_type>(__first.__ctz_, __clz_f - __dn);
     if (_FillVal)
       *__first.__seg_ |= __m;
     else
@@ -56,7 +56,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);
+    __storage_type __m = std::__trailing_mask<__storage_type>(__bits_per_word - __n);
     if (_FillVal)
       *__first.__seg_ |= __m;
     else

libcxx/include/__fwd/bit_reference.h

@@ -10,6 +10,8 @@
 #define _LIBCPP___FWD_BIT_REFERENCE_H
 
 #include <__config>
+#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 +25,22 @@ class __bit_iterator;
 template <class, class = void>
 struct __size_difference_type_traits;
 
+template <class _StorageType, __enable_if_t<is_unsigned<_StorageType>::value, int> = 0>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __leading_mask(unsigned __shift) {
+  return static_cast<_StorageType>(static_cast<_StorageType>(~static_cast<_StorageType>(0)) << __shift);
+}
+
+template <class _StorageType, __enable_if_t<is_unsigned<_StorageType>::value, int> = 0>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __trailing_mask(unsigned __shift) {
+  return static_cast<_StorageType>(static_cast<_StorageType>(~static_cast<_StorageType>(0)) >> __shift);
+}
+
+template <class _StorageType, __enable_if_t<is_unsigned<_StorageType>::value, int> = 0>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _StorageType __middle_mask(unsigned __lshift, unsigned __rshift) {
+  return static_cast<_StorageType>(
+      std::__leading_mask<_StorageType>(__lshift) & std::__trailing_mask<_StorageType>(__rshift));
+}
+
 _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,37 @@ struct Test {
   }
 };
 
+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 +125,9 @@ TEST_CONSTEXPR_CXX20 bool test() {
       assert(in == expected);
     }
   }
+
+  test_bititer_with_custom_sized_types();
+
   return true;
 }
 

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

@@ -15,159 +15,175 @@
 
 #include <algorithm>
 #include <cassert>
+#include <vector>
 
+#include "sized_allocator.h"
 #include "test_macros.h"
 #include "test_iterators.h"
 #include "user_defined_integral.h"
 
 #if TEST_STD_VER > 17
 TEST_CONSTEXPR bool test_constexpr() {
-    const std::size_t N = 5;
-    int ib[] = {0, 0, 0, 0, 0, 0}; // one bigger than N
-
-    auto it = std::fill_n(std::begin(ib), N, 5);
-    return it == (std::begin(ib) + N)
-        && std::all_of(std::begin(ib), it, [](int a) {return a == 5; })
-        && *it == 0 // don't overwrite the last value in the output array
-        ;
-    }
+  const std::size_t N = 5;
+  int ib[]            = {0, 0, 0, 0, 0, 0}; // one bigger than N
+
+  auto it = std::fill_n(std::begin(ib), N, 5);
+  return it == (std::begin(ib) + N) && std::all_of(std::begin(ib), it, [](int a) { return a == 5; }) &&
+       *it == 0 // don't overwrite the last value in the output array
+      ;
+}
 #endif
 
 typedef UserDefinedIntegral<unsigned> UDI;
 
 template <class Iter>
-void
-test_char()
-{
-    char a[4] = {};
-    Iter it = std::fill_n(Iter(a), UDI(4), char(1));
-    assert(base(it) == a + 4);
-    assert(a[0] == 1);
-    assert(a[1] == 1);
-    assert(a[2] == 1);
-    assert(a[3] == 1);
+void test_char() {
+  char a[4] = {};
+  Iter it   = std::fill_n(Iter(a), UDI(4), char(1));
+  assert(base(it) == a + 4);
+  assert(a[0] == 1);
+  assert(a[1] == 1);
+  assert(a[2] == 1);
+  assert(a[3] == 1);
 }
 
 template <class Iter>
-void
-test_int()
-{
-    int a[4] = {};
-    Iter it = std::fill_n(Iter(a), UDI(4), 1);
-    assert(base(it) == a + 4);
-    assert(a[0] == 1);
-    assert(a[1] == 1);
-    assert(a[2] == 1);
-    assert(a[3] == 1);
+void test_int() {
+  int a[4] = {};
+  Iter it  = std::fill_n(Iter(a), UDI(4), 1);
+  assert(base(it) == a + 4);
+  assert(a[0] == 1);
+  assert(a[1] == 1);
+  assert(a[2] == 1);
+  assert(a[3] == 1);
 }
 
-void
-test_int_array()
-{
-    int a[4] = {};
-    assert(std::fill_n(a, UDI(4), static_cast<char>(1)) == a + 4);
-    assert(a[0] == 1);
-    assert(a[1] == 1);
-    assert(a[2] == 1);
-    assert(a[3] == 1);
+void test_int_array() {
+  int a[4] = {};
+  assert(std::fill_n(a, UDI(4), static_cast<char>(1)) == a + 4);
+  assert(a[0] == 1);
+  assert(a[1] == 1);
+  assert(a[2] == 1);
+  assert(a[3] == 1);
 }
 
 struct source {
-    source() : i(0) { }
+  source() : i(0) {}
 
-    operator int() const { return i++; }
-    mutable int i;
+  operator int() const { return i++; }
+  mutable int i;
 };
 
-void
-test_int_array_struct_source()
-{
-    int a[4] = {};
-    assert(std::fill_n(a, UDI(4), source()) == a + 4);
-    assert(a[0] == 0);
-    assert(a[1] == 1);
-    assert(a[2] == 2);
-    assert(a[3] == 3);
+void test_int_array_struct_source() {
+  int a[4] = {};
+  assert(std::fill_n(a, UDI(4), source()) == a + 4);
+  assert(a[0] == 0);
+  assert(a[1] == 1);
+  assert(a[2] == 2);
+  assert(a[3] == 3);
 }
 
 struct test1 {
-    test1() : c(0) { }
-    test1(char xc) : c(xc + 1) { }
-    char c;
+  test1() : c(0) {}
+  test1(char xc) : c(xc + 1) {}
+  char c;
 };
 
-void
-test_struct_array()
-{
-    test1 test1a[4] = {};
-    assert(std::fill_n(test1a, UDI(4), static_cast<char>(10)) == test1a + 4);
-    assert(test1a[0].c == 11);
-    assert(test1a[1].c == 11);
-    assert(test1a[2].c == 11);
-    assert(test1a[3].c == 11);
+void test_struct_array() {
+  test1 test1a[4] = {};
+  assert(std::fill_n(test1a, UDI(4), static_cast<char>(10)) == test1a + 4);
+  assert(test1a[0].c == 11);
+  assert(test1a[1].c == 11);
+  assert(test1a[2].c == 11);
+  assert(test1a[3].c == 11);
 }
 
-class A
-{
-    char a_;
+class A {
+  char a_;
+
 public:
-    A() {}
-    explicit A(char a) : a_(a) {}
-    operator unsigned char() const {return 'b';}
+  A() {}
+  explicit A(char a) : a_(a) {}
+  operator unsigned char() const { return 'b'; }
 
-    friend bool operator==(const A& x, const A& y)
-        {return x.a_ == y.a_;}
+  friend bool operator==(const A& x, const A& y) { return x.a_ == y.a_; }
 };
 
-void
-test5()
-{
-    A a[3];
-    assert(std::fill_n(&a[0], UDI(3), A('a')) == a+3);
-    assert(a[0] == A('a'));
-    assert(a[1] == A('a'));
-    assert(a[2] == A('a'));
+void test5() {
+  A a[3];
+  assert(std::fill_n(&a[0], UDI(3), A('a')) == a + 3);
+  assert(a[0] == A('a'));
+  assert(a[1] == A('a'));
+  assert(a[2] == A('a'));
 }
 
-struct Storage
-{
-  union
-  {
+struct Storage {
+  union {
     unsigned char a;
     unsigned char b;
   };
 };
 
-void test6()
-{
+void test6() {
   Storage foo[5];
   std::fill_n(&foo[0], UDI(5), Storage());
 }
 
+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_n(in.begin(), in.size(), 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_n(in.begin(), in.size(), 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_n(in.begin(), in.size(), 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_n(in.begin(), in.size(), true);
+    assert(in == expected);
+  }
+}
+
+int main(int, char**) {
+  test_char<cpp17_output_iterator<char*> >();
+  test_char<forward_iterator<char*> >();
+  test_char<bidirectional_iterator<char*> >();
+  test_char<random_access_iterator<char*> >();
+  test_char<char*>();
 
-int main(int, char**)
-{
-    test_char<cpp17_output_iterator<char*> >();
-    test_char<forward_iterator<char*> >();
-    test_char<bidirectional_iterator<char*> >();
-    test_char<random_access_iterator<char*> >();
-    test_char<char*>();
+  test_int<cpp17_output_iterator<int*> >();
+  test_int<forward_iterator<int*> >();
+  test_int<bidirectional_iterator<int*> >();
+  test_int<random_access_iterator<int*> >();
+  test_int<int*>();
 
-    test_int<cpp17_output_iterator<int*> >();
-    test_int<forward_iterator<int*> >();
-    test_int<bidirectional_iterator<int*> >();
-    test_int<random_access_iterator<int*> >();
-    test_int<int*>();
+  test_int_array();
+  test_int_array_struct_source();
+  test_struct_array();
 
-    test_int_array();
-    test_int_array_struct_source();
-    test_struct_array();
+  test5();
+  test6();
 
-    test5();
-    test6();
+  test_bititer_with_custom_sized_types();
 
 #if TEST_STD_VER > 17
-    static_assert(test_constexpr());
+  static_assert(test_constexpr());
 #endif
 
   return 0;