FlatMatrix.hpp

#pragma once

#include <type_traits>
#include <vector>

namespace AMDiS
{
  template <class T, class Allocator = std::allocator<T>>
  class FlatMatrix
      : public std::vector<T,Allocator>
  {
    using Super = std::vector<T,Allocator>;

  public:
    using value_type = T;
    using size_type = typename Super::size_type;
    using reference = typename Super::reference;
    using const_reference = typename Super::const_reference;

  private:
    template <class Vector>
    struct AccessProxy
    {
      template <class V = Vector, std::enable_if_t<not std::is_const_v<V>,int> = 0>
      reference operator[](size_type col)
      {
        return (*vector_)[row_*cols_ + col];
      }

      const_reference operator[](size_type col) const
      {
        return (*vector_)[row_*cols_ + col];
      }
      // @}

      Vector* vector_;
      size_type row_;
      size_type cols_;
    };

  public:
    FlatMatrix() = default;

    FlatMatrix(size_type r, size_type c, value_type v = {})
      : Super(r*c, v)
      , rows_(r)
      , cols_(c)
    {}

    FlatMatrix& operator=(value_type s)
    {
      Super::assign(Super::size(), s);
      return *this;
    }


    void resize(size_type r, size_type c, value_type v = {})
    {
      Super::resize(r*c, v);
      rows_ = r;
      cols_ = c;
    }

    size_type rows() const noexcept
    {
      return rows_;
    }

    size_type cols() const noexcept
    {
      return cols_;
    }

    using Super::size;

    AccessProxy<Super> operator[](size_type row)
    {
      return AccessProxy<Super>{static_cast<Super*>(this), row, cols_};
    }

    AccessProxy<Super const> operator[](size_type row) const
    {
      return AccessProxy<Super const>{static_cast<Super const*>(this), row, cols_};
    }

  private:
    size_type rows_ = 0;
    size_type cols_ = 0;
  };

} // end namespace AMDiS