Implement ranges::search_n

stl/inc/algorithm

@@ -2223,6 +2223,169 @@ _NODISCARD _FwdIt search_n(_ExPo&& _Exec, const _FwdIt _First, const _FwdIt _Las
 }
 #endif // _HAS_CXX17
 
+#ifdef __cpp_lib_concepts
+namespace ranges {
+    // VARIABLE ranges::search_n
+    class _Search_n_fn : private _Not_quite_object {
+    public:
+        using _Not_quite_object::_Not_quite_object;
+
+        // clang-format off
+        template <forward_iterator _It, sentinel_for<_It> _Se, class _Ty, class _Pr = ranges::equal_to,
+            class _Pj = identity>
+            requires indirectly_comparable<_It, const _Ty*, _Pr, _Pj>
+        _NODISCARD constexpr subrange<_It> operator()(_It _First, _Se _Last, const iter_difference_t<_It> _Count,
+            const _Ty& _Val, _Pr _Pred = {}, _Pj _Proj = {}) const {
+            _Adl_verify_range(_First, _Last);
+
+            if (_Count <= 0) {
+                return {_First, _First};
+            }
+
+            auto _UFirst = _Get_unwrapped(_STD move(_First));
+            auto _ULast  = _Get_unwrapped(_STD move(_Last));
+
+            if constexpr (sized_sentinel_for<_Se, _It>) {
+                const auto _Dist = _ULast - _UFirst;
+                auto _UResult = _Search_n_sized(_STD move(_UFirst), _Dist, _Val, _Count, _Pass_fn(_Pred), _Pass_fn(_Proj));
+                return _Rewrap_subrange<subrange<_It>>(_First, _STD move(_UResult));
+            } else {
+                auto _UResult = _Search_n_unsized(
+                    _STD move(_UFirst), _STD move(_ULast), _Val, _Count, _Pass_fn(_Pred), _Pass_fn(_Proj));
+                return _Rewrap_subrange<subrange<_It>>(_First, _STD move(_UResult));
+            }
+        }
+
+        template <forward_range _Rng, class _Ty, class _Pr = ranges::equal_to, class _Pj = identity>
+            requires indirectly_comparable<iterator_t<_Rng>, const _Ty*, _Pr, _Pj>
+        _NODISCARD constexpr borrowed_subrange_t<_Rng> operator()(
+            _Rng&& _Range, const range_difference_t<_Rng> _Count, const _Ty& _Val, _Pr _Pred = {}, _Pj _Proj = {}) const {
+            auto _First = _RANGES begin(_Range);
+
+            if (_Count <= 0) {
+                return {_First, _First};
+            }
+
+            if constexpr (sized_range<_Rng>) {
+                const auto _Dist = _RANGES distance(_Range);
+
+                auto _UResult =
+                    _Search_n_sized(_Get_unwrapped(_First), _Dist, _Val, _Count, _Pass_fn(_Pred), _Pass_fn(_Proj));
+                return _Rewrap_subrange<borrowed_subrange_t<_Rng>>(_First, _STD move(_UResult));
+            } else {
+                auto _UResult = _Search_n_unsized(_Get_unwrapped(_First), _Uend(_Range), _Val,
+                    _Count, _Pass_fn(_Pred), _Pass_fn(_Proj));
+                return _Rewrap_subrange<borrowed_subrange_t<_Rng>>(_First, _STD move(_UResult));
+            }
+        }
+        // clang-format on
+
+    private:
+        template <class _It, class _Ty, class _Pr, class _Pj>
+        _NODISCARD static constexpr subrange<_It> _Search_n_sized(_It _First, iter_difference_t<_It> _Dist,
+            const _Ty& _Val, const iter_difference_t<_It> _Count, _Pr _Pred, _Pj _Proj) {
+            _STL_INTERNAL_STATIC_ASSERT(forward_iterator<_It>);
+            _STL_INTERNAL_STATIC_ASSERT(indirectly_comparable<_It, const _Ty*, _Pr, _Pj>);
+            _STL_INTERNAL_CHECK(_Count > 0);
+            // pre: _First + [0, _Dist) is a valid counted range
+
+            if constexpr (bidirectional_iterator<_It>) {
+                if (_Dist < _Count) {
+                    _RANGES advance(_First, _Dist);
+                    return {_First, _First};
+                }
+
+                auto _Last = _RANGES next(_First, _Count);
+                auto _Mid1 = _First;
+                auto _Mid2 = _Last;
+                for (;;) {
+                    // Invariants: _Last - _First == _Count, [_First, _Mid1) and [_Mid2, _Last) match _Val:
+                    //
+                    // _First  _Mid1    _Mid2    _Last
+                    // |=======|????????|========|??????...
+
+                    --_Mid2;
+                    if (!_STD invoke(_Pred, _STD invoke(_Proj, *_Mid2), _Val)) { // mismatch; skip past it
+                        ++_Mid2;
+                        const auto _Delta = _RANGES distance(_First, _Mid2);
+
+                        if (_Dist - _Delta < _Count) { // not enough space left
+                            _First = _STD move(_Last);
+                            _Dist -= _Count;
+                            break;
+                        }
+
+                        _First = _STD move(_Mid2);
+                        _Dist -= _Delta;
+                        _Mid1 = _Last;
+                        _RANGES advance(_Last, _Delta);
+                        _Mid2 = _Last;
+                        continue;
+                    }
+
+                    if (_Mid2 == _Mid1) { // [_Mid1, _Mid2) is empty, so [_First, _Last) all match
+                        return {_STD move(_First), _STD move(_Last)};
+                    }
+                }
+            } else {
+                for (; _Dist >= _Count; ++_First, (void) --_Dist) {
+                    if (_STD invoke(_Pred, _STD invoke(_Proj, *_First), _Val)) {
+                        auto _Saved = _First;
+                        for (iter_difference_t<_It> _Len = 0;;) {
+                            ++_First;
+                            if (++_Len == _Count) { // match
+                                return {_STD move(_Saved), _STD move(_First)};
+                            }
+
+                            if (!_STD invoke(_Pred, _STD invoke(_Proj, *_First), _Val)) { // mismatch
+                                _Dist -= _Len;
+                                break;
+                            }
+                        }
+                    }
+                }
+            }
+
+            _RANGES advance(_First, _Dist);
+            return {_First, _First};
+        }
+
+        template <class _It, class _Se, class _Ty, class _Pr, class _Pj>
+        _NODISCARD static constexpr subrange<_It> _Search_n_unsized(
+            _It _First, const _Se _Last, const _Ty& _Val, const iter_difference_t<_It> _Count, _Pr _Pred, _Pj _Proj) {
+            _STL_INTERNAL_STATIC_ASSERT(forward_iterator<_It>);
+            _STL_INTERNAL_STATIC_ASSERT(sentinel_for<_Se, _It>);
+            _STL_INTERNAL_STATIC_ASSERT(indirectly_comparable<_It, const _Ty*, _Pr, _Pj>);
+            _STL_INTERNAL_CHECK(_Count > 0);
+
+            for (; _First != _Last; ++_First) {
+                if (_STD invoke(_Pred, _STD invoke(_Proj, *_First), _Val)) {
+                    auto _Saved = _First;
+                    for (auto _Len = _Count;;) {
+                        ++_First;
+                        if (--_Len == 0) { // match
+                            return {_STD move(_Saved), _STD move(_First)};
+                        }
+
+                        if (_First == _Last) { // no more to match against
+                            return {_First, _First};
+                        }
+
+                        if (!_STD invoke(_Pred, _STD invoke(_Proj, *_First), _Val)) { // mismatch
+                            break;
+                        }
+                    }
+                }
+            }
+
+            return {_First, _First};
+        }
+    };
+
+    inline constexpr _Search_n_fn search_n{_Not_quite_object::_Construct_tag{}};
+} // namespace ranges
+#endif // __cpp_lib_concepts
+
 // FUNCTION TEMPLATE find_end
 #if _HAS_IF_CONSTEXPR
 template <class _FwdIt1, class _FwdIt2, class _Pr>

tests/std/test.lst

@@ -257,6 +257,7 @@ tests\P0896R4_ranges_alg_mismatch
 tests\P0896R4_ranges_alg_move
 tests\P0896R4_ranges_alg_none_of
 tests\P0896R4_ranges_alg_search
+tests\P0896R4_ranges_alg_search_n
 tests\P0896R4_ranges_iterator_machinery
 tests\P0896R4_ranges_range_machinery
 tests\P0896R4_ranges_subrange

tests/std/tests/P0896R4_ranges_alg_search_n/env.lst

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+RUNALL_INCLUDE ..\concepts_matrix.lst

tests/std/tests/P0896R4_ranges_alg_search_n/test.cpp

@@ -0,0 +1,135 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <concepts>
+#include <ranges>
+#include <utility>
+
+#include <range_algorithm_support.hpp>
+
+using namespace std;
+
+// Validate dangling story
+STATIC_ASSERT(same_as<decltype(ranges::search_n(borrowed<false>{}, 13, 42)), ranges::dangling>);
+STATIC_ASSERT(same_as<decltype(ranges::search_n(borrowed<true>{}, 42, 13)), ranges::subrange<int*>>);
+
+using P = pair<int, int>;
+
+struct instantiator {
+    static constexpr array<P, 11> pairs = {
+        P{0, 42}, P{1, 42}, P{1, 42}, P{0, 42}, P{1, 42}, P{1, 42}, P{0, 42}, P{1, 42}, P{1, 42}, P{1, 42}, P{0, 42}};
+
+    template <class Fwd>
+    static constexpr void call() {
+        const Fwd range{pairs};
+
+        // defaulted predicate and projections
+        {
+            const auto result = ranges::search_n(range, 2, 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(), 3));
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), 2, 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(), 3));
+        }
+
+        // explicit predicate
+        {
+            const auto result = ranges::search_n(range, 2, P{1, 42}, ranges::equal_to{});
+            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(), 3));
+        }
+        {
+            const auto result =
+                ranges::search_n(ranges::begin(range), ranges::end(range), 2, P{1, 42}, ranges::equal_to{});
+            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(), 3));
+        }
+
+        // explicit predicate and projection
+        constexpr auto cmp = [](auto&& x, auto&& y) { return x == y + 1; };
+        {
+            const auto result = ranges::search_n(range, 3, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == ranges::next(range.begin(), 7));
+            assert(result.end() == ranges::next(range.begin(), 10));
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), 3, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == ranges::next(range.begin(), 7));
+            assert(result.end() == ranges::next(range.begin(), 10));
+        }
+
+        // negative case
+        {
+            const auto result = ranges::search_n(range, 4, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.end());
+            assert(result.end() == range.end());
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), 4, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.end());
+            assert(result.end() == range.end());
+        }
+
+        // trivial case: empty needle
+        {
+            const auto result = ranges::search_n(range, 0, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.begin());
+            assert(result.end() == range.begin());
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), 0, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.begin());
+            assert(result.end() == range.begin());
+        }
+
+        // trivial case: range too small
+        {
+            const auto result = ranges::search_n(range, 99999, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.end());
+            assert(result.end() == range.end());
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), 99999, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.end());
+            assert(result.end() == range.end());
+        }
+
+        // trivial case: negative count
+        {
+            const auto result = ranges::search_n(range, -42, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.begin());
+            assert(result.end() == range.begin());
+        }
+        {
+            const auto result = ranges::search_n(ranges::begin(range), ranges::end(range), -42, 0, cmp, get_first);
+            STATIC_ASSERT(same_as<decltype(result), const ranges::subrange<ranges::iterator_t<Fwd>>>);
+            assert(result.begin() == range.begin());
+            assert(result.end() == range.begin());
+        }
+    }
+};
+
+int main() {
+    using Elem = const P;
+    STATIC_ASSERT((test_fwd<instantiator, Elem>(), true));
+    test_fwd<instantiator, Elem>();
+}