Refactor _Copy_memmove and _Copy_memmove_n (#5046)

CaseyCarter · StephanTLavavej · web-flow · commit 645cac9cad15 · 2024-11-08T09:03:59.000-08:00
Co-authored-by: Stephan T. Lavavej &lt;stl@nuwen.net&gt;
diff --git a/stl/inc/algorithm b/stl/inc/algorithm
@@ -1557,24 +1557,26 @@ namespace ranges {
             requires indirectly_copyable<_It, _Out>
         _STATIC_CALL_OPERATOR constexpr copy_n_result<_It, _Out> operator()(
             _It _First, iter_difference_t<_It> _Count, _Out _Output) _CONST_CALL_OPERATOR {
-            auto _UFirst  = _STD _Get_unwrapped_n(_STD move(_First), _Count);
-            auto _UOutput = _STD _Get_unwrapped_n(_STD move(_Output), _Count);
-            if constexpr (_Iter_copy_cat<decltype(_UFirst), decltype(_UOutput)>::_Bitcopy_assignable) {
-                if (!_STD is_constant_evaluated()) {
-                    _UOutput = _STD _Copy_memmove_n(_UFirst, static_cast<size_t>(_Count), _STD move(_UOutput));
-                    _UFirst += _Count;
-                    _STD _Seek_wrapped(_First, _STD move(_UFirst));
-                    _STD _Seek_wrapped(_Output, _STD move(_UOutput));
-                    return {_STD move(_First), _STD move(_Output)};
+            if (_Count > 0) {
+                auto _UFirst  = _STD _Get_unwrapped_n(_STD move(_First), _Count);
+                auto _UOutput = _STD _Get_unwrapped_n(_STD move(_Output), _Count);
+                if constexpr (_Iter_copy_cat<decltype(_UFirst), decltype(_UOutput)>::_Bitcopy_assignable) {
+                    if (!_STD is_constant_evaluated()) {
+                        _UOutput = _STD _Copy_memmove_n(_UFirst, static_cast<size_t>(_Count), _STD move(_UOutput));
+                        _UFirst += _Count;
+                        _STD _Seek_wrapped(_First, _STD move(_UFirst));
+                        _STD _Seek_wrapped(_Output, _STD move(_UOutput));
+                        return {_STD move(_First), _STD move(_Output)};
+                    }
                 }
-            }
 
-            for (; _Count > 0; ++_UFirst, (void) ++_UOutput, --_Count) {
-                *_UOutput = *_UFirst;
-            }
+                for (; _Count > 0; ++_UFirst, (void) ++_UOutput, --_Count) {
+                    *_UOutput = *_UFirst;
+                }
 
-            _STD _Seek_wrapped(_First, _STD move(_UFirst));
-            _STD _Seek_wrapped(_Output, _STD move(_UOutput));
+                _STD _Seek_wrapped(_First, _STD move(_UFirst));
+                _STD _Seek_wrapped(_Output, _STD move(_UOutput));
+            }
             return {_STD move(_First), _STD move(_Output)};
         }
     };
diff --git a/stl/inc/xutility b/stl/inc/xutility
@@ -4708,32 +4708,40 @@ _CONSTEXPR20 void _Verify_ranges_do_not_overlap(const _Iter1& _First1, const _Se
 #endif // _ITERATOR_DEBUG_LEVEL != 2 ^^^
 }
 
-template <class _CtgIt, class _OutCtgIt>
-_OutCtgIt _Copy_memmove(_CtgIt _First, _CtgIt _Last, _OutCtgIt _Dest) {
-    auto _FirstPtr              = _STD _To_address(_First);
-    auto _LastPtr               = _STD _To_address(_Last);
-    auto _DestPtr               = _STD _To_address(_Dest);
-    const char* const _First_ch = const_cast<const char*>(reinterpret_cast<const volatile char*>(_FirstPtr));
-    const char* const _Last_ch  = const_cast<const char*>(reinterpret_cast<const volatile char*>(_LastPtr));
-    char* const _Dest_ch        = const_cast<char*>(reinterpret_cast<const volatile char*>(_DestPtr));
-    const auto _Count           = static_cast<size_t>(_Last_ch - _First_ch);
-    _CSTD memmove(_Dest_ch, _First_ch, _Count);
+template <class _OutCtgIt>
+_OutCtgIt _Copy_memmove_tail(
+    const char* const _First_ch, const _OutCtgIt _Dest, const size_t _Byte_count, const size_t _Object_count) {
+    _STL_INTERNAL_CHECK(_Byte_count == _Object_count * sizeof(*_Dest));
+    const auto _Dest_ptr = _STD _To_address(_Dest);
+    const auto _Dest_ch  = const_cast<char*>(reinterpret_cast<const volatile char*>(_Dest_ptr));
+    _CSTD memmove(_Dest_ch, _First_ch, _Byte_count);
     if constexpr (is_pointer_v<_OutCtgIt>) {
+        (void) _Object_count;
         // CodeQL [SM02986] This cast is correct: we're bypassing pointer arithmetic for performance.
-        return reinterpret_cast<_OutCtgIt>(_Dest_ch + _Count);
+        return reinterpret_cast<_OutCtgIt>(_Dest_ch + _Byte_count);
     } else {
-        return _Dest + static_cast<_Iter_diff_t<_OutCtgIt>>(_LastPtr - _FirstPtr);
+        return _Dest + static_cast<_Iter_diff_t<_OutCtgIt>>(_Object_count);
     }
 }
 
 template <class _CtgIt, class _OutCtgIt>
-_OutCtgIt _Copy_memmove_n(_CtgIt _First, const size_t _Count, _OutCtgIt _Dest) {
-    const auto _Result = _STD _Copy_memmove(_First, _First + _Count, _Dest);
-    if constexpr (is_pointer_v<_OutCtgIt>) {
-        return _Result;
-    } else { // _Result is unused so the compiler can optimize it away
-        return _Dest + static_cast<_Iter_diff_t<_OutCtgIt>>(_Count);
-    }
+_OutCtgIt _Copy_memmove(_CtgIt _First, _CtgIt _Last, _OutCtgIt _Dest) {
+    _STL_INTERNAL_CHECK(_First <= _Last);
+    const auto _First_ptr    = _STD _To_address(_First);
+    const auto _Last_ptr     = _STD _To_address(_Last);
+    const auto _Object_count = static_cast<size_t>(_Last_ptr - _First_ptr);
+    const auto _First_ch     = const_cast<const char*>(reinterpret_cast<const volatile char*>(_First_ptr));
+    const auto _Last_ch      = const_cast<const char*>(reinterpret_cast<const volatile char*>(_Last_ptr));
+    const auto _Byte_count   = static_cast<size_t>(_Last_ch - _First_ch);
+    return _STD _Copy_memmove_tail(_First_ch, _STD move(_Dest), _Byte_count, _Object_count);
+}
+
+template <class _CtgIt, class _OutCtgIt>
+_OutCtgIt _Copy_memmove_n(_CtgIt _First, const size_t _Object_count, _OutCtgIt _Dest) {
+    const auto _First_ptr  = _STD _To_address(_First);
+    const auto _First_ch   = const_cast<const char*>(reinterpret_cast<const volatile char*>(_First_ptr));
+    const auto _Byte_count = _Object_count * sizeof(*_First_ptr);
+    return _STD _Copy_memmove_tail(_First_ch, _STD move(_Dest), _Byte_count, _Object_count);
 }
 
 template <class _It, bool _RequiresMutable = false>
@@ -4755,9 +4763,8 @@ template <class _InIt, class _SizeTy, class _OutIt>
 _CONSTEXPR20 _OutIt _Copy_n_unchecked4(_InIt _First, _SizeTy _Count, _OutIt _Dest) {
     // copy _First + [0, _Count) to _Dest + [0, _Count), returning _Dest + _Count
     // note: has callers outside the copy family
-#if _HAS_CXX20
     _STL_INTERNAL_STATIC_ASSERT(_Integer_like<_SizeTy>);
-#endif // _HAS_CXX20
+    _STL_INTERNAL_CHECK(_Count >= 0);
 
     if constexpr (_Iter_copy_cat<_InIt, _OutIt>::_Bitcopy_assignable) {
 #if _HAS_CXX20
@@ -5041,18 +5048,18 @@ _FwdIt2 copy_n(_ExPo&&, _FwdIt1 _First, _Diff _Count_raw, _FwdIt2 _Dest) noexcep
 template <class _CtgIt1, class _CtgIt2>
 _CtgIt2 _Copy_backward_memmove(_CtgIt1 _First, _CtgIt1 _Last, _CtgIt2 _Dest) {
     // implement copy_backward-like function as memmove
-    auto _FirstPtr              = _STD _To_address(_First);
-    auto _LastPtr               = _STD _To_address(_Last);
-    auto _DestPtr               = _STD _To_address(_Dest);
-    const char* const _First_ch = const_cast<const char*>(reinterpret_cast<const volatile char*>(_FirstPtr));
-    const char* const _Last_ch  = const_cast<const char*>(reinterpret_cast<const volatile char*>(_LastPtr));
-    char* const _Dest_ch        = const_cast<char*>(reinterpret_cast<const volatile char*>(_DestPtr));
-    const auto _Count           = static_cast<size_t>(_Last_ch - _First_ch);
-    auto _Result                = _CSTD memmove(_Dest_ch - _Count, _First_ch, _Count);
+    const auto _First_ptr = _STD _To_address(_First);
+    const auto _Last_ptr  = _STD _To_address(_Last);
+    const auto _Dest_ptr  = _STD _To_address(_Dest);
+    const auto _First_ch  = const_cast<const char*>(reinterpret_cast<const volatile char*>(_First_ptr));
+    const auto _Last_ch   = const_cast<const char*>(reinterpret_cast<const volatile char*>(_Last_ptr));
+    const auto _Dest_ch   = const_cast<char*>(reinterpret_cast<const volatile char*>(_Dest_ptr));
+    const auto _Count     = static_cast<size_t>(_Last_ch - _First_ch);
+    const auto _Result    = _CSTD memmove(_Dest_ch - _Count, _First_ch, _Count);
     if constexpr (is_pointer_v<_CtgIt2>) {
         return static_cast<_CtgIt2>(_Result);
     } else {
-        return _Dest - static_cast<_Iter_diff_t<_CtgIt2>>(_LastPtr - _FirstPtr);
+        return _Dest - static_cast<_Iter_diff_t<_CtgIt2>>(_Last_ptr - _First_ptr);
     }
 }
 
diff --git a/tests/std/tests/GH_000431_copy_move_family/test.cpp b/tests/std/tests/GH_000431_copy_move_family/test.cpp
@@ -17,6 +17,7 @@
 #endif // _HAS_CXX17
 
 #if _HAS_CXX20
+#include <compare>
 #include <span>
 #endif // _HAS_CXX20
 
@@ -641,6 +642,111 @@ void test_algorithms(CopyFn copy_fn) {
     }
 }
 
+#if _HAS_CXX20
+template <class T = int>
+struct MyIterator { // A contiguous iterator with a weirdly narrow difference type
+    using iterator_concept = contiguous_iterator_tag;
+    using value_type       = T;
+    using difference_type  = short;
+
+    value_type* ptr = nullptr;
+
+    MyIterator() = default;
+    explicit MyIterator(value_type* p) : ptr{p} {}
+
+    value_type& operator*() const {
+        return *ptr;
+    }
+    value_type* operator->() const {
+        return ptr;
+    }
+
+    MyIterator& operator++() {
+        ++ptr;
+        return *this;
+    }
+    MyIterator operator++(int) {
+        auto tmp = *this;
+        ++ptr;
+        return tmp;
+    }
+
+    MyIterator& operator--() {
+        --ptr;
+        return *this;
+    }
+    MyIterator operator--(int) {
+        auto tmp = *this;
+        --ptr;
+        return tmp;
+    }
+
+    auto operator<=>(const MyIterator&) const = default;
+
+    value_type& operator[](difference_type n) const {
+        return ptr[n];
+    }
+
+    MyIterator& operator+=(difference_type n) {
+        ptr += n;
+        return *this;
+    }
+    MyIterator& operator-=(difference_type n) {
+        ptr -= n;
+        return *this;
+    }
+    difference_type operator-(const MyIterator& that) const {
+        return ptr - that.ptr;
+    }
+    MyIterator operator+(difference_type n) const {
+        return MyIterator{ptr + n};
+    }
+    MyIterator operator-(difference_type n) const {
+        return MyIterator{ptr - n};
+    }
+    friend MyIterator operator+(difference_type n, const MyIterator& i) {
+        return i + n;
+    }
+};
+static_assert(contiguous_iterator<MyIterator<int>>);
+
+void test_copy_n_regressions() {
+    const int src = 1729;
+    int x         = 100;
+
+    // _Copy_memmove_n was adding a size_t to an iterator without converting to its difference_type:
+    //  warning C4267: 'argument': conversion from 'size_t' to 'MyIterator::difference_type', possible loss of data
+    copy_n(MyIterator{}, -42, MyIterator{});
+    copy_n(MyIterator{}, 0, MyIterator{});
+    copy_n(MyIterator{&src}, 1, MyIterator{&x});
+    assert(x == 1729);
+    x = 100;
+
+    copy_n(&src, -42, &x);
+    assert(x == 100);
+    copy_n(&src, 0, &x);
+    assert(x == 100);
+    copy_n(&src, 1, &x);
+    assert(x == 1729);
+    x = 100;
+
+    // ranges::copy_n wasn't guarding against negative n when calling the memmove optimization
+    ranges::copy_n(MyIterator{}, -42, MyIterator{});
+    ranges::copy_n(MyIterator{}, 0, MyIterator{});
+    ranges::copy_n(MyIterator{&src}, 1, MyIterator{&x});
+    assert(x == 1729);
+    x = 100;
+
+    ranges::copy_n(&src, -42, &x);
+    assert(x == 100);
+    ranges::copy_n(&src, 0, &x);
+    assert(x == 100);
+    ranges::copy_n(&src, 1, &x);
+    assert(x == 1729);
+    x = 100;
+}
+#endif // _HAS_CXX20
+
 int main() {
     test_algorithms([](auto begin, auto end, auto out) { copy(begin, end, out); });
     test_algorithms([](auto begin, auto end, auto out) { copy_n(begin, distance(begin, end), out); });
@@ -710,5 +816,7 @@ int main() {
     test_algorithms([](auto begin, auto end, auto out) {
         ranges::uninitialized_move_n(begin, distance(begin, end), out, unreachable_sentinel);
     });
+
+    test_copy_n_regressions();
 #endif // _HAS_CXX20
 }
