Table.hpp

#ifndef TABLE_HPP
#define TABLE_HPP

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

#include <iostream>

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

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

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

 public:
  class [[nodiscard]] UnsafeRowView
  {
   private:
    // @note for performances reasons it is very important to keep
    // m_table as a reference. Kokkos' reference counting mechanism
    // seems quite expensive
    const Table<DataType>& m_table;
    const size_t m_row;

   public:
    PUGS_INLINE size_t
    size() const noexcept
    {
      return m_table.numberOfColumns();
    }

    PUGS_INLINE
    DataType&
    operator[](size_t i) const
    {
      Assert(i < m_table.numberOfColumns(), "invalid index");
      return m_table(m_row, i);
    }

    PUGS_INLINE void
    fill(const DataType& data) const
    {
      for (size_t i = 0; i < this->size(); ++i) {
        m_table(m_row, i) = data;
      }
    }

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

    UnsafeRowView(const Table<DataType>& table, index_type row) : m_table{table}, m_row{row}
    {
      Assert(row < table.numberOfRows(), "required row view is not contained in the Table");
    }

    friend std::ostream&
    operator<<(std::ostream& os, const UnsafeRowView& 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;
    }

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

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

  class [[nodiscard]] UnsafeTableView
  {
   private:
    // @note for performances reasons it is very important to keep
    // m_table as a reference. Kokkos' reference counting mechanism
    // seems quite expensive
    const Table<DataType>& m_table;
    const size_t m_row_begin;
    const size_t m_column_begin;
    const size_t m_row_size;
    const size_t m_column_size;

   public:
    class [[nodiscard]] RowView
    {
     private:
      const UnsafeTableView& m_table_view;
      const size_t m_row;

     public:
      [[nodiscard]] PUGS_INLINE size_t
      size() const
      {
        return m_table_view.numberOfColumns();
      }

      [[nodiscard]] PUGS_INLINE DataType&
      operator[](size_t i) const
      {
        Assert(i < m_table_view.numberOfColumns(), "invalid index");
        return m_table_view(m_row, i);
      }

      PUGS_INLINE void
      fill(const DataType& data) const
      {
        for (size_t i = 0; i < this->size(); ++i) {
          m_table_view(m_row, i) = data;
        }
      }

      friend std::ostream&
      operator<<(std::ostream& os, const RowView& 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;
      }

      RowView(const UnsafeTableView& table_view, index_type row) : m_table_view{table_view}, m_row{row}
      {
        Assert(row < m_table_view.numberOfRows(), "required row view is not contained in the Table view");
      }

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

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

    [[nodiscard]] PUGS_INLINE size_t
    numberOfRows() const noexcept
    {
      return m_row_size;
    }

    [[nodiscard]] PUGS_INLINE size_t
    numberOfColumns() const noexcept
    {
      return m_column_size;
    }

    [[nodiscard]] PUGS_INLINE RowView
    operator[](size_t i) const
    {
      Assert(i < this->numberOfRows(), "invalid index");
      return RowView(*this, i);
    }

    [[nodiscard]] PUGS_INLINE DataType&
    operator()(size_t i, size_t j) const
    {
      Assert(i < m_row_size, "invalid row index");
      Assert(j < m_column_size, "invalid column index");

      return m_table(m_row_begin + i, m_column_begin + j);
    }

    friend std::ostream&
    operator<<(std::ostream& os, const UnsafeTableView& t)
    {
      for (size_t i = 0; i < t.numberOfRows(); ++i) {
        os << i << '|';
        for (size_t j = 0; j < t.numberOfColumns(); ++j) {
          os << ' ' << j << ':' << NaNHelper(t(i, j));
        }
        os << '\n';
      }
      return os;
    }

    PUGS_INLINE void
    fill(const DataType& data) const
    {
      for (size_t i = 0; i < m_row_size; ++i) {
        for (size_t j = 0; j < m_column_size; ++j) {
          m_table(m_row_begin + i, m_column_begin + j) = data;
        }
      }
    }
    UnsafeTableView& operator=(const UnsafeTableView&) = delete;
    UnsafeTableView& operator=(UnsafeTableView&&)      = delete;

    UnsafeTableView(const Table<DataType>& table,
                    index_type row_begin,
                    index_type row_size,
                    index_type column_begin,
                    index_type column_size)
      : m_table{table},
        m_row_begin{row_begin},
        m_column_begin{column_begin},
        m_row_size{row_size},
        m_column_size{column_size}
    {
      Assert((row_begin < table.numberOfRows()) and (row_begin + row_size <= table.numberOfRows()),
             "required view has rows not contained in the Table");
      Assert((column_begin < table.numberOfColumns()) and (column_begin + column_size <= table.numberOfColumns()),
             "required view has columns not contained in the Table");
    }

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

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

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

  [[nodiscard]] PUGS_INLINE size_t
  numberOfColumns() const noexcept
  {
    return m_values.extent(1);
  }

  [[nodiscard]] PUGS_INLINE Table<DataType>::UnsafeRowView
  operator[](index_type i) const
  {
    return UnsafeRowView(*this, i);
  }

  [[nodiscard]] friend PUGS_INLINE Table<std::remove_const_t<DataType>>
  copy(const Table<DataType>& source)
  {
    Table<std::remove_const_t<DataType>> image(source.numberOfRows(), source.numberOfColumns());
    Kokkos::deep_copy(image.m_values, source.m_values);

    return image;
  }

  friend PUGS_INLINE void
  copy_to(const Table<DataType>& source, const Table<std::remove_const_t<DataType>>& destination)
  {
    Assert(source.numberOfRows() == destination.numberOfRows(), "incompatible number of rows");
    Assert(source.numberOfColumns() == destination.numberOfColumns(), "incompatible number of columns");
    Kokkos::deep_copy(destination.m_values, source.m_values);
  }

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

  template <typename DataType2>
  friend PUGS_INLINE typename Table<DataType2>::UnsafeTableView subTableView(
    const Table<DataType2>& table,
    typename Table<DataType2>::index_type row_begin,
    typename Table<DataType2>::index_type row_size,
    typename Table<DataType2>::index_type column_begin,
    typename Table<DataType2>::index_type column_size);

  [[nodiscard]] PUGS_INLINE DataType&
  operator()(index_type i, index_type j) const noexcept(NO_ASSERT)
  {
    Assert(i < this->numberOfRows(), "invalid row index");
    Assert(j < this->numberOfColumns(), "invalid column index");
    return m_values(i, j);
  }

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

    Kokkos::deep_copy(m_values, data);
  }

  template <typename DataType2>
  PUGS_INLINE Table&
  operator=(const Table<DataType2>& table) 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 Table 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 Table of const to Table of non-const");
    m_values = table.m_values;
    return *this;
  }

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

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

  PUGS_INLINE
  explicit Table(size_t nb_lines, size_t nb_columns)
  {
    static_assert(not std::is_const<DataType>(), "Cannot allocate Table 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"), nb_lines, nb_columns};
    } else {
      m_values = Kokkos::View<DataType**>{"anonymous", nb_lines, nb_columns};
    }

#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 Table& t)
  {
    for (size_t i = 0; i < t.numberOfRows(); ++i) {
      os << i << '|';
      for (size_t j = 0; j < t.numberOfColumns(); ++j) {
        os << ' ' << j << ':' << NaNHelper(t(i, j));
      }
      os << '\n';
    }
    return os;
  }

  PUGS_INLINE
  Table() = default;

  PUGS_INLINE
  Table(const Table&) = default;

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

  PUGS_INLINE
  Table(Table&&) = default;

  PUGS_INLINE
  ~Table() = default;
};

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

template <typename DataType>
[[nodiscard]] PUGS_INLINE typename Table<DataType>::UnsafeTableView
subTableView(const Table<DataType>& table,
             typename Table<DataType>::index_type row_begin,
             typename Table<DataType>::index_type row_size,
             typename Table<DataType>::index_type column_begin,
             typename Table<DataType>::index_type column_size)
{
  return typename Table<DataType>::UnsafeTableView(table, row_begin, row_size, column_begin, column_size);
}

#endif   // TABLE_HPP