docs/iterator_8hpp_source.html
| | Taskflow: A General-purpose Task-parallel Programming System |
Loading...
Searching...
No Matches
iterator.hpp
1#pragma once
2
3#include <array>
4#include <cassert>
5#include <concepts>
6#include <cstddef>
7#include <tuple>
8#include <type_traits>
9
10namespace tf {
11
28template <std::integral T>
29constexpr bool is_index_range_invalid(T beg, T end, T step) {
30return ((step == T{0} && beg != end) ||
31 (beg < end && step <= T{0}) || // positive range
32 (beg > end && step >= T{0})); // negative range
33}
34
70template <std::integral T>
71constexpr size_t distance(T beg, T end, T step) {
72if constexpr (std::is_unsigned_v<T>) {
73return end > beg
74 ? static_cast<size_t>((end - beg + step - T{1}) / step)
75 : size_t{0};
76 } else {
77return static_cast<size_t>(
78 std::max(T{0}, (end - beg + step + (step > T{0} ? T{-1} : T{1})) / step)
79 );
80 }
81}
82
83// ============================================================================
84// IndexRanges<T, N>
85//
86// A single class template representing an N-dimensional index range, the
87// Cartesian product of N independent 1D ranges (begin, end, step). Each
88// dimension is stored as a std::tuple<T, T, T>, so dim(d) returns a tuple
89// that can be read or mutated directly, including via structured bindings:
90//
91// auto& [beg, end, step] = ranges.dim(0);
92//
93// Members that only make sense for a single dimension (begin(), end(),
94// step_size(), reset(), unravel()) are gated with requires (N == 1).
95// Members that only make sense for more than one dimension (ceil(), floor(),
96// upper_slice(), lower_slice()) are
97// gated with requires (N \> 1). This keeps
98// everything in one class body — instead of a primary template plus a
99// partial specialization — which Doxygen can parse and cross-reference
100// cleanly.
101//
102// tf::IndexRange<T> is an alias for tf::IndexRanges<T, 1>, the common 1D
103// case.
104//
105// Iteration order for N > 1 is row-major (the last dimension is innermost /
106// fastest), matching the natural loop nesting:
107//
108// for i in dim[0]: // outermost
109// for j in dim[1]:
110// ...
111// for k in dim[N-1]: // innermost
112//
113// Flat index 0 corresponds to (beg[0], beg[1], ..., beg[N-1]).
114// ============================================================================
115
187template <std::integral T, size_t N = 1>
188class IndexRanges {
189
190public:
191
195using index_type = T;
196
200static constexpr size_t rank = N;
201
202// --------------------------------------------------------------------------
203// Construction
204// --------------------------------------------------------------------------
205
213IndexRanges() = default;
214
226explicit IndexRanges(T beg, T end, T step_size) requires (N == 1)
227 : _dims{ std::tuple<T, T, T>{beg, end, step_size} } {}
228
251template <typename... Ranges>
252requires (sizeof...(Ranges) == N) &&
253 (std::same_as<std::decay_t<Ranges>, IndexRanges<T, 1>> && ...)
254explicit IndexRanges(Ranges&&... ranges)
255 : _dims{ ranges.dim(0)... } {}
256
273explicit IndexRanges(const std::array<std::tuple<T, T, T>, N>& dims) : _dims{dims} {}
274
275// --------------------------------------------------------------------------
276// Dimension access
277// --------------------------------------------------------------------------
278
297const std::tuple<T, T, T>& dim(size_t d) const { return _dims[d]; }
298
330 std::tuple<T, T, T>& dim(size_t d) { return _dims[d]; }
331
332// --------------------------------------------------------------------------
333// 1D convenience accessors (only available when N == 1)
334// --------------------------------------------------------------------------
335
346 T begin() const requires (N == 1) { return std::get<0>(_dims[0]); }
347
358 T end() const requires (N == 1) { return std::get<1>(_dims[0]); }
359
370 T step_size() const requires (N == 1) { return std::get<2>(_dims[0]); }
371
396IndexRanges& reset(T beg, T end, T step_size) requires (N == 1);
397
410IndexRanges& begin(T new_begin) requires (N == 1);
411
424IndexRanges& end(T new_end) requires (N == 1);
425
438IndexRanges& step_size(T new_step_size) requires (N == 1);
439
465IndexRanges unravel(size_t part_beg, size_t part_end) const requires (N == 1);
466
467// --------------------------------------------------------------------------
468// Size queries (available for any N)
469// --------------------------------------------------------------------------
470
491size_t size(size_t d) const { return std::apply(distance<T>, _dims[d]); }
492
522
523// --------------------------------------------------------------------------
524// N-dimensional algorithms (only available when N > 1)
525// --------------------------------------------------------------------------
526
561size_t ceil(size_t chunk_size) const requires (N > 1);
562
591size_t floor(size_t chunk_size) const requires (N > 1);
592
619IndexRanges upper_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1);
620
640IndexRanges lower_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1);
641
655IndexRanges empty_box() const requires (N > 1);
656
657private:
658
659 std::array<std::tuple<T, T, T>, N> _dims;
660};
661
662// ============================================================================
663// Out-of-class definitions — 1D convenience accessors and size queries
664// ============================================================================
665
666template <std::integral T, size_t N>
667IndexRanges<T, N>&
668IndexRanges<T, N>::reset(T beg, T end, T step_size) requires (N == 1) {
669 _dims[0] = {beg, end, step_size};
670return *this;
671}
672
673template <std::integral T, size_t N>
675IndexRanges<T, N>::begin(T new_begin) requires (N == 1) {
676 std::get<0>(_dims[0]) = new_begin;
677return *this;
678}
679
680template <std::integral T, size_t N>
682IndexRanges<T, N>::end(T new_end) requires (N == 1) {
683 std::get<1>(_dims[0]) = new_end;
684return *this;
685}
686
687template <std::integral T, size_t N>
689IndexRanges<T, N>::step_size(T new_step_size) requires (N == 1) {
690 std::get<2>(_dims[0]) = new_step_size;
691return *this;
692}
693
694template <std::integral T, size_t N>
696IndexRanges<T, N>::unravel(size_t part_beg, size_t part_end) const requires (N == 1) {
697auto [beg, end, step] = _dims[0];
698return IndexRanges(
699static_cast<T>(part_beg) * step + beg,
700static_cast<T>(part_end) * step + beg,
701 step
702 );
703}
704
705template <std::integral T, size_t N>
706size_t IndexRanges<T, N>::size() const {
707if constexpr (N == 1) {
708return size(0);
709 } else {
710// Compile-time-unrolled recursion: D is a template parameter, so each
711// depth is a distinct instantiation and the recursion fully unrolls
712// (no runtime loop). self is passed explicitly because C++20 lambdas
713// cannot otherwise name themselves for recursion.
714auto compute = [this]<size_t D>(auto& self, size_t total) -> size_t {
715if constexpr (D == N) {
716return total;
717 } else {
718size_t s = this->size(D);
719if (s == 0) return D == 0 ? 0 : total; // outermost zero -> 0, inner zero -> outer product
720return self.template operator()<D + 1>(self, total * s);
721 }
722 };
723return compute.template operator()<0>(compute, 1);
724 }
725}
726
727// ============================================================================
728// Out-of-class definitions — N-dimensional algorithms (N > 1)
729// ============================================================================
730
731template <std::integral T, size_t N>
733IndexRanges<T, N>::empty_box() const requires (N > 1) {
734 std::array<std::tuple<T, T, T>, N> box_dims = _dims;
735 box_dims[0] = std::tuple<T, T, T>{
736 std::get<0>(_dims[0]), std::get<0>(_dims[0]), std::get<2>(_dims[0])
737 };
738return IndexRanges(box_dims);
739}
740
741template <std::integral T, size_t N>
742size_t IndexRanges<T, N>::ceil(size_t chunk_size) const requires (N > 1) {
743// Pass 1: cache all dim sizes (one call each) and find d_active.
744size_t dim_sizes[N];
745size_t d_active = N;
746for (size_t d = 0; d < N; ++d) {
747 dim_sizes[d] = size(d);
748if (dim_sizes[d] == 0) { d_active = d; break; }
749 }
750if (d_active == 0) return 0;
751
752// Pass 2: walk suffix products backward within [0, d_active) only.
753// Return the first suffix product >= chunk_size, capped at size().
754size_t inner_volume = 1;
755if (inner_volume >= chunk_size) return inner_volume;
756for (size_t d = d_active; d-- > 0; ) {
757 inner_volume *= dim_sizes[d];
758if (inner_volume >= chunk_size) return inner_volume;
759 }
760return inner_volume; // == size() of active dims, capped
761}
762
763template <std::integral T, size_t N>
764size_t IndexRanges<T, N>::floor(size_t chunk_size) const requires (N > 1) {
765// Pass 1: cache all dim sizes (one call each) and find d_active.
766size_t dim_sizes[N];
767size_t d_active = N;
768for (size_t d = 0; d < N; ++d) {
769 dim_sizes[d] = size(d);
770if (dim_sizes[d] == 0) { d_active = d; break; }
771 }
772if (d_active == 0) return 0;
773
774// Pass 2: walk suffix products backward within [0, d_active) only.
775// Return the largest suffix product <= chunk_size.
776size_t inner_volume = 1;
777size_t last_fit = 1;
778for (size_t d = d_active; d-- > 0; ) {
779size_t next = inner_volume * dim_sizes[d];
780if (next > chunk_size) break;
781 inner_volume = next;
782 last_fit = inner_volume;
783 }
784return last_fit;
785}
786
787template <std::integral T, size_t N>
789IndexRanges<T, N>::upper_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1) {
790
791if (chunk_size == 0) {
792return empty_box();
793 }
794
795// Pass 1 (forward): cache dim sizes and find d_active — the first zero-size
796// dimension. Only [0, d_active) participates in the flat iteration space;
797// dims [d_active, N) are copied as full extent into the returned box.
798size_t dim_sizes[N];
799size_t d_active = N;
800for (size_t d = 0; d < N; ++d) {
801 dim_sizes[d] = size(d);
802if (dim_sizes[d] == 0) { d_active = d; break; }
803 }
804
805if (d_active == 0) return *this;
806
807// Pass 2a (backward): decode flat_beg into per-dimension coords.
808size_t coords[N] = {};
809size_t temp = flat_beg;
810for (size_t d = d_active; d-- > 0; ) {
811 coords[d] = temp % dim_sizes[d];
812 temp /= dim_sizes[d];
813 }
814
815// Pass 2b (backward, fused with grow_dim search): find grow_dim.
816// Ceil variant: commit grow_dim BEFORE the budget check (round up).
817size_t grow_dim = d_active - 1;
818size_t inner_volume = 1;
819size_t active_inner_vol = 1;
820
821for (size_t d = d_active; d-- > 0; ) {
822// trailing-zeros rule: inner coord non-zero -> can't expand further out
823if (d + 1 < d_active && coords[d + 1] != 0) break;
824
825// ceil: commit BEFORE budget check
826 grow_dim = d;
827 active_inner_vol = inner_volume;
828
829if (inner_volume >= chunk_size) break;
830
831 inner_volume *= dim_sizes[d];
832 }
833
834// Steps along grow_dim: ceil division, at least 1 for forward progress.
835size_t steps_left = dim_sizes[grow_dim] - coords[grow_dim];
836size_t steps_needed = (std::max)(size_t{1}, chunk_size / active_inner_vol);
837size_t steps_to_take = (std::min)(steps_left, steps_needed);
838
839// Pass 3: construct the sub-box in three clean segments:
840// [0, grow_dim) — locked dims (one element each)
841// [grow_dim] — the grow dimension
842// [grow_dim+1, N) — full inner/inactive extent
843 std::array<std::tuple<T, T, T>, N> box_dims;
844
845for (size_t d = 0; d < grow_dim; ++d) {
846const T beg = std::get<0>(_dims[d]);
847const T step = std::get<2>(_dims[d]);
848 box_dims[d] = std::tuple<T, T, T>{
849 beg + static_cast<T>(coords[d]) * step,
850 beg + static_cast<T>(coords[d] + 1) * step,
851 step
852 };
853 }
854
855 {
856const T beg = std::get<0>(_dims[grow_dim]);
857const T step = std::get<2>(_dims[grow_dim]);
858 box_dims[grow_dim] = std::tuple<T, T, T>{
859 beg + static_cast<T>(coords[grow_dim]) * step,
860 beg + static_cast<T>(coords[grow_dim] + steps_to_take) * step,
861 step
862 };
863 }
864
865for (size_t d = grow_dim + 1; d < N; ++d) {
866 box_dims[d] = _dims[d]; // full inner or inactive extent
867 }
868
869return IndexRanges<T, N>(box_dims);
870}
871
872template <std::integral T, size_t N>
874IndexRanges<T, N>::lower_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1) {
875
876if (chunk_size == 0) {
877return empty_box();
878 }
879
880// Pass 1 (forward): cache dim sizes and find d_active.
881size_t dim_sizes[N];
882size_t d_active = N;
883for (size_t d = 0; d < N; ++d) {
884 dim_sizes[d] = size(d);
885if (dim_sizes[d] == 0) { d_active = d; break; }
886 }
887
888if (d_active == 0) return *this;
889
890// Pass 2a (backward): decode flat_beg into per-dimension coords.
891size_t coords[N] = {};
892size_t temp = flat_beg;
893for (size_t d = d_active; d-- > 0; ) {
894 coords[d] = temp % dim_sizes[d];
895 temp /= dim_sizes[d];
896 }
897
898// Pass 2b (backward, fused with grow_dim search): find grow_dim.
899// Floor variant: commit grow_dim AFTER the budget check (round down).
900size_t grow_dim = d_active - 1;
901size_t inner_volume = 1;
902size_t active_inner_vol = 1;
903
904for (size_t d = d_active; d-- > 0; ) {
905// trailing-zeros rule
906if (d + 1 < d_active && coords[d + 1] != 0) break;
907
908// floor: budget check BEFORE committing
909size_t next_vol = inner_volume * dim_sizes[d];
910if (next_vol > chunk_size && inner_volume > 1) break;
911
912 grow_dim = d;
913 active_inner_vol = inner_volume;
914 inner_volume = next_vol;
915 }
916
917// Steps along grow_dim: floor division, at least 1 for forward progress.
918size_t steps_left = dim_sizes[grow_dim] - coords[grow_dim];
919size_t steps_needed = (std::max)(size_t{1}, chunk_size / active_inner_vol);
920size_t steps_to_take = (std::min)(steps_left, steps_needed);
921
922// Pass 3: construct the sub-box in three clean segments.
923 std::array<std::tuple<T, T, T>, N> box_dims;
924
925for (size_t d = 0; d < grow_dim; ++d) {
926const T beg = std::get<0>(_dims[d]);
927const T step = std::get<2>(_dims[d]);
928 box_dims[d] = std::tuple<T, T, T>{
929 beg + static_cast<T>(coords[d]) * step,
930 beg + static_cast<T>(coords[d] + 1) * step,
931 step
932 };
933 }
934
935 {
936const T beg = std::get<0>(_dims[grow_dim]);
937const T step = std::get<2>(_dims[grow_dim]);
938 box_dims[grow_dim] = std::tuple<T, T, T>{
939 beg + static_cast<T>(coords[grow_dim]) * step,
940 beg + static_cast<T>(coords[grow_dim] + steps_to_take) * step,
941 step
942 };
943 }
944
945for (size_t d = grow_dim + 1; d < N; ++d) {
946 box_dims[d] = _dims[d]; // full inner or inactive extent
947 }
948
949return IndexRanges<T, N>(box_dims);
950}
951
952// ----------------------------------------------------------------------------
953// IndexRange<T> — alias for the common 1D case
954// ----------------------------------------------------------------------------
955
970template <std::integral T>
971using IndexRange = IndexRanges<T, 1>;
972
973// ==========================================
974// traits
975// ==========================================
976
981template <typename>
982constexpr bool is_index_ranges_v = false;
983
992template <typename T, size_t N>
993constexpr bool is_index_ranges_v<IndexRanges<T, N>> = true;
994
1003template <typename R>
1004concept IndexRangesLike = is_index_ranges_v<std::decay_t<std::unwrap_ref_decay_t<R>>>;
1005
1016template <typename R>
1017concept IndexRanges1DLike = IndexRangesLike<R> &&
1018 (std::decay_t<std::unwrap_ref_decay_t<R>>::rank == 1);
1019
1031template <typename R>
1032concept IndexRangesMDLike = IndexRangesLike<R> &&
1033 (std::decay_t<std::unwrap_ref_decay_t<R>>::rank > 1);
1034
1035} // end of namespace tf -----------------------------------------------------
class to create an N-dimensional index range of integral indices
Definition iterator.hpp:188
IndexRanges & reset(T beg, T end, T step_size)
updates the range with a new starting index, ending index, and step size (only available when N == 1)
Definition iterator.hpp:668
IndexRanges upper_slice(size_t flat_beg, size_t chunk_size) const
returns the smallest perfectly-aligned hyperbox reachable from flat_beg, rounding up to the next hype...
Definition iterator.hpp:789
T end() const
Definition iterator.hpp:358
const std::tuple< T, T, T > & dim(size_t d) const
Definition iterator.hpp:297
IndexRanges & end(T new_end)
updates the ending index of the range (only available when N == 1)
Definition iterator.hpp:682
size_t ceil(size_t chunk_size) const
returns the smallest hyperplane-aligned chunk size that is >= chunk_size, capped at size() when chunk...
Definition iterator.hpp:742
std::tuple< T, T, T > & dim(size_t d)
returns the (begin, end, step) tuple for dimension d (mutable)
Definition iterator.hpp:330
IndexRanges unravel(size_t part_beg, size_t part_end) const
maps a contiguous index partition back to the corresponding subrange (only available when N == 1)
Definition iterator.hpp:696
IndexRanges & step_size(T new_step_size)
updates the step size of the range (only available when N == 1)
Definition iterator.hpp:689
static constexpr size_t rank
Definition iterator.hpp:200
IndexRanges()=default
constructs an index range without initialization
size_t size(size_t d) const
returns the number of iterations along dimension d
Definition iterator.hpp:491
IndexRanges(Ranges &&... ranges)
constructs an N-D index range from N 1D ranges
Definition iterator.hpp:254
IndexRanges lower_slice(size_t flat_beg, size_t chunk_size) const
returns the largest perfectly-aligned hyperbox reachable from flat_beg whose size does not exceed chu...
Definition iterator.hpp:874
size_t size() const
returns the number of active flat iterations
Definition iterator.hpp:706
T index_type
alias for the index type
Definition iterator.hpp:195
IndexRanges(const std::array< std::tuple< T, T, T >, N > &dims)
constructs an index range from an array of (begin, end, step) tuples
Definition iterator.hpp:273
size_t floor(size_t chunk_size) const
returns the largest hyperplane-aligned chunk size that is <= chunk_size (only available when N > 1)
Definition iterator.hpp:764
IndexRanges empty_box() const
returns a box whose size() is 0 (only available when N > 1)
Definition iterator.hpp:733
IndexRanges(T beg, T end, T step_size)
constructs a 1D index range (only available when N == 1)
Definition iterator.hpp:226
IndexRanges & begin(T new_begin)
updates the starting index of the range (only available when N == 1)
Definition iterator.hpp:675
T begin() const
queries the starting index of the range (only available when N == 1)
Definition iterator.hpp:346
tf::IndexRanges< T, 1 >::step_size
T step_size() const
Definition iterator.hpp:370
concept to check if a type is a tf::IndexRanges<T, 1> (i.e., tf::IndexRange<T>)
Definition iterator.hpp:1017
concept to check if a type is a tf::IndexRanges, regardless of dimensionality
Definition iterator.hpp:1004
concept to check if a type is a tf::IndexRanges<T, N> with rank > 1
Definition iterator.hpp:1032
taskflow namespace
Definition small_vector.hpp:20
IndexRanges< T, 1 > IndexRange
alias for the common 1D case of tf::IndexRanges
Definition iterator.hpp:971
constexpr bool is_index_ranges_v
base type trait to detect if a type is a tf::IndexRanges
Definition iterator.hpp:982
constexpr size_t distance(T beg, T end, T step)
calculates the number of iterations in the given index range
Definition iterator.hpp:71
constexpr bool is_index_range_invalid(T beg, T end, T step)
checks if the given index range is invalid
Definition iterator.hpp:29