Array.hpp

#ifndef ARRAY_HPP
#define ARRAY_HPP

#include <utils/InvalidData.hpp>
#include <utils/NaNHelper.hpp>
#include <utils/PugsAssert.hpp>
#include <utils/PugsMacros.hpp>
#include <utils/PugsUtils.hpp>
#include <utils/Types.hpp>

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

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

  class [[nodiscard]] UnsafeArrayView
  {
   private:
    DataType* const m_values;
    const size_t m_size;

   public:
    [[nodiscard]] PUGS_INLINE size_t size() const
    {
      return m_size;
    }

    [[nodiscard]] PUGS_INLINE DataType& operator[](size_t i) const
    {
      Assert(i < m_size, "invalid index");
      return m_values[i];
    }

    PUGS_INLINE void fill(const DataType& data) const
    {
      for (size_t i = 0; i < m_size; ++i) {
        m_values[i] = data;
      }
    }

    UnsafeArrayView& operator=(const UnsafeArrayView&) = delete;
    UnsafeArrayView& operator=(UnsafeArrayView&&) = delete;

    UnsafeArrayView(const Array<DataType>& array, index_type begin, index_type size)
      : m_values{&array[begin]}, m_size{size}
    {
      Assert((begin < array.size()) and (begin + size <= array.size()), "required view is not contained in the Array");
    }

    // To try to keep these views close to the initial array one
    // forbids copy constructor and take benefits of C++-17 copy
    // elisions.
    UnsafeArrayView(const UnsafeArrayView&) = delete;

    UnsafeArrayView()  = delete;
    ~UnsafeArrayView() = default;
  };

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

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

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

  [[nodiscard]] 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;
  }

  friend PUGS_INLINE void copy_to(const Array<DataType>& source,
                                  const Array<std::remove_const_t<DataType>>& destination)
  {
    Assert(source.size() == destination.size(), "incompatible Array sizes");
    Kokkos::deep_copy(destination.m_values, source.m_values);
  }

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

  template <typename DataType2>
  PUGS_INLINE typename Array<DataType2>::UnsafeArrayView subArrayView(const Array<DataType2>& array,
                                                                      typename Array<DataType2>::index_type begin,
                                                                      typename Array<DataType2>::index_type size);

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

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

    Kokkos::deep_copy(m_values, 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_v<DataType2> and std::is_const_v<DataType>) or not std::is_const_v<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
  explicit Array(size_t size)
  {
    static_assert(not std::is_const<DataType>(), "Cannot allocate Array of const data: only view is "
                                                 "supported");
    if constexpr (std::is_arithmetic_v<DataType>) {
      m_values = Kokkos::View<DataType*>{Kokkos::view_alloc(Kokkos::WithoutInitializing, "anonymous"), size};
    } else {
      m_values = Kokkos::View<DataType*>{"anonymous", size};
    }

#ifndef NDEBUG
    if constexpr (not std::is_const_v<DataType>) {
      using T = std::decay_t<DataType>;
      if constexpr (std::is_arithmetic_v<T>) {
        this->fill(invalid_data_v<T>);
      } else if constexpr (is_tiny_vector_v<T>) {
        this->fill(T{});
      } else if constexpr (is_tiny_matrix_v<T>) {
        this->fill(T{});
      }
    }
#endif   // NDEBUG
  }

  friend std::ostream& operator<<(std::ostream& os, const Array& x)
  {
    if (x.size() > 0) {
      os << 0 << ':' << NaNHelper(x[0]);
    }
    for (size_t i = 1; i < x.size(); ++i) {
      os << ' ' << i << ':' << NaNHelper(x[i]);
    }
    return os;
  }

  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>
[[nodiscard]] PUGS_INLINE size_t
size(const Array<DataType>& array)
{
  return array.size();
}

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

template <typename DataType>
[[nodiscard]] PUGS_INLINE typename Array<DataType>::UnsafeArrayView
subArrayView(const Array<DataType>& array,
             typename Array<DataType>::index_type begin,
             typename Array<DataType>::index_type size)
{
  return typename Array<DataType>::UnsafeArrayView(array, begin, size);
}

// map, multimap, unordered_map and stack cannot be copied this way
template <typename Container>
[[nodiscard]] 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;
}

template <typename DataType>
[[nodiscard]] std::remove_const_t<DataType>
min(const Array<DataType>& array)
{
  using ArrayType  = Array<DataType>;
  using data_type  = std::remove_const_t<typename ArrayType::data_type>;
  using index_type = typename ArrayType::index_type;

  static_assert(std::is_arithmetic_v<data_type>, "min cannot be called on non-arithmetic data");
  static_assert(not std::is_same_v<data_type, bool>, "min cannot be called on boolean data");

  class ArrayMin
  {
   private:
    const ArrayType m_array;

   public:
    PUGS_INLINE
    operator data_type()
    {
      data_type reduced_value;
      parallel_reduce(m_array.size(), *this, reduced_value);
      return reduced_value;
    }

    PUGS_INLINE
    void
    operator()(const index_type& i, data_type& data) const
    {
      if (m_array[i] < data) {
        data = m_array[i];
      }
    }

    PUGS_INLINE
    void
    join(volatile data_type& dst, const volatile data_type& src) const
    {
      if (src < dst) {
        dst = src;
      }
    }

    PUGS_INLINE
    void
    init(data_type& value) const
    {
      value = std::numeric_limits<data_type>::max();
    }

    PUGS_INLINE
    ArrayMin(const ArrayType& array) : m_array(array)
    {
      ;
    }

    PUGS_INLINE
    ~ArrayMin() = default;
  };

  return ArrayMin(array);
}

template <typename DataType>
[[nodiscard]] std::remove_const_t<DataType>
max(const Array<DataType>& array)
{
  using ArrayType  = Array<DataType>;
  using data_type  = std::remove_const_t<typename ArrayType::data_type>;
  using index_type = typename ArrayType::index_type;

  static_assert(std::is_arithmetic_v<data_type>, "max cannot be called on non-arithmetic data");
  static_assert(not std::is_same_v<data_type, bool>, "max cannot be called on boolean data");

  class ArrayMax
  {
   private:
    const ArrayType m_array;

   public:
    PUGS_INLINE
    operator data_type()
    {
      data_type reduced_value;
      parallel_reduce(m_array.size(), *this, reduced_value);
      return reduced_value;
    }

    PUGS_INLINE
    void
    operator()(const index_type& i, data_type& data) const
    {
      if (m_array[i] > data) {
        data = m_array[i];
      }
    }

    PUGS_INLINE
    void
    join(volatile data_type& dst, const volatile data_type& src) const
    {
      if (src > dst) {
        dst = src;
      }
    }

    PUGS_INLINE
    void
    init(data_type& value) const
    {
      value = std::numeric_limits<data_type>::min();
    }

    PUGS_INLINE
    ArrayMax(const ArrayType& array) : m_array(array)
    {
      ;
    }

    PUGS_INLINE
    ~ArrayMax() = default;
  };

  return ArrayMax(array);
}

template <typename DataType>
[[nodiscard]] std::remove_const_t<DataType>
sum(const Array<DataType>& array)
{
  using ArrayType  = Array<DataType>;
  using data_type  = std::remove_const_t<DataType>;
  using index_type = typename ArrayType::index_type;

  static_assert(not std::is_same_v<data_type, bool>, "sum cannot be called on boolean data");

  class ArraySum
  {
   private:
    const ArrayType m_array;

   public:
    PUGS_INLINE
    operator data_type()
    {
      data_type reduced_value;
      parallel_reduce(m_array.size(), *this, reduced_value);
      return reduced_value;
    }

    PUGS_INLINE
    void
    operator()(const index_type& i, data_type& data) const
    {
      data += m_array[i];
    }

    PUGS_INLINE
    void
    join(volatile data_type& dst, const volatile data_type& src) const
    {
      dst += src;
    }

    PUGS_INLINE
    void
    init(data_type& value) const
    {
      if constexpr (std::is_arithmetic_v<data_type>) {
        value = 0;
      } else {
        static_assert(is_tiny_vector_v<data_type> or is_tiny_matrix_v<data_type>, "invalid data type");
        value = zero;
      }
    }

    PUGS_INLINE
    ArraySum(const ArrayType& array) : m_array(array)
    {
      ;
    }

    PUGS_INLINE
    ~ArraySum() = default;
  };

  return ArraySum(array);
}

#endif   // ARRAY_HPP