Line data    Source code

       1           1 : // Distributed under the MIT License.
       2             : // See LICENSE.txt for details.
       3             : 
       4             : /// \file
       5             : /// Defines class template Index.
       6             : 
       7             : #pragma once
       8             : 
       9             : #include <array>
      10             : #include <concepts>
      11             : #include <cstddef>
      12             : #include <limits>
      13             : #include <ostream>
      14             : 
      15             : #include "Utilities/ErrorHandling/Assert.hpp"
      16             : #include "Utilities/ForceInline.hpp"
      17             : #include "Utilities/Gsl.hpp"
      18             : #include "Utilities/MakeArray.hpp"
      19             : 
      20             : namespace PUP {
      21             : class er;
      22             : }  // namespace PUP
      23             : 
      24             : /// \ingroup DataStructuresGroup
      25             : /// An integer multi-index.
      26             : ///
      27             : /// \tparam Dim the number of integers in the Index.
      28             : template <size_t Dim>
      29           1 : class Index {
      30             :  public:
      31             :   /// Construct with each element set to the same value.
      32           1 :   explicit Index(const size_t i0 = std::numeric_limits<size_t>::max())
      33             :       : indices_(make_array<Dim>(i0)) {}
      34             : 
      35             :   /// Construct specifying value in each dimension
      36             :   template <std::integral... I>
      37             :     requires(sizeof...(I) == Dim and sizeof...(I) > 1)
      38           1 :   explicit Index(I... i) : indices_(make_array(static_cast<size_t>(i)...)) {}
      39             : 
      40           0 :   explicit Index(std::array<size_t, Dim> i) : indices_(std::move(i)) {}
      41             : 
      42           0 :   size_t operator[](const size_t d) const { return gsl::at(indices_, d); }
      43           0 :   size_t& operator[](const size_t d) { return gsl::at(indices_, d); }
      44             : 
      45           0 :   typename std::array<size_t, Dim>::iterator begin() {
      46             :     return indices_.begin();
      47             :   }
      48           0 :   typename std::array<size_t, Dim>::const_iterator begin() const {
      49             :     return indices_.begin();
      50             :   }
      51             : 
      52           0 :   typename std::array<size_t, Dim>::iterator end() { return indices_.end(); }
      53           0 :   typename std::array<size_t, Dim>::const_iterator end() const {
      54             :     return indices_.end();
      55             :   }
      56             : 
      57           0 :   size_t size() const { return Dim; }
      58             : 
      59             :   /// The product of the indices.
      60             :   /// If Dim = 0, the product is defined as 1.
      61             :   template <size_t N = Dim>
      62             :     requires(N > 0)
      63           1 :   constexpr size_t product() const {
      64             :     return indices_[N - 1] * product<N - 1>();
      65             :   }
      66             :   /// \cond
      67             :   // Specialization for N = 0 to stop recursion
      68             :   template <size_t N = Dim>
      69             :     requires(N == 0)
      70             :   constexpr size_t product() const {
      71             :     return 1;
      72             :   }
      73             :   /// \endcond
      74             : 
      75             :   /// Return a smaller Index with the d-th element removed.
      76             :   ///
      77             :   /// \param d the element to remove.
      78             :   template <size_t N = Dim>
      79             :     requires(N > 0)
      80           1 :   Index<Dim - 1> slice_away(const size_t d) const {
      81             :     ASSERT(d < Dim,
      82             :            "Can't slice dimension " << d << " from an Index<" << Dim << ">");
      83             :     std::array<size_t, Dim - 1> t{};
      84             :     for (size_t i = 0; i < d; ++i) {
      85             :       gsl::at(t, i) = gsl::at(indices_, i);
      86             :     }
      87             :     for (size_t i = d + 1; i < Dim; ++i) {
      88             :       gsl::at(t, i - 1) = gsl::at(indices_, i);
      89             :     }
      90             :     return Index<Dim - 1>(t);
      91             :   }
      92             : 
      93             :   /// \cond
      94             :   // NOLINTNEXTLINE(google-runtime-references)
      95             :   void pup(PUP::er& p);
      96             :   /// \endcond
      97             : 
      98             :   template <size_t N>
      99           0 :   friend std::ostream& operator<<(std::ostream& os,  // NOLINT
     100             :                                   const Index<N>& i);
     101             : 
     102           0 :   const size_t* data() const { return indices_.data(); }
     103           0 :   size_t* data() { return indices_.data(); }
     104             : 
     105           0 :   const std::array<size_t, Dim>& indices() const { return indices_; }
     106             : 
     107             :  private:
     108           0 :   std::array<size_t, Dim> indices_;
     109             : };
     110             : 
     111             : /// \ingroup DataStructuresGroup
     112             : /// Get the collapsed index into a 1D array of the data corresponding to this
     113             : /// Index. Note that the first dimension of the Index varies fastest when
     114             : /// computing the collapsed index.
     115             : template <size_t N>
     116           1 : size_t collapsed_index(const Index<N>& index, const Index<N>& extents);
     117             : 
     118             : /// \ingroup DataStructuresGroup
     119             : /// This is the inverse function of collapsed_index. Given a collapsed (1D)
     120             : /// index and the extents of the array, return the Index corresponding to that
     121             : /// collapsed index. Note that the first dimension of the Index varies fastest
     122             : /// when computing the collapsed index.
     123             : template <size_t Dim>
     124           1 : Index<Dim> expanded_index(size_t index, const Index<Dim>& extents);
     125             : 
     126             : template <size_t N>
     127           0 : std::ostream& operator<<(std::ostream& os, const Index<N>& i);
     128             : 
     129             : /// \cond HIDDEN_SYMBOLS
     130             : #ifdef SPECTRE_DEBUG
     131             : namespace Index_detail {
     132             : template <size_t Dim>
     133             : void collapsed_index_check(const Index<Dim>& index, const Index<Dim>& extents) {
     134             :   for (size_t d = 0; d < Dim; ++d) {
     135             :     ASSERT(index[d] < extents[d], "The requested index in the dimension "
     136             :                                       << d << " with value " << index[d]
     137             :                                       << " exceeds the number of grid "
     138             :                                          "points "
     139             :                                       << extents[d]);
     140             :   }
     141             : }
     142             : }  // namespace Index_detail
     143             : #endif
     144             : 
     145             : // the specializations are in the header file so they can be inlined. We use
     146             : // specializations to avoid having loops since this computation is very
     147             : // straightforward.
     148             : template <>
     149             : SPECTRE_ALWAYS_INLINE size_t collapsed_index(const Index<0>& /*index*/,
     150             :                                              const Index<0>& /*extents*/) {
     151             :   return 0;
     152             : }
     153             : 
     154             : template <>
     155             : SPECTRE_ALWAYS_INLINE size_t collapsed_index(const Index<1>& index,
     156             :                                              const Index<1>& extents) {
     157             :   (void)extents;
     158             : #ifdef SPECTRE_DEBUG
     159             :   Index_detail::collapsed_index_check(index, extents);
     160             : #endif
     161             :   return index[0];
     162             : }
     163             : 
     164             : template <>
     165             : SPECTRE_ALWAYS_INLINE size_t collapsed_index(const Index<2>& index,
     166             :                                              const Index<2>& extents) {
     167             : #ifdef SPECTRE_DEBUG
     168             :   Index_detail::collapsed_index_check(index, extents);
     169             : #endif
     170             :   return index[0] + extents[0] * index[1];
     171             : }
     172             : 
     173             : template <>
     174             : SPECTRE_ALWAYS_INLINE size_t collapsed_index(const Index<3>& index,
     175             :                                              const Index<3>& extents) {
     176             : #ifdef SPECTRE_DEBUG
     177             :   Index_detail::collapsed_index_check(index, extents);
     178             : #endif
     179             :   return index[0] + extents[0] * (index[1] + extents[1] * index[2]);
     180             : }
     181             : 
     182             : template <>
     183             : SPECTRE_ALWAYS_INLINE size_t collapsed_index(const Index<4>& index,
     184             :                                              const Index<4>& extents) {
     185             : #ifdef SPECTRE_DEBUG
     186             :   Index_detail::collapsed_index_check(index, extents);
     187             : #endif
     188             :   return index[0] +
     189             :          extents[0] *
     190             :              (index[1] + extents[1] * (index[2] + extents[2] * index[3]));
     191             : }
     192             : 
     193             : template <>
     194             : SPECTRE_ALWAYS_INLINE Index<1> expanded_index(size_t index,
     195             :                                               const Index<1>& extents) {
     196             :   (void)extents;
     197             :   return Index<1>{index};
     198             : }
     199             : 
     200             : template <>
     201             : SPECTRE_ALWAYS_INLINE Index<2> expanded_index(size_t index,
     202             :                                               const Index<2>& extents) {
     203             :   const size_t ix = index % extents[0];
     204             :   const size_t iy = index / extents[0];
     205             :   return Index<2>{{ix, iy}};
     206             : }
     207             : 
     208             : template <>
     209             : SPECTRE_ALWAYS_INLINE Index<3> expanded_index(size_t index,
     210             :                                               const Index<3>& extents) {
     211             :   const size_t ix = index % extents[0];
     212             :   const size_t iy = (index / extents[0]) % extents[1];
     213             :   const size_t iz = (index / extents[0]) / extents[1];
     214             :   return Index<3>{{ix, iy, iz}};
     215             : }
     216             : 
     217             : template <size_t N>
     218             : bool operator==(const Index<N>& lhs, const Index<N>& rhs);
     219             : 
     220             : template <size_t N>
     221             : bool operator!=(const Index<N>& lhs, const Index<N>& rhs);
     222             : /// \endcond