Array.hpp

#ifndef ARRAY_HPP
#define ARRAY_HPP

#include <utils/PugsMacros.hpp>
#include <utils/PugsUtils.hpp>

#include <utils/PugsAssert.hpp>

#include <Kokkos_CopyViews.hpp>
#include <algorithm>

template <typename DataType>
class [[nodiscard]] Array
{
 public:
  using data_type  = DataType;
  using index_type = size_t;

 private:
  Kokkos::View<DataType*> m_values;

  // Allows const version to access our data
  friend Array<std::add_const_t<DataType>>;

 public:
  PUGS_INLINE size_t size() const noexcept
  {
    return m_values.extent(0);
  }

  friend PUGS_INLINE Array<std::remove_const_t<DataType>> copy(const Array<DataType>& source)
  {
    Array<std::remove_const_t<DataType>> image(source.size());
    Kokkos::deep_copy(image.m_values, source.m_values);

    return image;
  }

  template <typename DataType2, typename... RT>
  friend PUGS_INLINE Array<DataType2> encapsulate(const Kokkos::View<DataType2*, RT...>& values);

  PUGS_INLINE DataType& operator[](index_type i) const noexcept(NO_ASSERT)
  {
    Assert(i < m_values.extent(0));
    return m_values[i];
  }

  PUGS_INLINE
  void fill(const DataType& data) const
  {
    static_assert(not std::is_const<DataType>(), "Cannot modify Array of const");

    // could consider to use std::fill
    parallel_for(
      this->size(), PUGS_LAMBDA(index_type i) { m_values[i] = data; });
  }

  template <typename DataType2>
  PUGS_INLINE Array& operator=(const Array<DataType2>& array) noexcept
  {
    // ensures that DataType is the same as source DataType2
    static_assert(std::is_same<std::remove_const_t<DataType>, std::remove_const_t<DataType2>>(),
                  "Cannot assign Array of different type");
    // ensures that const is not lost through copy
    static_assert(((std::is_const<DataType2>() and std::is_const<DataType>()) or not std::is_const<DataType2>()),
                  "Cannot assign Array of const to Array of non-const");
    m_values = array.m_values;
    return *this;
  }

  PUGS_INLINE
  Array& operator=(const Array&) = default;

  PUGS_INLINE
  Array& operator=(Array&&) = default;

  PUGS_INLINE
  Array(size_t size) : m_values("anonymous", size)
  {
    static_assert(not std::is_const<DataType>(), "Cannot allocate Array of const data: only view is "
                                                 "supported");
  }

  PUGS_INLINE
  Array() = default;

  PUGS_INLINE
  Array(const Array&) = default;

  template <typename DataType2>
  PUGS_INLINE Array(const Array<DataType2>& array) noexcept
  {
    this->operator=(array);
  }

  PUGS_INLINE
  Array(Array &&) = default;

  PUGS_INLINE
  ~Array() = default;
};

template <typename DataType, typename... RT>
PUGS_INLINE Array<DataType>
encapsulate(const Kokkos::View<DataType*, RT...>& values)
{
  Array<DataType> array;
  array.m_values = values;
  return array;
}

// map, multimap, unordered_map and stack cannot be copied this way
template <typename Container>
PUGS_INLINE Array<std::remove_const_t<typename Container::value_type>>
convert_to_array(const Container& given_container)
{
  using DataType = typename Container::value_type;
  Array<std::remove_const_t<DataType>> array(given_container.size());
  if (given_container.size() > 0) {
    std::copy(begin(given_container), end(given_container), &(array[0]));
  }
  return array;
}

#endif   // ARRAY_HPP