[libcxx-commits] [libcxx] [libc++] Optimize search_n (PR #171389)

[Nikolas Klauser via libcxx-commits]

https://github.com/philnik777 updated https://github.com/llvm/llvm-project/pull/171389

>From 88a5755d2d0b3997edd5360d4ff58522d4a56e18 Mon Sep 17 00:00:00 2001
From: Nikolas Klauser <nikolasklauser at berlin.de>
Date: Mon, 8 Dec 2025 19:03:05 +0100
Subject: [PATCH] [libc++] Optimize search_n

---
 libcxx/include/__algorithm/ranges_search_n.h |  4 +-
 libcxx/include/__algorithm/search_n.h        | 65 ++++++++++----------
 2 files changed, 36 insertions(+), 33 deletions(-)

diff --git a/libcxx/include/__algorithm/ranges_search_n.h b/libcxx/include/__algorithm/ranges_search_n.h
index 81b568c0965fd..746bfcc3d1a8f 100644
--- a/libcxx/include/__algorithm/ranges_search_n.h
+++ b/libcxx/include/__algorithm/ranges_search_n.h
@@ -54,8 +54,8 @@ struct __search_n {
       }
 
       if constexpr (random_access_iterator<_Iter1>) {
-        auto __ret = std::__search_n_random_access_impl<_RangeAlgPolicy>(
-            __first, __last, __count, __value, __pred, __proj, __size);
+        auto __ret =
+            std::__search_n_random_access_impl<_RangeAlgPolicy>(__first, __count, __value, __pred, __proj, __size);
         return {std::move(__ret.first), std::move(__ret.second)};
       }
     }
diff --git a/libcxx/include/__algorithm/search_n.h b/libcxx/include/__algorithm/search_n.h
index 38474e1b2379d..8d961489a46c2 100644
--- a/libcxx/include/__algorithm/search_n.h
+++ b/libcxx/include/__algorithm/search_n.h
@@ -11,7 +11,9 @@
 #define _LIBCPP___ALGORITHM_SEARCH_N_H
 
 #include <__algorithm/comp.h>
+#include <__algorithm/find.h>
 #include <__algorithm/iterator_operations.h>
+#include <__algorithm/unwrap_iter.h>
 #include <__config>
 #include <__functional/identity.h>
 #include <__iterator/advance.h>
@@ -68,43 +70,35 @@ _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 pair<_Iter, _Iter> __search_
   }
 }
 
-template <class _AlgPolicy, class _Pred, class _Iter, class _Sent, class _SizeT, class _Type, class _Proj, class _DiffT>
+template <class _AlgPolicy, class _Pred, class _Iter, class _SizeT, class _Type, class _Proj, class _DiffT>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 std::pair<_Iter, _Iter> __search_n_random_access_impl(
-    _Iter __first, _Sent __last, _SizeT __count, const _Type& __value, _Pred& __pred, _Proj& __proj, _DiffT __size1) {
-  using difference_type = typename iterator_traits<_Iter>::difference_type;
+    _Iter __first, _SizeT __count_in, const _Type& __value, _Pred& __pred, _Proj& __proj, _DiffT __size) {
+  auto __last  = __first + __size;
+  auto __count = static_cast<_DiffT>(__count_in);
+
   if (__count == 0)
     return std::make_pair(__first, __first);
-  if (__size1 < static_cast<_DiffT>(__count)) {
-    _IterOps<_AlgPolicy>::__advance_to(__first, __last);
-    return std::make_pair(__first, __first);
-  }
+  if (__size < __count)
+    return std::make_pair(__last, __last);
 
-  const auto __s = __first + __size1 - difference_type(__count - 1); // Start of pattern match can't go beyond here
+  _Iter __try_match_until = __last - __count;
+  _Iter __match_start     = __first;
+  _Iter __matched_until   = __first;
   while (true) {
-    // Find first element in sequence that matchs __value, with a mininum of loop checks
-    while (true) {
-      if (__first >= __s) { // return __last if no element matches __value
-        _IterOps<_AlgPolicy>::__advance_to(__first, __last);
-        return std::make_pair(__first, __first);
-      }
-      if (std::__invoke(__pred, std::__invoke(__proj, *__first), __value))
-        break;
-      ++__first;
-    }
-    // *__first matches __value_, now match elements after here
-    auto __m = __first;
-    _SizeT __c(0);
+    if (__match_start > __try_match_until)
+      return std::make_pair(__last, __last);
+
+    _Iter __check = __match_start + __count;
+
     while (true) {
-      if (++__c == __count) // If pattern exhausted, __first is the answer (works for 1 element pattern)
-        return std::make_pair(__first, __first + _DiffT(__count));
-      ++__m; // no need to check range on __m because __s guarantees we have enough source
+      if (__check == __matched_until)
+        return std::make_pair(__match_start, __match_start + __count);
 
-      // if there is a mismatch, restart with a new __first
-      if (!std::__invoke(__pred, std::__invoke(__proj, *__m), __value)) {
-        __first = __m;
-        ++__first;
+      if (!std::__invoke(__pred, std::__invoke(__proj, *--__check), __value)) {
+        __matched_until = __match_start + __count;
+        __match_start   = ++__check;
         break;
-      } // else there is a match, check next elements
+      }
     }
   }
 }
@@ -119,7 +113,7 @@ template <class _Iter,
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 pair<_Iter, _Iter>
 __search_n_impl(_Iter __first, _Sent __last, _DiffT __count, const _Type& __value, _Pred& __pred, _Proj& __proj) {
   return std::__search_n_random_access_impl<_ClassicAlgPolicy>(
-      __first, __last, __count, __value, __pred, __proj, __last - __first);
+      __first, __count, __value, __pred, __proj, __last - __first);
 }
 
 template <class _Iter1,
@@ -142,7 +136,16 @@ template <class _ForwardIterator, class _Size, class _Tp, class _BinaryPredicate
   static_assert(
       __is_callable<_BinaryPredicate&, decltype(*__first), const _Tp&>::value, "The comparator has to be callable");
   auto __proj = __identity();
-  return std::__search_n_impl(__first, __last, std::__convert_to_integral(__count), __value, __pred, __proj).first;
+  return std::__rewrap_iter(
+      __first,
+      std::__search_n_impl(
+          std::__unwrap_iter(__first),
+          std::__unwrap_iter(__last),
+          std::__convert_to_integral(__count),
+          __value,
+          __pred,
+          __proj)
+          .first);
 }
 
 template <class _ForwardIterator, class _Size, class _Tp>