iterator.hpp

#pragma once
 
#include <array>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <tuple>
#include <type_traits>
 
namespace tf {

template <std::integral T>
constexpr bool is_index_range_invalid(T beg, T end, T step) {
  return ((step == T{0} && beg != end) ||
          (beg < end && step <= T{0}) ||  // positive range
          (beg > end && step >= T{0}));   // negative range
}

template <std::integral T>
constexpr size_t distance(T beg, T end, T step) {
  if constexpr (std::is_unsigned_v<T>) {
    return end > beg
      ? static_cast<size_t>((end - beg + step - T{1}) / step)
      : size_t{0};
  } else {
    return static_cast<size_t>(
      std::max(T{0}, (end - beg + step + (step > T{0} ? T{-1} : T{1})) / step)
    );
  }
}
 
// ============================================================================
// IndexRanges<T, N>
//
// A single class template representing an N-dimensional index range, the
// Cartesian product of N independent 1D ranges (begin, end, step).  Each
// dimension is stored as a std::tuple<T, T, T>, so dim(d) returns a tuple
// that can be read or mutated directly, including via structured bindings:
//
//   auto& [beg, end, step] = ranges.dim(0);
//
// Members that only make sense for a single dimension (begin(), end(),
// step_size(), reset(), unravel()) are gated with `requires (N == 1)`.
// Members that only make sense for more than one dimension (ceil(), floor(),
// upper_slice(), lower_slice()) are
// gated with `requires (N > 1)`.  This keeps
// everything in one class body — instead of a primary template plus a
// partial specialization — which Doxygen can parse and cross-reference
// cleanly.
//
// tf::IndexRange<T> is an alias for tf::IndexRanges<T, 1>, the common 1D
// case.
//
// Iteration order for N > 1 is row-major (the last dimension is innermost /
// fastest), matching the natural loop nesting:
//
//   for i in dim[0]:        // outermost
//     for j in dim[1]:
//       ...
//         for k in dim[N-1]: // innermost
//
// Flat index 0 corresponds to (beg[0], beg[1], ..., beg[N-1]).
// ============================================================================

template <std::integral T, size_t N = 1>
class IndexRanges {
 
public:

  using index_type = T;

  static constexpr size_t rank = N;
 
  // --------------------------------------------------------------------------
  // Construction
  // --------------------------------------------------------------------------

  IndexRanges() = default;

  explicit IndexRanges(T beg, T end, T step_size) requires (N == 1)
    : _dims{ std::tuple<T, T, T>{beg, end, step_size} } {}

  template <typename... Ranges>
  requires (sizeof...(Ranges) == N) && 
           (std::same_as<std::decay_t<Ranges>, IndexRanges<T, 1>> && ...)
  explicit IndexRanges(Ranges&&... ranges)
    : _dims{ ranges.dim(0)... } {}

  explicit IndexRanges(const std::array<std::tuple<T, T, T>, N>& dims) : _dims{dims} {}
 
  // --------------------------------------------------------------------------
  // Dimension access (available for any N)
  // --------------------------------------------------------------------------

  const std::tuple<T, T, T>& dim(size_t d) const { return _dims[d]; }

  std::tuple<T, T, T>& dim(size_t d) { return _dims[d]; }
 
  // --------------------------------------------------------------------------
  // 1D convenience accessors (only available when N == 1)
  // --------------------------------------------------------------------------

  T begin() const requires (N == 1) { return std::get<0>(_dims[0]); }

  T end() const requires (N == 1) { return std::get<1>(_dims[0]); }

  T step_size() const requires (N == 1) { return std::get<2>(_dims[0]); }

  IndexRanges& reset(T beg, T end, T step_size) requires (N == 1);

  IndexRanges& begin(T new_begin) requires (N == 1);

  IndexRanges& end(T new_end) requires (N == 1);

  IndexRanges& step_size(T new_step_size) requires (N == 1);

  IndexRanges unravel(size_t part_beg, size_t part_end) const requires (N == 1);
 
  // --------------------------------------------------------------------------
  // Size queries (available for any N)
  // --------------------------------------------------------------------------

  size_t size(size_t d) const { return std::apply(distance<T>, _dims[d]); }

  size_t size() const;
 
private:
 
  std::array<std::tuple<T, T, T>, N> _dims;
};
 
// ============================================================================
// Out-of-class definitions — 1D convenience accessors and size queries
// ============================================================================
 
template <std::integral T, size_t N>
IndexRanges<T, N>&
IndexRanges<T, N>::reset(T beg, T end, T step_size) requires (N == 1) {
  _dims[0] = {beg, end, step_size};
  return *this;
}
 
template <std::integral T, size_t N>
IndexRanges<T, N>&
IndexRanges<T, N>::begin(T new_begin) requires (N == 1) {
  std::get<0>(_dims[0]) = new_begin;
  return *this;
}
 
template <std::integral T, size_t N>
IndexRanges<T, N>&
IndexRanges<T, N>::end(T new_end) requires (N == 1) {
  std::get<1>(_dims[0]) = new_end;
  return *this;
}
 
template <std::integral T, size_t N>
IndexRanges<T, N>&
IndexRanges<T, N>::step_size(T new_step_size) requires (N == 1) {
  std::get<2>(_dims[0]) = new_step_size;
  return *this;
}
 
template <std::integral T, size_t N>
IndexRanges<T, N>
IndexRanges<T, N>::unravel(size_t part_beg, size_t part_end) const requires (N == 1) {
  auto [beg, end, step] = _dims[0];
  return IndexRanges(
    static_cast<T>(part_beg) * step + beg,
    static_cast<T>(part_end) * step + beg,
    step
  );
}
 
template <std::integral T, size_t N>
size_t IndexRanges<T, N>::size() const {
  if constexpr (N == 1) {
    return size(0);
  } else {
    // Compile-time-unrolled recursion: D is a template parameter, so each
    // depth is a distinct instantiation and the recursion fully unrolls
    // (no runtime loop). `self` is passed explicitly because C++20 lambdas
    // cannot otherwise name themselves for recursion.
    auto compute = [this]<size_t D>(auto& self, size_t total) -> size_t {
      if constexpr (D == N) {
        return total;
      } else {
        size_t s = this->size(D);
        if (s == 0) return D == 0 ? 0 : total;  // outermost zero -> 0, inner zero -> outer product
        return self.template operator()<D + 1>(self, total * s);
      }
    };
    return compute.template operator()<0>(compute, 1);
  }
}
 
// ----------------------------------------------------------------------------
// IndexRange<T> — alias for the common 1D case
// ----------------------------------------------------------------------------

template <std::integral T>
using IndexRange = IndexRanges<T, 1>;
 
// ==========================================
// traits
// ==========================================

template <typename>
constexpr bool is_index_ranges_v = false;

template <typename T, size_t N>
constexpr bool is_index_ranges_v<IndexRanges<T, N>> = true;

template <typename R>
concept IndexRangesLike = is_index_ranges_v<std::decay_t<std::unwrap_ref_decay_t<R>>>;
 
// ------------------------------------------------------------------------------------------------
// Input Iterator Concept
// ------------------------------------------------------------------------------------------------

template<typename T>
concept InputIteratorLike = std::input_iterator<std::decay_t<std::unwrap_ref_decay_t<T>>>;
 
 
}  // end of namespace tf -----------------------------------------------------