for_each.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./permutation.hpp"
#include "../stdutil/array.hpp"
 
#include <array>
#include <concepts>
#include <cstdint>
#include <utility>
 
namespace nda {
 
  namespace detail {
 
    // Get the i-th slowest moving dimension from a given encoded stride order.
    template <int R>
    constexpr int index_from_stride_order(uint64_t stride_order, int i) {
      if (stride_order == 0) return i;                 // default C-order
      auto stride_order_arr = decode<R>(stride_order); // FIXME C++20
      return stride_order_arr[i];
    }
 
    // Get the extent of an array along its i-th dimension.
    template <int I, int R, uint64_t StaticExtents, std::integral Int = long>
    constexpr long get_extent(std::array<Int, R> const &shape) {
      if constexpr (StaticExtents == 0) {
        // dynamic extents
        return shape[I];
      } else {
        // full/partial static extents
        constexpr auto static_extents = decode<R>(StaticExtents); // FIXME C++20
        if constexpr (static_extents[I] == 0)
          return shape[I];
        else
          return static_extents[I];
      }
    }
 
    // Apply a callable object recursively to all possible index values of a given shape.
    template <int I, uint64_t StaticExtents, uint64_t StrideOrder, typename F, size_t R, std::integral Int = long>
    FORCEINLINE void for_each_static_impl(std::array<Int, R> const &shape, std::array<long, R> &idxs, F &f) {
      if constexpr (I == R) {
        // end of recursion
        std::apply(f, idxs);
      } else {
        // get the dimension over which to iterate and its extent
        static constexpr int J = index_from_stride_order<R>(StrideOrder, I);
        const long imax        = get_extent<J, R, StaticExtents>(shape);
 
        // loop over all indices of the current dimension
        for (long i = 0; i < imax; ++i) {
          // recursive call for the next dimension
          for_each_static_impl<I + 1, StaticExtents, StrideOrder>(shape, idxs, f);
          ++idxs[J];
        }
        idxs[J] = 0;
      }
    }
 
  } // namespace detail
 
  template <uint64_t StaticExtents, uint64_t StrideOrder, typename F, auto R, std::integral Int = long>
  FORCEINLINE void for_each_static(std::array<Int, R> const &shape, F &&f) { // NOLINT (we do not want to forward here)
    auto idxs = nda::stdutil::make_initialized_array<R>(0l);
    detail::for_each_static_impl<0, StaticExtents, StrideOrder>(shape, idxs, f);
  }
 
  template <typename F, auto R, std::integral Int = long>
  FORCEINLINE void for_each(std::array<Int, R> const &shape, F &&f) { // NOLINT (we do not want to forward here)
    auto idxs = nda::stdutil::make_initialized_array<R>(0l);
    detail::for_each_static_impl<0, 0, 0>(shape, idxs, f);
  }
 
} // namespace nda