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>