Use wmemchr in optimizations (#4894)

AlexGuteniev · StephanTLavavej · web-flow · commit f3567b8ae385 · 2024-08-25T10:52:27.000-07:00
Co-authored-by: Stephan T. Lavavej &lt;stl@nuwen.net&gt;
diff --git a/stl/inc/xutility b/stl/inc/xutility
@@ -12,6 +12,7 @@
 #include <climits>
 #include <cstdlib>
 #include <cstring>
+#include <cwchar>
 
 #pragma pack(push, _CRT_PACKING)
 #pragma warning(push, _STL_WARNING_LEVEL)
@@ -6009,17 +6010,27 @@ _NODISCARD _CONSTEXPR20 _InIt _Find_unchecked(_InIt _First, const _InIt _Last, c
                 return _First + (_Result - _First_ptr);
             }
 #else // ^^^ _USE_STD_VECTOR_ALGORITHMS / !_USE_STD_VECTOR_ALGORITHMS vvv
-            if constexpr (sizeof(_Iter_value_t<_InIt>) == 1) {
+            if constexpr (sizeof(_Iter_value_t<_InIt>) <= 2) {
                 const auto _First_ptr = _STD _To_address(_First);
-                const auto _Result    = static_cast<remove_reference_t<_Iter_ref_t<_InIt>>*>(
-                    _CSTD memchr(_First_ptr, static_cast<unsigned char>(_Val), static_cast<size_t>(_Last - _First)));
+                const auto _Count     = static_cast<size_t>(_Last - _First);
+
+                using _Ptr_t = remove_reference_t<_Iter_ref_t<_InIt>>*;
+                _Ptr_t _Result;
+
+                if constexpr (sizeof(_Iter_value_t<_InIt>) == 1) {
+                    _Result = static_cast<_Ptr_t>(_CSTD memchr(_First_ptr, static_cast<unsigned char>(_Val), _Count));
+                } else {
+                    _STL_INTERNAL_STATIC_ASSERT(sizeof(_Iter_value_t<_InIt>) == 2);
+                    _Result = reinterpret_cast<_Ptr_t>(const_cast<wchar_t*>(_CSTD wmemchr(
+                        reinterpret_cast<const wchar_t*>(_First_ptr), static_cast<wchar_t>(_Val), _Count)));
+                }
+
                 if constexpr (is_pointer_v<_InIt>) {
                     return _Result ? _Result : _Last;
                 } else {
                     return _Result ? _First + (_Result - _First_ptr) : _Last;
                 }
             }
-            // TRANSITION, DevCom-1614562: not trying wmemchr
 #endif // ^^^ !_USE_STD_VECTOR_ALGORITHMS ^^^
         }
     }
@@ -6062,14 +6073,12 @@ namespace ranges {
     template <input_iterator _It, sentinel_for<_It> _Se, class _Ty, class _Pj = identity>
         requires indirect_binary_predicate<ranges::equal_to, projected<_It, _Pj>, const _Ty*>
     _NODISCARD constexpr _It _Find_unchecked(_It _First, const _Se _Last, const _Ty& _Val, _Pj _Proj = {}) {
-        // TRANSITION, DevCom-1614562: not trying wmemchr
-        // Only single-byte elements are suitable for unsized optimization
-        constexpr bool _Single_byte_elements = sizeof(_Iter_value_t<_It>) == 1;
-        constexpr bool _Is_sized             = sized_sentinel_for<_Se, _It>;
+        constexpr bool _Elements_are_1_or_2_bytes = sizeof(_Iter_value_t<_It>) <= 2;
+        constexpr bool _Is_sized                  = sized_sentinel_for<_Se, _It>;
 
         if constexpr (_Vector_alg_in_find_is_safe<_It, _Ty>
-                      && (_Single_byte_elements ? _Is_sized || same_as<_Se, unreachable_sentinel_t>
-                                                : _Is_sized && _USE_STD_VECTOR_ALGORITHMS)
+                      && (_Elements_are_1_or_2_bytes ? _Is_sized || same_as<_Se, unreachable_sentinel_t>
+                                                     : _Is_sized && _USE_STD_VECTOR_ALGORITHMS)
                       && same_as<_Pj, identity>) {
             if (!_STD is_constant_evaluated()) {
                 if (!_STD _Could_compare_equal_to_value_type<_It>(_Val)) {
@@ -6087,21 +6096,28 @@ namespace ranges {
                 _Ptr_t _Result;
 #if _USE_STD_VECTOR_ALGORITHMS
                 if constexpr (_Is_sized) {
-                    // When _Is_sized && _Single_byte_elements, prefer this over memchr() for performance
+                    // When _Is_sized && _Elements_are_1_or_2_bytes, prefer this over memchr()/wmemchr() for performance
                     const auto _Last_ptr = _First_ptr + (_Last - _First);
                     _Result              = _STD _Find_vectorized(_First_ptr, _Last_ptr, _Val);
                 } else
 #endif // ^^^ _USE_STD_VECTOR_ALGORITHMS ^^^
                 {
-                    _STL_INTERNAL_STATIC_ASSERT(_Single_byte_elements);
+                    _STL_INTERNAL_STATIC_ASSERT(_Elements_are_1_or_2_bytes);
                     size_t _Count;
                     if constexpr (_Is_sized) {
                         _Count = static_cast<size_t>(_Last - _First);
                     } else {
                         _Count = SIZE_MAX;
                     }
 
-                    _Result = static_cast<_Ptr_t>(_CSTD memchr(_First_ptr, static_cast<unsigned char>(_Val), _Count));
+                    if constexpr (sizeof(_Iter_value_t<_It>) == 1) {
+                        _Result =
+                            static_cast<_Ptr_t>(_CSTD memchr(_First_ptr, static_cast<unsigned char>(_Val), _Count));
+                    } else {
+                        _STL_INTERNAL_STATIC_ASSERT(sizeof(_Iter_value_t<_It>) == 2);
+                        _Result = reinterpret_cast<_Ptr_t>(const_cast<wchar_t*>(_CSTD wmemchr(
+                            reinterpret_cast<const wchar_t*>(_First_ptr), static_cast<wchar_t>(_Val), _Count)));
+                    }
 
                     if constexpr (_Is_sized) {
                         if (_Result == nullptr) {
diff --git a/stl/src/vector_algorithms.cpp b/stl/src/vector_algorithms.cpp
@@ -9,6 +9,7 @@
 #include <__msvc_minmax.hpp>
 #include <cstdint>
 #include <cstring>
+#include <cwchar>
 #include <xtr1common>
 
 #ifndef _M_ARM64EC
@@ -3271,8 +3272,7 @@ const void* __stdcall __std_find_trivial_unsized_1(const void* const _First, con
 
 // TRANSITION, ABI: preserved for binary compatibility
 const void* __stdcall __std_find_trivial_unsized_2(const void* const _First, const uint16_t _Val) noexcept {
-    // TRANSITION, DevCom-1614562: not trying wmemchr
-    return __std_find_trivial_unsized_impl(_First, _Val);
+    return wmemchr(static_cast<const wchar_t*>(_First), static_cast<wchar_t>(_Val), SIZE_MAX);
 }
 
 // TRANSITION, ABI: preserved for binary compatibility
