submdspan

submdspan-test.cpp

#include <cassert>
#include <concepts>
#include <span>
#include <tuple>
#include <utility>
#include <ranges>
#if 1
#include <https://raw.githubusercontent.com/kokkos/mdspan/single-header/mdspan.hpp>
#else
#include <mdspan>
#endif

template <int N>
constexpr auto Int = std::integral_constant<int, N>{};

int main()
{
  int a[60];
  std::ranges::iota(a, 1);

  // indexed, full_extent
  {
    using Ext = std::extents<size_t, 3, 4, 5>;
    std::mdspan<int, Ext> m0{a};

    auto m1 = std::submdspan(m0, 1, std::full_extent, std::full_extent);
    static_assert(m1.rank() == 2);
    static_assert(std::same_as<decltype(m1)::extents_type, std::extents<size_t, 4, 5>>);
    static_assert(std::same_as<decltype(m1)::layout_type, std::layout_right>);

    auto m2 = std::submdspan(m0, 1, 2, std::full_extent);
    static_assert(m2.rank() == 1);
    static_assert(std::same_as<decltype(m2)::extents_type, std::extents<size_t, 5>>);
    static_assert(std::same_as<decltype(m2)::layout_type, std::layout_right>);

    auto m3 = std::submdspan(m0, 1, 2, 3);
    static_assert(m3.rank() == 0);
    static_assert(std::same_as<decltype(m3)::extents_type, std::extents<size_t>>);
    static_assert(std::same_as<decltype(m3)::layout_type, std::layout_right>);

    auto m4 = std::submdspan(m0, std::full_extent, std::full_extent, 0);
    static_assert(m4.rank() == 2);
    static_assert(std::same_as<decltype(m4)::extents_type, std::extents<size_t, 3, 4>>);
    static_assert(std::same_as<decltype(m4)::layout_type, std::layout_stride>);
    assert(m4.mapping().stride(0) == 20 && m4.mapping().stride(1) == 5);
  }
  // pair-like
  {
    using Ext = std::extents<size_t, 6, 10>;
    std::mdspan<int, Ext> m0{a};

    auto m1 = std::submdspan(m0, std::pair{2, 4}, std::tuple{3, 7});
    static_assert(m1.rank() == 2);
    static_assert(std::same_as<decltype(m1)::extents_type, std::dextents<size_t, 2>>);
    assert(m1.extent(0) == 2 && m1.extent(1) == 4);

    auto m2 = std::submdspan(m0, std::pair{Int<2>, Int<4>}, std::tuple{Int<3>, Int<7>});
    static_assert(m2.rank() == 2);
    static_assert(std::same_as<decltype(m2)::extents_type, std::extents<size_t, 2, 4>>);
  }
  // strided_slice
  {
    using Ext = std::extents<size_t, 6, 10>;
    std::mdspan<int, Ext> m0{a};

    auto m1 = std::submdspan(m0, std::strided_slice{0, 6, 5}, std::strided_slice{2, 8, 3});
    static_assert(m1.rank() == 2);
    static_assert(std::same_as<decltype(m1)::extents_type, std::dextents<size_t, 2>>);
    assert(m1.extent(0) == 2 && m1.extent(1) == 3);

    auto m1d = std::submdspan(m0,
        std::strided_slice{.offset=0, .extent=6, .stride=5},
        std::strided_slice{.offset=2, .extent=8, .stride=3});
    assert(m1.extents() == m1d.extents());

    auto m2 = std::submdspan(m0, std::strided_slice{0, Int<6>, Int<5>}, std::strided_slice{2, Int<8>, Int<3>});
    static_assert(m2.rank() == 2);
    static_assert(std::same_as<decltype(m2)::extents_type, std::extents<size_t, 2, 3>>);
  }
}