core/utility/nd_map.hpp

0001 /*
0002  * SPDX-PackageName: "covfie, a part of the ACTS project"
0003  * SPDX-FileCopyrightText: 2022 CERN
0004  * SPDX-License-Identifier: MPL-2.0
0005  */
0006
0007 #pragma once
0008
0009 #include <functional>
0010 #include <utility>
0011
0012 #include <covfie/core/array.hpp>
0013
0014 namespace covfie::utility {
0015 template <typename T, std::size_t N, std::size_t... Ns>
0016 auto tail_impl(
0017     std::index_sequence<Ns...>, [[maybe_unused]] const array::array<T, N> & t
0018 )
0019 {
0020     return array::array<T, N - 1>{t.at(Ns + 1u)...};
0021 }
0022
0023 template <typename T, std::size_t N>
0024 auto tail(const array::array<T, N> & t)
0025 {
0026     return tail_impl(std::make_index_sequence<N - 1u>(), t);
0027 }
0028
0029 template <
0030     typename T,
0031     std::size_t N1,
0032     std::size_t N2,
0033     std::size_t... Is1,
0034     std::size_t... Is2>
0035 array::array<T, N1 + N2>
0036 cat_impl(const array::array<T, N1> & a1, const array::array<T, N2> & a2, std::index_sequence<Is1...>, std::index_sequence<Is2...>)
0037 {
0038     return {a1.at(Is1)..., a2.at(Is2)...};
0039 }
0040
0041 template <typename T, std::size_t N1, std::size_t N2>
0042 array::array<T, N1 + N2>
0043 cat(const array::array<T, N1> & a1, const array::array<T, N2> & a2)
0044 {
0045     return cat_impl(
0046         a1, a2, std::make_index_sequence<N1>(), std::make_index_sequence<N2>()
0047     );
0048 }
0049
0050 template <typename Tuple>
0051 void nd_map(std::function<void(Tuple)> f, Tuple s)
0052 {
0053     if constexpr (Tuple::dimensions == 0u) {
0054         f({});
0055     } else if constexpr (Tuple::dimensions == 1u) {
0056         for (typename Tuple::value_type i =
0057                  static_cast<typename Tuple::value_type>(0);
0058              i < s.at(0);
0059              ++i)
0060         {
0061             f(array::array<typename Tuple::value_type, 1>{i});
0062         }
0063     } else {
0064         using tail_t = decltype(tail(std::declval<Tuple>()));
0065
0066         for (typename Tuple::value_type i =
0067                  static_cast<typename Tuple::value_type>(0);
0068              i < s.at(0);
0069              ++i)
0070         {
0071             nd_map<tail_t>(
0072                 [f, i](tail_t r) {
0073                     f(cat(array::array<typename Tuple::value_type, 1>{i}, r));
0074                 },
0075                 tail(s)
0076             );
0077         }
0078     }
0079 }
0080 }