cpp-utils/include/ranges/enumerate.hpp

#pragma once
#include "types.hpp"
#include <cstddef>
#include <ranges>
#include <type_traits>
#include <xutility>

namespace util {
namespace rg = std::ranges;

template <class R>
concept simple_view = // exposition only
    rg::view<R> && rg::range<const R> &&
    std::same_as<std::ranges::iterator_t<R>,
                 std::ranges::iterator_t<const R>> &&
    std::same_as<std::ranges::sentinel_t<R>, std::ranges::sentinel_t<const R>>;

template <rg::input_range R>
requires(rg::view<R>) class enumerate_view
    : public rg::view_interface<enumerate_view<R>> {
private:
  R base_ = R();

public:
  template <bool Const> class iterator;
  template <bool Const> class sentinel;

  constexpr enumerate_view() = default;
  constexpr enumerate_view(R base) : base_(std::move(base)) {}

  constexpr auto size() -> rg::range_difference_t<R>
  requires(rg::sized_range<R>) { return rg::size(base_); }
  constexpr auto size() const -> rg::range_difference_t<R>
  requires(rg::sized_range<const R>) { return rg::size(base_); }

  constexpr auto begin() -> iterator<false> {
    return iterator<false>{rg::begin(base_), 0};
  }

  constexpr auto begin() const -> iterator<true> {
    return iterator<true>{rg::begin(base_), 0};
  }

  constexpr auto end() { return sentinel<false>(rg::end(base_)); }

  constexpr auto end() requires(rg::common_range<R> &&rg::sized_range<R>) {
    return iterator<false>{rg::end(base_),
                           static_cast<rg::range_difference_t<R>>(size())};
  }

  constexpr auto end() const requires(rg::range<const R>) {
    return sentinel<true>{rg::end(base_)};
  }

  constexpr auto end() const
      requires(rg::common_range<const R> &&rg::sized_range<R>) {
    return iterator<true>{rg::end(base_),
                          static_cast<rg::range_difference_t<R>>(size())};
  }

  constexpr auto base() const -> R &requires(std::copy_constructible<R>) {
    return base_;
  }
  constexpr auto base() -> R && { return std::move(base_); }
};

template <class R> enumerate_view(R &&) -> enumerate_view<std::views::all_t<R>>;

template <rg::input_range R>
requires(rg::view<R>) template <bool Const> class enumerate_view<R>::iterator {
public:
  using Base = std::conditional_t<Const, const R, R>;
  using iterator_type = rg::iterator_t<Base>;
  using iterator_category =
      typename std::iterator_traits<iterator_type>::iterator_category;
  using iterator_concept = iterator_category;

  using difference_type = typename rg::range_difference_t<Base>;
  using enumerator_type = difference_type;

  struct result {
    const enumerator_type index;
    rg::range_reference_t<Base> value;
  };

  using reference = result;
  using value_type = result;

public:
  enumerator_type index_ = {};
  iterator_type base_ = iterator_type();

  iterator() = default;
  constexpr explicit iterator(iterator_type base, enumerator_type index = 0)
      : base_(std::move(base)), index_(index) {}
  constexpr iterator(iterator<!Const> i) requires Const
      && std::convertible_to<rg::iterator_t<R>, rg::iterator_t<Base>>
      : base_(std::move(i.base_)), index_(std::move(i.index_)) {}

  constexpr rg::iterator_t<Base>
  base() const &requires std::copyable<rg::iterator_t<Base>> {
    return base_;
  }
  constexpr rg::iterator_t<Base> base() && { return std::move(base_); }

  constexpr decltype(auto) operator*() const { return result{index_, *base_}; }

  constexpr auto operator++() -> iterator & {
    ++index_;
    ++base_;
    return *this;
  }

  constexpr auto operator++(int) -> void requires(!rg::forward_range<Base>) {
    ++index_;
    ++base_;
  }

  constexpr auto operator++(int) -> iterator requires(rg::forward_range<Base>) {
    const auto tmp = *this;
    ++index_;
    ++base_;
    return tmp;
  }

  constexpr auto operator--()
      -> iterator &requires(rg::bidirectional_range<Base>) {
    --index_;
    --base_;
    return *this;
  }

  constexpr auto operator--(int) -> iterator
      requires(rg::bidirectional_range<Base>) {
    const auto tmp = *this;
    --index_;
    --base_;
    return tmp;
  }

  constexpr auto operator+=(difference_type n)
      -> iterator &requires(rg::random_access_range<Base>) {
    index_ += n;
    base_ += n;
    return *this;
  }

  constexpr auto operator-=(difference_type n)
      -> iterator &requires(rg::random_access_range<Base>) {
    index_ -= n;
    base_ -= n;
    return *this;
  }

  friend constexpr auto operator==(const iterator &lhs, const iterator &rhs)
      -> bool requires(std::equality_comparable<rg::iterator_t<Base>>) {
    return lhs.base_ == rhs.base_;
  }

  friend constexpr auto operator<(const iterator &lhs, const iterator &rhs)
      -> bool requires(rg::random_access_range<Base>) {
    return lhs.base_ < rhs.base_;
  }
  friend constexpr auto operator>(const iterator &lhs, const iterator &rhs)
      -> bool requires(rg::random_access_range<Base>) {
    return rhs < lhs;
  }
  friend constexpr auto operator<=(const iterator &lhs, const iterator &rhs)
      -> bool requires(rg::random_access_range<Base>) {
    return !(rhs < lhs);
  }
  friend constexpr auto operator>=(const iterator &lhs, const iterator &rhs)
      -> bool requires(rg::random_access_range<Base>) {
    return !(lhs < rhs);
  }

  friend constexpr auto operator<=>(
      const iterator &lhs,
      const iterator &rhs) requires(std::three_way_comparable<iterator_type>) {
    return lhs.base_ <=> rhs.base_;
  }

  friend constexpr auto operator+(const iterator &lhs, difference_type n)
      -> iterator requires(rg::random_access_range<Base>) {
    return iterator{lhs} + n;
  }

  friend constexpr auto operator+(difference_type n, const iterator &rhs)
      -> iterator requires(rg::random_access_range<Base>) {
    return n + rhs;
  }

  friend constexpr auto operator-(const iterator &lhs, difference_type n)
      -> iterator requires(rg::random_access_range<Base>) {
    return iterator{lhs} -= n;
  }

  friend constexpr auto operator-(const iterator &lhs, const iterator &rhs)
      -> iterator requires(rg::random_access_range<Base>) {
    return lhs.base_ - rhs.base_;
  }
};

template <rg::input_range R>
requires(rg::view<R>) template <bool Const> class enumerate_view<R>::sentinel {
private:
  using Base = std::conditional_t<Const, const R, R>;
  rg::sentinel_t<Base> end_ = rg::sentinel_t<Base>();

public:
  sentinel() = default;
  constexpr explicit sentinel(rg::sentinel_t<Base> end) : end_(end) {}
  constexpr sentinel(sentinel<!Const> other) requires Const
      && std::convertible_to<rg::sentinel_t<R>, rg::sentinel_t<Base>>
      : end_(std::move(other.end_)) {}

  constexpr auto base() const -> rg::sentinel_t<Base> { return end_; }

  friend constexpr auto operator==(const iterator<Const> &rhs,
                                   const sentinel &lhs) -> bool {
    return rhs.base_ == lhs.end_;
  }

  friend constexpr auto operator-(const iterator<Const> &rhs,
                                  const sentinel &lhs)
      -> rg::range_difference_t<Base> {
    return rhs.base_ - lhs.end_;
  }
  friend constexpr auto operator-(const sentinel &rhs,
                                  const iterator<Const> &lhs)
      -> rg::range_difference_t<Base> {
    return rhs.end_ - lhs.base_;
  }
};

namespace detail {

struct enumerate_view_adaptor {
  template <typename R> constexpr auto operator()(R &&r) const {
    return enumerate_view{std::move(r)};
  }

  template <rg::range R>
  constexpr friend auto operator|(R &&rng, const enumerate_view_adaptor &) {
    return enumerate_view{std::move(rng)};
  }
};
} // namespace detail

namespace views {
inline detail::enumerate_view_adaptor enumerate;
} // namespace views

} // namespace util