SubItemValuePerItem.hpp

#ifndef SUBITEM_VALUE_PER_ITEM_HPP
#define SUBITEM_VALUE_PER_ITEM_HPP

#include <mesh/ConnectivityMatrix.hpp>
#include <mesh/IConnectivity.hpp>
#include <mesh/ItemId.hpp>
#include <mesh/ItemOfItemType.hpp>
#include <mesh/ItemType.hpp>
#include <utils/Array.hpp>
#include <utils/PugsAssert.hpp>

#include <memory>

template <typename DataType, typename ItemOfItem, typename ConnectivityPtr = std::shared_ptr<const IConnectivity>>
class SubItemValuePerItem
{
 public:
  static constexpr ItemType item_type{ItemOfItem::item_type};
  static constexpr ItemType sub_item_type{ItemOfItem::sub_item_type};

  using ItemOfItemType = ItemOfItem;
  using data_type      = DataType;
  using ItemId         = ItemIdT<item_type>;
  using index_type     = ItemId;

 private:
  using ConnectivitySharedPtr = std::shared_ptr<const IConnectivity>;
  using ConnectivityWeakPtr   = std::weak_ptr<const IConnectivity>;

  static_assert(std::is_same_v<ConnectivityPtr, ConnectivitySharedPtr> or
                std::is_same_v<ConnectivityPtr, ConnectivityWeakPtr>);

  ConnectivityPtr m_connectivity_ptr;

  typename ConnectivityMatrix::HostRowType m_host_row_map;
  Array<DataType> m_values;

  // Allow const std:shared_ptr version to access our data
  friend SubItemValuePerItem<std::add_const_t<DataType>, ItemOfItem, ConnectivitySharedPtr>;

  // Allow const std:weak_ptr version to access our data
  friend SubItemValuePerItem<std::add_const_t<DataType>, ItemOfItem, ConnectivityWeakPtr>;

  // Allow const std:shared_ptr version to access our data
  friend SubItemValuePerItem<std::remove_const_t<DataType>, ItemOfItem, ConnectivitySharedPtr>;

  // Allow const std:weak_ptr version to access our data
  friend SubItemValuePerItem<std::remove_const_t<DataType>, ItemOfItem, ConnectivityWeakPtr>;

 public:
  using ToShared = SubItemValuePerItem<DataType, ItemOfItem, ConnectivitySharedPtr>;

  class SubView
  {
   public:
    using data_type = DataType;

   private:
    PUGS_RESTRICT DataType* const m_sub_values;
    const size_t m_size;

   public:
    template <typename IndexType>
    PUGS_INLINE const DataType&
    operator[](IndexType i) const noexcept(NO_ASSERT)
    {
      static_assert(std::is_integral_v<IndexType>, "SubView is indexed by integral values");
      Assert(i < m_size);
      return m_sub_values[i];
    }

    template <typename IndexType>
    PUGS_FORCEINLINE DataType&
    operator[](IndexType i) noexcept(NO_ASSERT)
    {
      static_assert(std::is_integral_v<IndexType>, "SubView is indexed by integral values");
      Assert(i < m_size);
      return m_sub_values[i];
    }

    PUGS_INLINE
    size_t
    size() const noexcept
    {
      return m_size;
    }

    SubView(const SubView&)     = delete;
    SubView(SubView&&) noexcept = delete;

    PUGS_INLINE
    SubView(const Array<DataType>& values, size_t begin, size_t end) noexcept(NO_ASSERT)
      : m_sub_values(&(values[begin])), m_size(end - begin)
    {
      Assert(begin <= end);
      Assert(end <= values.size());
    }

    ~SubView() = default;
  };

  friend PUGS_INLINE SubItemValuePerItem<std::remove_const_t<DataType>, ItemOfItem, ConnectivityPtr>
  copy(SubItemValuePerItem<DataType, ItemOfItem, ConnectivityPtr>& source)
  {
    SubItemValuePerItem<std::remove_const_t<DataType>, ItemOfItem, ConnectivityPtr> image;

    image.m_connectivity_ptr = source.m_connectivity_ptr;
    image.m_host_row_map     = source.m_host_row_map;
    image.m_values           = copy(source.m_values);
    return image;
  }

  PUGS_INLINE
  bool
  isBuilt() const noexcept
  {
    return m_connectivity_ptr.use_count() != 0;
  }

  PUGS_INLINE
  std::shared_ptr<const IConnectivity>
  connectivity_ptr() const noexcept
  {
    if constexpr (std::is_same_v<ConnectivityPtr, ConnectivitySharedPtr>) {
      return m_connectivity_ptr;
    } else {
      return m_connectivity_ptr.lock();
    }
  }

  template <typename IndexType, typename SubIndexType>
  PUGS_FORCEINLINE DataType&
  operator()(IndexType item_id, SubIndexType i) const noexcept(NO_ASSERT)
  {
    static_assert(std::is_same_v<IndexType, ItemId>, "first index must be of the correct ItemId type");
    static_assert(std::is_integral_v<SubIndexType>, "second index must be an integral type");
    Assert(this->isBuilt());
    Assert(i + m_host_row_map(size_t{item_id}) < m_host_row_map(size_t{item_id} + 1));
    return m_values[m_host_row_map(size_t{item_id}) + i];
  }

  // Following Kokkos logic, these classes are view and const view does allow
  // changes in data
  template <typename ArrayIndexType>
  DataType&
  operator[](const ArrayIndexType& i) const noexcept(NO_ASSERT)
  {
    static_assert(std::is_integral_v<ArrayIndexType>, "index must be an integral type");
    Assert(this->isBuilt());
    Assert(static_cast<size_t>(i) < m_values.size());
    return m_values[i];
  }

  PUGS_INLINE
  size_t
  numberOfValues() const noexcept(NO_ASSERT)
  {
    Assert(this->isBuilt());
    return m_values.size();
  }

  PUGS_INLINE
  size_t
  numberOfItems() const noexcept(NO_ASSERT)
  {
    Assert(this->isBuilt());
    Assert(m_host_row_map.extent(0) > 0);
    return m_host_row_map.extent(0) - 1;
  }

  template <typename IndexType>
  PUGS_INLINE size_t
  numberOfSubValues(IndexType item_id) const noexcept(NO_ASSERT)
  {
    static_assert(std::is_same_v<IndexType, ItemId>, "index must be an ItemId");
    Assert(this->isBuilt());
    Assert(item_id < this->numberOfItems());
    return m_host_row_map(size_t{item_id} + 1) - m_host_row_map(size_t{item_id});
  }

  template <typename IndexType>
  PUGS_INLINE SubView
  itemValues(IndexType item_id) noexcept(NO_ASSERT)
  {
    static_assert(std::is_same_v<IndexType, ItemId>, "index must be an ItemId");
    Assert(this->isBuilt());
    Assert(item_id < this->numberOfItems());
    const auto& item_begin = m_host_row_map(size_t{item_id});
    const auto& item_end   = m_host_row_map(size_t{item_id} + 1);
    return SubView(m_values, item_begin, item_end);
  }

  // Following Kokkos logic, these classes are view and const view does allow
  // changes in data
  template <typename IndexType>
  PUGS_INLINE SubView
  itemValues(IndexType item_id) const noexcept(NO_ASSERT)
  {
    static_assert(std::is_same_v<IndexType, ItemId>, "index must be an ItemId");
    Assert(this->isBuilt());
    Assert(item_id < this->numberOfItems());
    const auto& item_begin = m_host_row_map(size_t{item_id});
    const auto& item_end   = m_host_row_map(size_t{item_id} + 1);
    return SubView(m_values, item_begin, item_end);
  }

  template <typename DataType2, typename ConnectivityPtr2>
  PUGS_INLINE SubItemValuePerItem&
  operator=(const SubItemValuePerItem<DataType2, ItemOfItem, ConnectivityPtr2>& sub_item_value_per_item) noexcept
  {
    // ensures that DataType is the same as source DataType2
    static_assert(std::is_same_v<std::remove_const_t<DataType>, std::remove_const_t<DataType2>>,
                  "Cannot assign SubItemValuePerItem 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 SubItemValuePerItem of const data to "
                  "SubItemValuePerItem of non-const data");
    m_host_row_map = sub_item_value_per_item.m_host_row_map;
    m_values       = sub_item_value_per_item.m_values;

    if constexpr (std::is_same_v<ConnectivityPtr, ConnectivitySharedPtr> and
                  std::is_same_v<ConnectivityPtr2, ConnectivityWeakPtr>) {
      m_connectivity_ptr = sub_item_value_per_item.m_connectivity_ptr.lock();
    } else {
      m_connectivity_ptr = sub_item_value_per_item.m_connectivity_ptr;
    }

    return *this;
  }

  template <typename DataType2, typename ConnectivityPtr2>
  PUGS_INLINE
  SubItemValuePerItem(
    const SubItemValuePerItem<DataType2, ItemOfItem, ConnectivityPtr2>& sub_item_value_per_item) noexcept
  {
    this->operator=(sub_item_value_per_item);
  }

  SubItemValuePerItem() = default;

  SubItemValuePerItem(const IConnectivity& connectivity) noexcept : m_connectivity_ptr{connectivity.shared_ptr()}
  {
    static_assert(not std::is_const_v<DataType>, "Cannot allocate SubItemValuePerItem of const data: only "
                                                 "view is supported");

    ConnectivityMatrix connectivity_matrix = connectivity._getMatrix(item_type, sub_item_type);

    m_host_row_map = connectivity_matrix.rowsMap();
    m_values       = Array<std::remove_const_t<DataType>>(connectivity_matrix.numEntries());
  }

  ~SubItemValuePerItem() = default;
};

// Item values at nodes

template <typename DataType>
using NodeValuePerEdge = SubItemValuePerItem<DataType, NodeOfEdge>;

template <typename DataType>
using NodeValuePerFace = SubItemValuePerItem<DataType, NodeOfFace>;

template <typename DataType>
using NodeValuePerCell = SubItemValuePerItem<DataType, NodeOfCell>;

// Item values at edges

template <typename DataType>
using EdgeValuePerNode = SubItemValuePerItem<DataType, EdgeOfNode>;

template <typename DataType>
using EdgeValuePerFace = SubItemValuePerItem<DataType, EdgeOfFace>;

template <typename DataType>
using EdgeValuePerCell = SubItemValuePerItem<DataType, EdgeOfCell>;

// Item values at faces

template <typename DataType>
using FaceValuePerNode = SubItemValuePerItem<DataType, FaceOfNode>;

template <typename DataType>
using FaceValuePerEdge = SubItemValuePerItem<DataType, FaceOfEdge>;

template <typename DataType>
using FaceValuePerCell = SubItemValuePerItem<DataType, FaceOfCell>;

// Item values at cells

template <typename DataType>
using CellValuePerNode = SubItemValuePerItem<DataType, CellOfNode>;

template <typename DataType>
using CellValuePerEdge = SubItemValuePerItem<DataType, CellOfEdge>;

template <typename DataType>
using CellValuePerFace = SubItemValuePerItem<DataType, CellOfFace>;

// Weak versions: should not be used outside of Connectivity
// Item values at nodes

template <typename DataType, typename ItemOfItem>
using WeakSubItemValuePerItem = SubItemValuePerItem<DataType, ItemOfItem, std::weak_ptr<const IConnectivity>>;

template <typename DataType>
using WeakNodeValuePerEdge = WeakSubItemValuePerItem<DataType, NodeOfEdge>;

template <typename DataType>
using WeakNodeValuePerFace = WeakSubItemValuePerItem<DataType, NodeOfFace>;

template <typename DataType>
using WeakNodeValuePerCell = WeakSubItemValuePerItem<DataType, NodeOfCell>;

// Item values at edges

template <typename DataType>
using WeakEdgeValuePerNode = WeakSubItemValuePerItem<DataType, EdgeOfNode>;

template <typename DataType>
using WeakEdgeValuePerFace = WeakSubItemValuePerItem<DataType, EdgeOfFace>;

template <typename DataType>
using WeakEdgeValuePerCell = WeakSubItemValuePerItem<DataType, EdgeOfCell>;

// Item values at faces

template <typename DataType>
using WeakFaceValuePerNode = WeakSubItemValuePerItem<DataType, FaceOfNode>;

template <typename DataType>
using WeakFaceValuePerEdge = WeakSubItemValuePerItem<DataType, FaceOfEdge>;

template <typename DataType>
using WeakFaceValuePerCell = WeakSubItemValuePerItem<DataType, FaceOfCell>;

// Item values at cells

template <typename DataType>
using WeakCellValuePerNode = WeakSubItemValuePerItem<DataType, CellOfNode>;

template <typename DataType>
using WeakCellValuePerEdge = WeakSubItemValuePerItem<DataType, CellOfEdge>;

template <typename DataType>
using WeakCellValuePerFace = WeakSubItemValuePerItem<DataType, CellOfFace>;

#endif   // SUBITEM_VALUE_PER_ITEM_HPP