array_getter.hpp

// This file is part of the ACTS project.
//
// Copyright (C) 2016 CERN for the benefit of the ACTS project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
#pragma once
 
// Project include(s).
#include "detray/algebra/common/math.hpp"
#include "detray/algebra/concepts.hpp"
#include "detray/definitions/detail/assert.hpp"
#include "detray/definitions/detail/qualifiers.hpp"
 
// System include(s).
#include <cstddef>
#include <type_traits>
 
namespace detray::algebra::array::storage {

struct element_getter {
  template <std::size_t ROWS, std::size_t COLS, concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr scalar_t &operator()(
      array_t<array_t<scalar_t, ROWS>, COLS> &m, std::size_t row,
      std::size_t col) const {
    assert(row < ROWS);
    assert(col < COLS);
    return m[col][row];
  }

  template <std::size_t ROWS, std::size_t COLS, concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr scalar_t operator()(
      const array_t<array_t<scalar_t, ROWS>, COLS> &m, std::size_t row,
      std::size_t col) const {
    assert(row < ROWS);
    assert(col < COLS);
    return m[col][row];
  }

  template <std::size_t N, concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr scalar_t &operator()(
      array_t<array_t<scalar_t, N>, 1> &m, std::size_t row) const {
    assert(row < N);
    return m[0][row];
  }

  template <std::size_t N, concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr scalar_t operator()(
      const array_t<array_t<scalar_t, N>, 1> &m, std::size_t row) const {
    assert(row < N);
    return m[0][row];
  }

  template <std::size_t N, concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr scalar_t &operator()(array_t<scalar_t, N> &m,
                                                    std::size_t row) const {
    assert(row < N);
    return m[row];
  }

  template <std::size_t N, concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr scalar_t operator()(
      const array_t<scalar_t, N> &m, std::size_t row) const {
    assert(row < N);
    return m[row];
  }
 
};  // struct element_getter

template <template <typename, std::size_t> class array_t,
          concepts::scalar scalar_t, std::size_t ROWS, std::size_t COLS>
DETRAY_HOST_DEVICE constexpr scalar_t element(
    const array_t<array_t<scalar_t, ROWS>, COLS> &m, std::size_t row,
    std::size_t col) {
  return element_getter()(m, row, col);
}

template <template <typename, std::size_t> class array_t,
          concepts::scalar scalar_t, std::size_t ROWS, std::size_t COLS>
DETRAY_HOST_DEVICE constexpr scalar_t &element(
    array_t<array_t<scalar_t, ROWS>, COLS> &m, std::size_t row,
    std::size_t col) {
  return element_getter()(m, row, col);
}

template <template <typename, std::size_t> class array_t,
          concepts::scalar scalar_t, std::size_t N>
DETRAY_HOST_DEVICE constexpr scalar_t element(
    const array_t<array_t<scalar_t, N>, 1> &m, std::size_t row) {
  return element_getter()(m, row);
}

template <template <typename, std::size_t> class array_t,
          concepts::scalar scalar_t, std::size_t N>
DETRAY_HOST_DEVICE constexpr scalar_t &element(
    array_t<array_t<scalar_t, N>, 1> &m, std::size_t row) {
  return element_getter()(m, row);
}

template <template <typename, std::size_t> class array_t,
          concepts::scalar scalar_t, std::size_t N>
DETRAY_HOST_DEVICE constexpr scalar_t element(const array_t<scalar_t, N> &v,
                                              std::size_t row) {
  return element_getter()(v, row);
}

template <template <typename, std::size_t> class array_t,
          concepts::scalar scalar_t, std::size_t N>
DETRAY_HOST_DEVICE constexpr scalar_t &element(array_t<scalar_t, N> &v,
                                               std::size_t row) {
  return element_getter()(v, row);
}
 
template <std::size_t I, std::size_t J, concepts::matrix M>
DETRAY_HOST_DEVICE decltype(auto) element(M &matrix) {
  if constexpr (concepts::has_compile_time_2d_access<M>) {
    return matrix.template element<I, J>();
  } else {
    using index_t = detray::traits::index_t<std::decay_t<M>>;
    return element(matrix, static_cast<index_t>(I), static_cast<index_t>(J));
  }
}
 
template <std::size_t I, concepts::vector V>
DETRAY_HOST_DEVICE decltype(auto) element(V &vector) {
  if constexpr (concepts::has_compile_time_1d_access<V>) {
    return vector.template element<I>();
  } else {
    using index_t = detray::traits::index_t<std::decay_t<V>>;
    return element(vector, static_cast<index_t>(I));
  }
}

struct block_getter {
  template <std::size_t ROWS, std::size_t COLS, std::size_t oROWS,
            std::size_t oCOLS, concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr auto operator()(
      const array_t<array_t<scalar_t, oROWS>, oCOLS> &m, std::size_t row,
      std::size_t col) const {
    array_t<array_t<scalar_t, ROWS>, COLS> submatrix{};
 
    for (std::size_t icol = col; icol < col + COLS; ++icol) {
      for (std::size_t irow = row; irow < row + ROWS; ++irow) {
        submatrix[icol - col][irow - row] = m[icol][irow];
      }
    }
    return submatrix;
  }

  template <std::size_t SIZE, std::size_t ROWS, std::size_t COLS,
            concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr array_t<scalar_t, SIZE> vector(
      const array_t<array_t<scalar_t, ROWS>, COLS> &m, std::size_t row,
      std::size_t col) const {
    assert(col < COLS);
    assert(row + SIZE <= ROWS);
 
    array_t<scalar_t, SIZE> subvector{};
 
    for (std::size_t irow = row; irow < row + SIZE; ++irow) {
      subvector[irow - row] = m[col][irow];
    }
 
    return subvector;
  }

  template <std::size_t ROWS, std::size_t COLS, class input_matrix_type,
            concepts::scalar scalar_t,
            template <typename, std::size_t> class array_t>
  DETRAY_HOST_DEVICE constexpr void set(
      input_matrix_type &m, const array_t<array_t<scalar_t, ROWS>, COLS> &b,
      std::size_t row, std::size_t col) const {
    for (std::size_t j = 0u; j < COLS; ++j) {
      for (std::size_t i = 0u; i < ROWS; ++i) {
        m[j + col][i + row] = b[j][i];
      }
    }
  }

  template <std::size_t ROWS, concepts::scalar scalar_t,
            template <typename, std::size_t> class vector_t,
            class input_matrix_type>
  DETRAY_HOST_DEVICE constexpr void set(input_matrix_type &m,
                                        const vector_t<scalar_t, ROWS> &b,
                                        std::size_t row,
                                        std::size_t col) const {
    for (std::size_t i = 0; i < ROWS; ++i) {
      m[col][i + row] = b[i];
    }
  }
 
};  // struct block_getter

template <std::size_t ROWS, std::size_t COLS, class input_matrix_type>
DETRAY_HOST_DEVICE constexpr decltype(auto) block(const input_matrix_type &m,
                                                  std::size_t row,
                                                  std::size_t col) {
  constexpr block_getter block{};
  return block.template operator()<ROWS, COLS>(m, row, col);
}

template <std::size_t SIZE, std::size_t ROWS, std::size_t COLS,
          concepts::scalar scalar_t,
          template <typename, std::size_t> class array_t>
DETRAY_HOST_DEVICE constexpr array_t<scalar_t, SIZE> vector(
    const array_t<array_t<scalar_t, ROWS>, COLS> &m, std::size_t row,
    std::size_t col) {
  constexpr block_getter block{};
  return block.template vector<SIZE>(m, row, col);
}

template <std::size_t ROWS, std::size_t COLS, class input_matrix_type,
          concepts::scalar scalar_t,
          template <typename, std::size_t> class array_t>
DETRAY_HOST_DEVICE constexpr void set_block(
    input_matrix_type &m, const array_t<array_t<scalar_t, ROWS>, COLS> &b,
    std::size_t row, std::size_t col) {
  constexpr block_getter block{};
  block.set(m, b, row, col);
}

template <std::size_t ROWS, concepts::scalar scalar_t,
          template <typename, std::size_t> class vector_t,
          class input_matrix_type>
DETRAY_HOST_DEVICE constexpr void set_block(input_matrix_type &m,
                                            const vector_t<scalar_t, ROWS> &b,
                                            std::size_t row, std::size_t col) {
  constexpr block_getter block{};
  block.set(m, b, row, col);
}
 
}  // namespace detray::algebra::array::storage