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
  // --------------------------------------------------------------------------

  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;
 
  // --------------------------------------------------------------------------
  // N-dimensional algorithms (only available when N > 1)
  // --------------------------------------------------------------------------

  size_t ceil(size_t chunk_size) const requires (N > 1);

  size_t floor(size_t chunk_size) const requires (N > 1);

  IndexRanges upper_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1);

  IndexRanges lower_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1);

  IndexRanges empty_box() const requires (N > 1);
 
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);
  }
}
 
// ============================================================================
// Out-of-class definitions — N-dimensional algorithms (N > 1)
// ============================================================================
 
template <std::integral T, size_t N>
IndexRanges<T, N>
IndexRanges<T, N>::empty_box() const requires (N > 1) {
  std::array<std::tuple<T, T, T>, N> box_dims = _dims;
  box_dims[0] = std::tuple<T, T, T>{
    std::get<0>(_dims[0]), std::get<0>(_dims[0]), std::get<2>(_dims[0])
  };
  return IndexRanges(box_dims);
}
 
template <std::integral T, size_t N>
size_t IndexRanges<T, N>::ceil(size_t chunk_size) const requires (N > 1) {
  // Pass 1: cache all dim sizes (one call each) and find d_active.
  size_t dim_sizes[N];
  size_t d_active = N;
  for (size_t d = 0; d < N; ++d) {
    dim_sizes[d] = size(d);
    if (dim_sizes[d] == 0) { d_active = d; break; }
  }
  if (d_active == 0) return 0;
 
  // Pass 2: walk suffix products backward within [0, d_active) only.
  // Return the first suffix product >= chunk_size, capped at size().
  size_t inner_volume = 1;
  if (inner_volume >= chunk_size) return inner_volume;
  for (size_t d = d_active; d-- > 0; ) {
    inner_volume *= dim_sizes[d];
    if (inner_volume >= chunk_size) return inner_volume;
  }
  return inner_volume;  // == size() of active dims, capped
}
 
template <std::integral T, size_t N>
size_t IndexRanges<T, N>::floor(size_t chunk_size) const requires (N > 1) {
  // Pass 1: cache all dim sizes (one call each) and find d_active.
  size_t dim_sizes[N];
  size_t d_active = N;
  for (size_t d = 0; d < N; ++d) {
    dim_sizes[d] = size(d);
    if (dim_sizes[d] == 0) { d_active = d; break; }
  }
  if (d_active == 0) return 0;
 
  // Pass 2: walk suffix products backward within [0, d_active) only.
  // Return the largest suffix product <= chunk_size.
  size_t inner_volume = 1;
  size_t last_fit     = 1;
  for (size_t d = d_active; d-- > 0; ) {
    size_t next = inner_volume * dim_sizes[d];
    if (next > chunk_size) break;
    inner_volume = next;
    last_fit     = inner_volume;
  }
  return last_fit;
}
 
template <std::integral T, size_t N>
IndexRanges<T, N>
IndexRanges<T, N>::upper_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1) {
 
  if (chunk_size == 0) {
    return empty_box();
  }
 
  // Pass 1 (forward): cache dim sizes and find d_active — the first zero-size
  // dimension.  Only [0, d_active) participates in the flat iteration space;
  // dims [d_active, N) are copied as full extent into the returned box.
  size_t dim_sizes[N];
  size_t d_active = N;
  for (size_t d = 0; d < N; ++d) {
    dim_sizes[d] = size(d);
    if (dim_sizes[d] == 0) { d_active = d; break; }
  }
 
  if (d_active == 0) return *this;
 
  // Pass 2a (backward): decode flat_beg into per-dimension coords.
  size_t coords[N] = {};
  size_t temp      = flat_beg;
  for (size_t d = d_active; d-- > 0; ) {
    coords[d] = temp % dim_sizes[d];
    temp     /= dim_sizes[d];
  }
 
  // Pass 2b (backward, fused with grow_dim search): find grow_dim.
  // Ceil variant: commit grow_dim BEFORE the budget check (round up).
  size_t grow_dim         = d_active - 1;
  size_t inner_volume     = 1;
  size_t active_inner_vol = 1;
 
  for (size_t d = d_active; d-- > 0; ) {
    // trailing-zeros rule: inner coord non-zero -> can't expand further out
    if (d + 1 < d_active && coords[d + 1] != 0) break;
 
    // ceil: commit BEFORE budget check
    grow_dim         = d;
    active_inner_vol = inner_volume;
 
    if (inner_volume >= chunk_size) break;
 
    inner_volume *= dim_sizes[d];
  }
 
  // Steps along grow_dim: ceil division, at least 1 for forward progress.
  size_t steps_left    = dim_sizes[grow_dim] - coords[grow_dim];
  size_t steps_needed  = (std::max)(size_t{1}, chunk_size / active_inner_vol);
  size_t steps_to_take = (std::min)(steps_left, steps_needed);
 
  // Pass 3: construct the sub-box in three clean segments:
  //   [0, grow_dim)      — locked dims (one element each)
  //   [grow_dim]         — the grow dimension
  //   [grow_dim+1, N)    — full inner/inactive extent
  std::array<std::tuple<T, T, T>, N> box_dims;
 
  for (size_t d = 0; d < grow_dim; ++d) {
    const T beg  = std::get<0>(_dims[d]);
    const T step = std::get<2>(_dims[d]);
    box_dims[d] = std::tuple<T, T, T>{
      beg + static_cast<T>(coords[d]) * step,
      beg + static_cast<T>(coords[d] + 1) * step,
      step
    };
  }
 
  {
    const T beg  = std::get<0>(_dims[grow_dim]);
    const T step = std::get<2>(_dims[grow_dim]);
    box_dims[grow_dim] = std::tuple<T, T, T>{
      beg + static_cast<T>(coords[grow_dim])                 * step,
      beg + static_cast<T>(coords[grow_dim] + steps_to_take) * step,
      step
    };
  }
 
  for (size_t d = grow_dim + 1; d < N; ++d) {
    box_dims[d] = _dims[d];  // full inner or inactive extent
  }
 
  return IndexRanges<T, N>(box_dims);
}
 
template <std::integral T, size_t N>
IndexRanges<T, N>
IndexRanges<T, N>::lower_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1) {
 
  if (chunk_size == 0) {
    return empty_box();
  }
 
  // Pass 1 (forward): cache dim sizes and find d_active.
  size_t dim_sizes[N];
  size_t d_active = N;
  for (size_t d = 0; d < N; ++d) {
    dim_sizes[d] = size(d);
    if (dim_sizes[d] == 0) { d_active = d; break; }
  }
 
  if (d_active == 0) return *this;
 
  // Pass 2a (backward): decode flat_beg into per-dimension coords.
  size_t coords[N] = {};
  size_t temp      = flat_beg;
  for (size_t d = d_active; d-- > 0; ) {
    coords[d] = temp % dim_sizes[d];
    temp     /= dim_sizes[d];
  }
 
  // Pass 2b (backward, fused with grow_dim search): find grow_dim.
  // Floor variant: commit grow_dim AFTER the budget check (round down).
  size_t grow_dim         = d_active - 1;
  size_t inner_volume     = 1;
  size_t active_inner_vol = 1;
 
  for (size_t d = d_active; d-- > 0; ) {
    // trailing-zeros rule
    if (d + 1 < d_active && coords[d + 1] != 0) break;
 
    // floor: budget check BEFORE committing
    size_t next_vol = inner_volume * dim_sizes[d];
    if (next_vol > chunk_size && inner_volume > 1) break;
 
    grow_dim         = d;
    active_inner_vol = inner_volume;
    inner_volume     = next_vol;
  }
 
  // Steps along grow_dim: floor division, at least 1 for forward progress.
  size_t steps_left    = dim_sizes[grow_dim] - coords[grow_dim];
  size_t steps_needed  = (std::max)(size_t{1}, chunk_size / active_inner_vol);
  size_t steps_to_take = (std::min)(steps_left, steps_needed);
 
  // Pass 3: construct the sub-box in three clean segments.
  std::array<std::tuple<T, T, T>, N> box_dims;
 
  for (size_t d = 0; d < grow_dim; ++d) {
    const T beg  = std::get<0>(_dims[d]);
    const T step = std::get<2>(_dims[d]);
    box_dims[d] = std::tuple<T, T, T>{
      beg + static_cast<T>(coords[d]) * step,
      beg + static_cast<T>(coords[d] + 1) * step,
      step
    };
  }
 
  {
    const T beg  = std::get<0>(_dims[grow_dim]);
    const T step = std::get<2>(_dims[grow_dim]);
    box_dims[grow_dim] = std::tuple<T, T, T>{
      beg + static_cast<T>(coords[grow_dim])                 * step,
      beg + static_cast<T>(coords[grow_dim] + steps_to_take) * step,
      step
    };
  }
 
  for (size_t d = grow_dim + 1; d < N; ++d) {
    box_dims[d] = _dims[d];  // full inner or inactive extent
  }
 
  return IndexRanges<T, N>(box_dims);
}
 
// ----------------------------------------------------------------------------
// 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>>>;

template <typename R>
concept IndexRanges1DLike = IndexRangesLike<R> &&
  (std::decay_t<std::unwrap_ref_decay_t<R>>::rank == 1);

template <typename R>
concept IndexRangesMDLike = IndexRangesLike<R> &&
  (std::decay_t<std::unwrap_ref_decay_t<R>>::rank > 1);
 
}  // end of namespace tf -----------------------------------------------------