Use find for search_n when n=1 (#5346)

AlexGuteniev · StephanTLavavej · web-flow · commit 0a0514cb0ecc · 2025-03-24T16:38:29.000-07:00
Co-authored-by: Stephan T. Lavavej &lt;stl@microsoft.com&gt;
diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt
@@ -125,6 +125,7 @@ add_benchmark(random_integer_generation src/random_integer_generation.cpp)
 add_benchmark(remove src/remove.cpp)
 add_benchmark(replace src/replace.cpp)
 add_benchmark(search src/search.cpp)
+add_benchmark(search_n src/search_n.cpp)
 add_benchmark(std_copy src/std_copy.cpp)
 add_benchmark(sv_equal src/sv_equal.cpp)
 add_benchmark(swap_ranges src/swap_ranges.cpp)
diff --git a/benchmarks/src/search_n.cpp b/benchmarks/src/search_n.cpp
@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <algorithm>
+#include <benchmark/benchmark.h>
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include "skewed_allocator.hpp"
+
+using namespace std;
+
+// NB: This particular algorithm has std and ranges implementations with different perf characteristics!
+
+enum class AlgType { Std, Rng };
+
+template <class T, AlgType Alg>
+void bm(benchmark::State& state) {
+    const auto size = static_cast<size_t>(state.range(0));
+
+    constexpr size_t N = 1;
+
+    constexpr T no_match{'-'};
+    constexpr T match{'*'};
+
+    vector<T, not_highly_aligned_allocator<T>> v(size, no_match);
+
+    fill(v.begin() + v.size() / 2, v.end(), match);
+
+    for (auto _ : state) {
+        if constexpr (Alg == AlgType::Std) {
+            benchmark::DoNotOptimize(search_n(v.begin(), v.end(), N, match));
+        } else if constexpr (Alg == AlgType::Rng) {
+            benchmark::DoNotOptimize(ranges::search_n(v, N, match));
+        }
+    }
+}
+
+void common_args(auto bm) {
+    bm->Arg(3000);
+}
+
+BENCHMARK(bm<uint8_t, AlgType::Std>)->Apply(common_args);
+BENCHMARK(bm<uint8_t, AlgType::Rng>)->Apply(common_args);
+
+BENCHMARK(bm<uint16_t, AlgType::Std>)->Apply(common_args);
+BENCHMARK(bm<uint16_t, AlgType::Rng>)->Apply(common_args);
+
+BENCHMARK(bm<uint32_t, AlgType::Std>)->Apply(common_args);
+BENCHMARK(bm<uint32_t, AlgType::Rng>)->Apply(common_args);
+
+BENCHMARK(bm<uint64_t, AlgType::Std>)->Apply(common_args);
+BENCHMARK(bm<uint64_t, AlgType::Rng>)->Apply(common_args);
+
+BENCHMARK_MAIN();
diff --git a/stl/inc/algorithm b/stl/inc/algorithm
@@ -2262,6 +2262,16 @@ _NODISCARD _CONSTEXPR20 _FwdIt search_n(
         return _First;
     }
 
+    if constexpr (_Is_any_of_v<_Pr,
+#if _HAS_CXX20
+                      _RANGES equal_to,
+#endif
+                      equal_to<>>) {
+        if (_Count == 1) {
+            return _STD find(_First, _Last, _Val);
+        }
+    }
+
     if (static_cast<uintmax_t>(_Count) > static_cast<uintmax_t>(_STD _Max_limit<_Iter_diff_t<_FwdIt>>())) {
         // if the number of _Vals searched for is larger than the longest possible sequence, we can't find it
         return _Last;
@@ -2362,6 +2372,19 @@ namespace ranges {
                 return {_First, _First};
             }
 
+            if constexpr (_Is_any_of_v<_Pr, _STD equal_to<>, _RANGES equal_to>) {
+                if (_Count == 1) {
+                    auto _Res = _RANGES find(_First, _Last, _Val, _Pass_fn(_Proj));
+                    if (_Res != _Last) {
+                        auto _Res_end = _Res;
+                        ++_Res_end;
+                        return {_STD move(_Res), _STD move(_Res_end)};
+                    } else {
+                        return {_Res, _Res};
+                    }
+                }
+            }
+
             auto _UFirst = _RANGES _Unwrap_iter<_Se>(_STD move(_First));
             auto _ULast  = _RANGES _Unwrap_sent<_It>(_STD move(_Last));
 
@@ -2388,6 +2411,20 @@ namespace ranges {
                 return {_First, _First};
             }
 
+            if constexpr (_Is_any_of_v<_Pr, _STD equal_to<>, _RANGES equal_to>) {
+                if (_Count == 1) {
+                    auto _Res  = _RANGES find(_Range, _Val, _Pass_fn(_Proj));
+                    auto _Last = _RANGES end(_Range);
+                    if (_Res != _Last) {
+                        auto _Res_end = _Res;
+                        ++_Res_end;
+                        return {_STD move(_Res), _STD move(_Res_end)};
+                    } else {
+                        return {_Res, _Res};
+                    }
+                }
+            }
+
             if constexpr (sized_range<_Rng>) {
                 const auto _Dist = _RANGES distance(_Range);
 
diff --git a/tests/std/tests/P0896R4_ranges_alg_search_n/test.cpp b/tests/std/tests/P0896R4_ranges_alg_search_n/test.cpp
@@ -98,6 +98,34 @@ struct instantiator {
             assert(result.end() == range.begin());
         }
 
+        // trivial case: unit needle
+        {
+            const auto result = ranges::search_n(range, 1, P{1, 42});
+            static_assert(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == ranges::next(range.begin(), 1));
+            assert(result.end() == ranges::next(range.begin(), 2));
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), 1, P{1, 42});
+            static_assert(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == ranges::next(range.begin(), 1));
+            assert(result.end() == ranges::next(range.begin(), 2));
+        }
+
+        // trivial case: unit needle with predicate
+        {
+            const auto result = ranges::search_n(range, 1, 0, cmp, get_first);
+            static_assert(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == ranges::next(range.begin(), 1));
+            assert(result.end() == ranges::next(range.begin(), 2));
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), 1, 0, cmp, get_first);
+            static_assert(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == ranges::next(range.begin(), 1));
+            assert(result.end() == ranges::next(range.begin(), 2));
+        }
+
         // trivial case: range too small
         {
             const auto result = ranges::search_n(range, 99999, 0, cmp, get_first);
