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 };
  }
 
  // Find d_active: the index of the first zero-size dimension (or N if none).
  // Only dimensions [0, d_active) participate in the flat iteration space.
  // Dimensions [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 };  // outermost dim is zero — no work
 
  // Decode flat_beg into coords for active dimensions [0, d_active) only.
  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];
  }
 
  // Find grow_dim within [0, d_active): record each dim BEFORE the budget check
  // (ceil: round up to the next boundary). Stop when budget met or trailing-zeros fires.
  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; ) {
    if (d + 1 < d_active && coords[d + 1] != 0) break;  // trailing-zeros rule
 
    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.
  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);
 
  // Construct the sub-box.  Active dims [0, d_active) are locked/grow/inner
  // as usual; dims [d_active, N) are copied as full extent (zero-size and beyond).
  std::array<IndexRange<T, 1>, N> box_dims;
  for (size_t d = 0; d < N; ++d) {
    const auto& dim  = _dims[d];
    const T     step = dim.step_size();
    const T     beg  = dim.begin();
    if (d >= d_active) {
      box_dims[d] = dim;                                           // full extent
    } else if (d < grow_dim) {
      T b = beg + static_cast<T>(coords[d]) * step;
      box_dims[d] = IndexRange<T, 1>(b, b + step, step);          // locked
    } else if (d == grow_dim) {
      box_dims[d] = IndexRange<T, 1>(
        beg + static_cast<T>(coords[d])                 * step,
        beg + static_cast<T>(coords[d] + steps_to_take) * step,
        step
      );
    } else {
      box_dims[d] = dim;                                           // full inner 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 };
  }
 
  // Find d_active: the index of the first zero-size dimension (or N if none).
  // Only dimensions [0, d_active) participate in the flat iteration space.
  // Dimensions [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 };  // outermost dim is zero — no work
 
  // Decode flat_beg into coords for active dimensions [0, d_active) only.
  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];
  }
 
  // Find grow_dim within [0, d_active): commit each dim AFTER the budget check
  // (floor: round down to the previous boundary).
  size_t grow_dim        = d_active - 1;
  size_t active_inner_vol = 1;
  size_t inner_volume    = 1;
 
  for (size_t d = d_active; d-- > 0; ) {
    if (d + 1 < d_active && coords[d + 1] != 0) break;  // trailing-zeros rule
 
    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);
 
  // Construct the sub-box.  Active dims [0, d_active) are locked/grow/inner
  // as usual; dims [d_active, N) are copied as full extent (zero-size and beyond).
  std::array<IndexRange<T, 1>, N> box_dims;
  for (size_t d = 0; d < N; ++d) {
    const auto& dim  = _dims[d];
    const T     step = dim.step_size();
    const T     beg  = dim.begin();
    if (d >= d_active) {
      box_dims[d] = dim;                                           // full extent
    } else if (d < grow_dim) {
      T b = beg + static_cast<T>(coords[d]) * step;
      box_dims[d] = IndexRange<T, 1>(b, b + step, step);          // locked
    } else if (d == grow_dim) {
      box_dims[d] = IndexRange<T, 1>(
        beg + static_cast<T>(coords[d])                 * step,
        beg + static_cast<T>(coords[d] + steps_to_take) * step,
        step
      );
    } else {
      box_dims[d] = dim;                                           // full inner 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 -----------------------------------------------------