libcruft-util/cruft/util/ranges/chunk.hpp

#pragma once

#include <ranges>
#include <type_traits>

#include <cruft/util/platform.hpp>
#include <cruft/util/debug/assert.hpp>


namespace cruft::ranges {
    // Largely based on the exposition elements of the specification at
    // https://en.cppreference.com/w/cpp/ranges/chunk_view
    //
    // A lot of the implementation is simplified by assuming we're only working with forward ranges.
    // But we should remove this code when clang-19 gets released.
    //
    // clang#18: remove me when clang gets std::views::chunk
    template <std::ranges::view V>
    requires std::ranges::forward_range<V>
    class chunk_view : public std::ranges::view_interface<chunk_view<V>> {
    private:
        V base_;
        std::ranges::range_difference_t<V> n_;

    public:
        template <bool Const>
        class iterator {
            friend chunk_view;

        private:
            using Parent = std::conditional_t<Const, const chunk_view, chunk_view>;
            using Base = std::conditional_t<Const, const V, V>;

            std::ranges::iterator_t<Base> current_;
            std::ranges::iterator_t<Base> end_;
            std::ranges::range_difference_t<Base> n_;
            std::ranges::range_difference_t<Base> missing_;

        public:
            using iterator_category = std::input_iterator_tag;
            // HACK: assuming we're modelling only forward iterator
            using iterator_concept = std::forward_iterator_tag;

            using value_type = decltype(
                std::views::take(
                    std::ranges::subrange(current_, end_),
                    n_
                )
            );
            using difference_type = std::ranges::range_difference_t<Base>;

            iterator() = default;

            constexpr
            iterator (iterator<!Const> i)
            requires
                    Const and
                    std::convertible_to<std::ranges::iterator_t<V>, std::ranges::iterator_t<Base>> and
                    std::convertible_to<std::ranges::sentinel_t<V>, std::ranges::sentinel_t<Base>>
                : current_ (std::move (i.current_))
                , end_ (std::move (i.end_))
                , n_ (i.n_)
                , missing_ (i.missing_)
            { ; }
        private:
            constexpr
            iterator (
                    Parent* parent,
                    std::ranges::iterator_t<Base> current,
                    std::ranges::range_difference_t<Base> missing = 0
            )   : current_ (current)
                , end_ (std::ranges::end (parent->base_))
                , n_ (parent->n_)
                , missing_ (missing)
            {}

        public:
            constexpr value_type
            operator*() const
            {
                CHECK_NEQ (current_, end_);
                return std::views::take (
                    std::ranges::subrange (current_, end_), n_
                );
            }

            constexpr iterator&
            operator++()
            {
                missing_ = std::ranges::advance(current_, n_, end_);
                return *this;
            }

            constexpr iterator
            operator++(int)
            {
                auto tmp = *this;
                ++*this;
                return tmp;
            }

            constexpr iterator&
            operator--()
                requires std::ranges::bidirectional_range<Base>
            {
                std::ranges::advance(current_, missing_ - n_);
                missing_ = 0;
                return *this;
            }

            constexpr iterator
            operator+= (difference_type x)
                requires std::ranges::random_access_range<Base>
            {
                if (x > 0) {
                    CHECK (std::ranges::distance(current_, end_) > n_ * (x - 1));
                    std::ranges::advance(current_, n_ * (x - 1));
                    missing_ = std::ranges::advance(current_, n_, end_);
                }
                else if (x < 0) {
                    std::ranges::advance(current_, n_ * x + missing_);
                    missing_ = 0;
                }

                return *this;
            }

            constexpr iterator&
            operator-= (difference_type x)
                requires std::ranges::random_access_range<Base>
            {
                return *this += -x;
            }

            friend constexpr
            bool
            operator== (iterator const& x, iterator const &y)
            {
                return x.current_ == y.current_;
            }

            friend constexpr
            bool
            operator== (iterator const& x, std::default_sentinel_t)
            {
                return x.current_ == x.end_;
            }

            friend constexpr
            auto
            operator<=> (iterator const& x, iterator const& y)
                requires
                    std::ranges::random_access_range<Base> &&
                    std::three_way_comparable<std::ranges::iterator_t<Base>>
            {
                return x.current_ <=> y.current_;
            }
        };

    public:
        constexpr explicit
        chunk_view (V base, std::ranges::range_difference_t<V> n)
            : base_ (std::move (base))
            , n_ (n)
        { ; }

        constexpr V
        base () const&
            requires std::copy_constructible<V>
        { return base_; }

        constexpr V base() &&
        { return std::move(base_); }

        constexpr auto
        begin ()
        {
            return iterator<false>(this, std::ranges::begin(base_));
        }

        constexpr auto
        begin () const
            requires std::ranges::forward_range<const V>
        {
            return iterator<true>(
                this,
                std::ranges::begin(base_)
            );
        }

        constexpr auto
        end ()
        {
            if constexpr (std::ranges::common_range<V> and std::ranges::sized_range<V>)
            {
                auto missing = (n_ - std::ranges::distance (base_) % n_) % n_;
                return iterator<false>(this, std::ranges::end (base_), missing);
            }
            else if constexpr (std::ranges::common_range<V> and !std::ranges::bidirectional_range<V>)
                return iterator<false>(this, std::ranges::end(base_));
            else
                return std::default_sentinel;
        }

        constexpr auto
        end () const requires std::ranges::forward_range<const V>
        {
            if constexpr (std::ranges::common_range<const V> and std::ranges::sized_range<const V>) {
                auto missing = (n_ - std::ranges::distance (base_) % n_) % n_;
                return iterator<true>(this, std::ranges::end (base_), missing);
            }
            else if constexpr (std::ranges::common_range<const V> and !std::ranges::bidirectional_range<const V>)
                return iterator<true>(this, std::ranges::end(base_));
            else
                return std::default_sentinel;
        }
    };

    template <class R>
    chunk_view (R&&, std::ranges::range_difference_t<R>) -> chunk_view<std::views::all_t<R>>;


    namespace detail {
        struct chunk_closure_t : std::ranges::range_adaptor_closure<chunk_closure_t> {
            int n;

            template <typename R>
            constexpr auto
            operator() (R &&t) const
            {
                return chunk_view (std::forward<R> (t), n);
            }
        };


        struct chunk_base {
            template <std::ranges::viewable_range R>
            constexpr auto
            operator() ( R&& r, std::ranges::range_difference_t<R> n ) const
            {
                return chunk_view (std::views::all (std::forward<R> (r)), n);
            }


            constexpr auto
            operator() (int n) const
            {
                return chunk_closure_t {.n = n};
            }
        };
    }

    constexpr detail::chunk_base chunk;
}

template< class V >
constexpr bool std::ranges::enable_borrowed_range<cruft::ranges::chunk_view<V>> =
    std::ranges::forward_range<V> && std::ranges::enable_borrowed_range<V>;