iterator.hpp

#pragma once
 
#include <array>
#include <cassert>
#include <concepts>
#include <cstddef>
#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>) {
    // step is always positive for unsigned types — standard ceiling division
    return static_cast<size_t>((end - beg + step - T{1}) / step);
  } else {
    // signed: step may be positive or negative
    return static_cast<size_t>((end - beg + step + (step > T{0} ? T{-1} : T{1})) / step);
  }
}
 
 
// ----------------------------------------------------------------------------
// IndexRange<T, N>  (primary template, N > 1)
// IndexRange<T, 1>  (partial specialization, mirrors original IndexRange<T>)
//
// Forward-declare the primary template so the specialization can reference it.
// ----------------------------------------------------------------------------
 
template <std::integral T, size_t N = 1>
class IndexRange;
 
 
// ============================================================================
// Primary template: N-dimensional index range (N > 1)
//
// Represents the Cartesian product of N 1D IndexRange<T, 1> objects.
//
// Iteration order is row-major (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>
class IndexRange {
 
public:
 
  using index_type = T;
 
  static constexpr size_t rank = N;
 
  // --------------------------------------------------------------------------
  // Construction
  // --------------------------------------------------------------------------
 
  IndexRange() = default;
 
  template <typename... Ranges>
    requires (sizeof...(Ranges) == N) &&
             (std::same_as<std::remove_cvref_t<Ranges>, IndexRange<T, 1>> && ...)
  explicit IndexRange(Ranges&&... ranges)
    : _dims{ std::forward<Ranges>(ranges)... } {}
 
  explicit IndexRange(const std::array<IndexRange<T, 1>, N>& dims);
 
  // --------------------------------------------------------------------------
  // Dimension access
  // --------------------------------------------------------------------------
 
  const IndexRange<T, 1>& dim(size_t d) const { return _dims[d]; }
 
  IndexRange<T, 1>& dim(size_t d) { return _dims[d]; }
 
  const std::array<IndexRange<T, 1>, N>& dims() const { return _dims; }
 
  // --------------------------------------------------------------------------
  // Size queries
  // --------------------------------------------------------------------------
 
  size_t size(size_t d) const { return _dims[d].size(); }
 
  size_t size() const;
 
  std::pair<IndexRange<T, N>, size_t> slice_ceil(size_t flat_beg, size_t chunk_size) const;
 
  std::pair<IndexRange<T, N>, size_t> slice_floor(size_t flat_beg, size_t chunk_size) const;
 
  size_t ceil(size_t chunk_size) const;
 
  size_t floor(size_t chunk_size) const;
 
  private:
 
  std::array<IndexRange<T, 1>, N> _dims;
};
 
// ============================================================================
// Out-of-class definitions — IndexRange<T, N>  (primary template)
// ============================================================================
 
template <std::integral T, size_t N>
IndexRange<T, N>::IndexRange(const std::array<IndexRange<T, 1>, N>& dims) : _dims{dims} {}
 
template <std::integral T, size_t N>
size_t IndexRange<T, N>::size() const {
  size_t total = 1;
  for (size_t d = 0; d < N; ++d) {
    size_t s = _dims[d].size();
    if (s == 0) return d == 0 ? 0 : total;  // outermost zero -> 0, inner zero -> outer product
    total *= s;
  }
  return total;
}
 
template <std::integral T, size_t N>
size_t IndexRange<T, N>::ceil(size_t chunk_size) const {
  // Pass 1: cache all dim sizes (one .size() 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] = _dims[d].size();
    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 IndexRange<T, N>::floor(size_t chunk_size) const {
  // Pass 1: cache all dim sizes (one .size() 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] = _dims[d].size();
    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>
std::pair<IndexRange<T, N>, size_t>
IndexRange<T, N>::slice_ceil(size_t flat_beg, size_t chunk_size) const {
 
  if (chunk_size == 0) {
    return { *this, 0 };
  }
 
  // 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] = _dims[d].size();
    if (dim_sizes[d] == 0) { d_active = d; break; }
  }
 
  if (d_active == 0) return { *this, 0 };
 
  // Pass 2a (backward): decode flat_beg into per-dimension coords.
  // Must complete fully — all coords needed for box construction.
  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 with no branching
  // inside each loop:
  //   [0, grow_dim)      — locked dims (one element each)
  //   [grow_dim]         — the grow dimension
  //   [grow_dim+1, N)    — full inner/inactive extent
  std::array<IndexRange<T, 1>, N> box_dims;
 
  for (size_t d = 0; d < grow_dim; ++d) {
    const T beg  = _dims[d].begin();
    const T step = _dims[d].step_size();
    box_dims[d] = IndexRange<T, 1>(
      beg + static_cast<T>(coords[d]) * step,
      beg + static_cast<T>(coords[d] + 1) * step,
      step
    );
  }
 
  {
    const T beg  = _dims[grow_dim].begin();
    const T step = _dims[grow_dim].step_size();
    box_dims[grow_dim] = IndexRange<T, 1>(
      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 { IndexRange<T, N>(box_dims), steps_to_take * active_inner_vol };
}
 
template <std::integral T, size_t N>
std::pair<IndexRange<T, N>, size_t>
IndexRange<T, N>::slice_floor(size_t flat_beg, size_t chunk_size) const {
 
  if (chunk_size == 0) {
    return { *this, 0 };
  }
 
  // 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] = _dims[d].size();
    if (dim_sizes[d] == 0) { d_active = d; break; }
  }
 
  if (d_active == 0) return { *this, 0 };
 
  // Pass 2a (backward): decode flat_beg into per-dimension coords.
  // Must complete fully — all coords needed for box construction.
  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<IndexRange<T, 1>, N> box_dims;
 
  for (size_t d = 0; d < grow_dim; ++d) {
    const T beg  = _dims[d].begin();
    const T step = _dims[d].step_size();
    box_dims[d] = IndexRange<T, 1>(
      beg + static_cast<T>(coords[d]) * step,
      beg + static_cast<T>(coords[d] + 1) * step,
      step
    );
  }
 
  {
    const T beg  = _dims[grow_dim].begin();
    const T step = _dims[grow_dim].step_size();
    box_dims[grow_dim] = IndexRange<T, 1>(
      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 { IndexRange<T, N>(box_dims), steps_to_take * active_inner_vol };
}
 
 
// ============================================================================
// Partial specialization: IndexRange<T, 1>  (1D case)
//
// This is the original IndexRange<T> class, re-expressed as the N=1
// specialization so that IndexRange<int, 1> and IndexRange<int> (via the
// default argument N=1 on the primary template forward declaration above) both
// refer to the same type.
// ============================================================================
 
template <std::integral T>
class IndexRange<T, 1> {
 
  public:
 
  using index_type = T;
 
  static constexpr size_t rank = 1;
 
  IndexRange() = default;
 
  explicit IndexRange(T beg, T end, T step_size)
    : _beg{beg}, _end{end}, _step_size{step_size} {}
 
  T begin() const { return _beg; }
 
  T end() const { return _end; }
 
  T step_size() const { return _step_size; }
 
  IndexRange<T, 1>& reset(T begin, T end, T step_size) {
    _beg = begin;
    _end = end;
    _step_size = step_size;
    return *this;
  }
 
  IndexRange<T, 1>& begin(T new_begin) { _beg = new_begin; return *this; }
 
  IndexRange<T, 1>& end(T new_end) { _end = new_end; return *this; }
 
  IndexRange<T, 1>& step_size(T new_step_size) { _step_size = new_step_size; return *this; }
 
  size_t size() const;
 
  IndexRange<T, 1> unravel(size_t part_beg, size_t part_end) const;
 
  private:
 
  T _beg;
  T _end;
  T _step_size;
 
};
 
template <std::integral T>
size_t IndexRange<T, 1>::size() const {
  return distance(_beg, _end, _step_size);
}
 
template <std::integral T>
IndexRange<T, 1> IndexRange<T, 1>::unravel(size_t part_beg, size_t part_end) const {
  return IndexRange<T, 1>(
    static_cast<T>(part_beg) * _step_size + _beg,
    static_cast<T>(part_end) * _step_size + _beg,
    _step_size
  );
}
 
// ----------------------------------------------------------------------------
// Deduction guide: IndexRange(beg, end, step) -> IndexRange<T, 1>
//
// Required because IndexRange is now a two-parameter template (T, N).
// Without this guide, CTAD cannot deduce N from a three-argument constructor
// call such as `tf::IndexRange range(0, 10, 2)`.
// ----------------------------------------------------------------------------
 
template <std::integral T>
IndexRange(T, T, T) -> IndexRange<T, 1>;
 
// ==========================================
// traits
// ==========================================
 
template <typename>
constexpr bool is_index_range_v = false;
 
template <typename T, size_t N>
constexpr bool is_index_range_v<IndexRange<T, N>> = true;
 
template <typename R>
concept IndexRangeLike = is_index_range_v<std::decay_t<std::unwrap_ref_decay_t<R>>>;
 
template <typename T>
constexpr bool is_1d_index_range_v = false;
 
template <typename T>
constexpr bool is_1d_index_range_v<IndexRange<T, 1>> = true;
 
template <typename R>
concept IndexRange1DLike = is_1d_index_range_v<std::decay_t<std::unwrap_ref_decay_t<R>>>;
 
template <typename T>
constexpr bool is_md_index_range_v = false;
 
template <typename T, size_t N>
requires (N > 1)
constexpr bool is_md_index_range_v<IndexRange<T, N>> = true;
 
template <typename R>
concept IndexRangeMDLike = is_md_index_range_v<std::decay_t<std::unwrap_ref_decay_t<R>>>;
 
}  // end of namespace tf -----------------------------------------------------