for_each.hpp

Go to the documentation of this file.

// Copyright (c) 2018--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#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