Hilbert curve layout for std::mdspan

mdspan_hilbert.cpp

#include <cassert>
#include <bit>
#include <iostream>
#include <utility>
#if 1
#include <https://raw.githubusercontent.com/kokkos/mdspan/single-header/mdspan.hpp>
using std::experimental::mdspan;
using std::experimental::dextents;
#else
#include <mdspan>
using std::mdspan;
using std::dextents;
#endif

struct layout_hilbert {
  template <class Extents>
  class mapping {
  public:
    using extents_type = Extents;
    using index_type = typename extents_type::index_type;
    using rank_type = typename extents_type::rank_type;
    using layout_type = layout_hilbert;

    mapping(const extents_type& e)
      : extents_{e}
    {
      static_assert(extents_type::rank() == 2);
      assert(e.extent(0) == e.extent(1));
      assert(std::has_single_bit(e.extent(0)));
    }
    constexpr const extents_type& extents() const noexcept { return extents_; }
    template <class I0, class I1>
    constexpr index_type operator()(I0 i, I1 j) const noexcept
    {
      const index_type N = std::bit_ceil(extents_.extent(0));
      return xy2d(N, j, i);  // x,y=j,i
    }
    constexpr index_type required_span_size() const noexcept
    {
      const index_type N = std::bit_ceil(extents_.extent(0));
      return N * N;
    }
    static constexpr bool is_always_unique() noexcept { return true; }
    static constexpr bool is_always_exhaustive() noexcept { return true; }
    static constexpr bool is_always_strided() noexcept { return false; }
    static constexpr bool is_unique() noexcept { return true; }
    static constexpr bool is_exhaustive() noexcept { return true; }
    static constexpr bool is_strided() noexcept { return false; }
  private:
    // https://en.wikipedia.org/wiki/Hilbert_curve
    static constexpr index_type xy2d(index_type n, index_type x, index_type y)
    {
      index_type rx, ry, s, d = 0;
      for (s = n / 2; s > 0; s /= 2) {
        rx = (x & s) > 0;
        ry = (y & s) > 0;
        d += s * s * ((3 * rx) ^ ry);
        rot(n, &x, &y, rx, ry);
      }
      return d;
    }
#if 0
    static constexpr void d2xy(index_type n, index_type d, index_type *x, index_type *y)
    {
        index_type rx, ry, s, t = d;
        *x = *y = 0;
        for (s = 1; s < n; s *= 2) {
            rx = 1 & (t / 2);
            ry = 1 & (t ^ rx);
            rot(s, x, y, rx, ry);
            *x += s * rx;
            *y += s * ry;
            t /= 4;
        }
    }
#endif
    static constexpr void rot(index_type n, index_type *x, index_type *y, index_type rx, index_type ry)
    {
      if (ry == 0) {
        if (rx == 1) {
          *x = n-1 - *x;
          *y = n-1 - *y;
        }
        std::swap(*x, *y);
      }
    }
  private:
    extents_type extents_{};
  };
};

int main()
{
  auto dump_mat = [](const auto& mat) {
    std::cout << "N=" << mat.extent(0) << '\n';
    for (size_t i = 0; i < mat.extent(0); i++) {
      for (size_t j = 0; j < mat.extent(1); j++) {
        std::cout << ' ' << mat[i,j];
      }
      std::cout << std::endl;
    }
  };
  auto verify_order = [](const auto& mat, const auto& order) {
    mdspan ordmat{order, mat.extent(0), mat.extent(1)};
    for (size_t i = 0; i < mat.extent(0); i++) {
      for (size_t j = 0; j < mat.extent(1); j++) {
        assert((mat[i,j] == ordmat[i,j]));
      }
    }
  };

  using MatExts = dextents<size_t, 2>;

  {
    constexpr int N = 2;
    const int order[N * N] = {
      1, 4,
      2, 3,
    };
    int arr[N * N];    
    std::iota(std::begin(arr), std::end(arr), 1);
    mdspan mat{arr, layout_hilbert::mapping{MatExts{N, N}}};
    dump_mat(mat);
    verify_order(mat, order);
  }
  {
    constexpr int N = 4;
    const int order[N * N] = {
      1, 2, 15, 16,
      4, 3, 14, 13,
      5, 8,  9, 12,
      6, 7, 10, 11,
    };
    int arr[N * N];
    std::iota(std::begin(arr), std::end(arr), 1);
    mdspan mat{arr, layout_hilbert::mapping{MatExts{N, N}}};
    dump_mat(mat);
    verify_order(mat, order);
  }
  {
    constexpr int N = 8;
    const int order[N * N] = {
       1,  4,  5,  6, 59, 60, 61, 64,
       2,  3,  8,  7, 58, 57, 62, 63,
      15, 14,  9, 10, 55, 56, 51, 50,
      16, 13, 12, 11, 54, 53, 52, 49,
      17, 18, 31, 32, 33, 34, 47, 48,
      20, 19, 30, 29, 36, 35, 46, 45,
      21, 24, 25, 28, 37, 40, 41, 44,
      22, 23, 26, 27, 38, 39, 42, 43,
    };
    int arr[N * N];
    std::iota(std::begin(arr), std::end(arr), 1);
    mdspan mat{arr, layout_hilbert::mapping{MatExts{N, N}}};
    dump_mat(mat);
    verify_order(mat, order);
  }
}