array2d.h

// Copyright (C) 2021 Igor Baratta
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <array>
#include <cassert>
#include <vector>
#include <xtl/xspan.hpp>
 
namespace dolfinx
{
 
template <typename T, class Allocator = std::allocator<T>>
class array2d
{
public:
  using value_type = T;
  using allocator_type = Allocator;
  using size_type = typename std::vector<T, Allocator>::size_type;
  using reference = typename std::vector<T, Allocator>::reference;
  using const_reference = typename std::vector<T, Allocator>::const_reference;
  using pointer = typename std::vector<T, Allocator>::pointer;
  using iterator = typename std::vector<T, Allocator>::iterator;
  using const_iterator = typename std::vector<T, Allocator>::const_iterator;
 
  array2d(std::array<size_type, 2> shape, value_type value = T(),
          const Allocator& alloc = Allocator())
      : shape(shape)
  {
    _storage = std::vector<T, Allocator>(shape[0] * shape[1], value, alloc);
  }
 
  array2d(size_type rows, size_type cols, value_type value = T(),
          const Allocator& alloc = Allocator())
      : shape({rows, cols})
  {
    _storage = std::vector<T, Allocator>(shape[0] * shape[1], value, alloc);
  }
 
  template <typename Vector>
  array2d(std::array<size_type, 2> shape, Vector&& x)
      : shape(shape), _storage(std::forward<Vector>(x))
  {
    // Do nothing
  }
 
  constexpr array2d(std::initializer_list<std::initializer_list<T>> list)
      : shape({list.size(), (*list.begin()).size()})
  {
    _storage.reserve(shape[0] * shape[1]);
    for (std::initializer_list<T> l : list)
      for (const T val : l)
        _storage.push_back(val);
  }
 
  array2d(const array2d& x) = default;
 
  array2d(array2d&& x) = default;
 
  ~array2d() = default;
 
  array2d& operator=(const array2d& x) = default;
 
  array2d& operator=(array2d&& x) = default;
 
  constexpr reference operator()(size_type i, size_type j)
  {
    return _storage[i * shape[1] + j];
  }
 
  constexpr const_reference operator()(size_type i, size_type j) const
  {
    return _storage[i * shape[1] + j];
  }
 
  constexpr xtl::span<value_type> row(size_type i)
  {
    return xtl::span<value_type>(std::next(_storage.data(), i * shape[1]),
                                 shape[1]);
  }
 
  constexpr xtl::span<const value_type> row(size_type i) const
  {
    return xtl::span<const value_type>(std::next(_storage.data(), i * shape[1]),
                                       shape[1]);
  }
 
  constexpr value_type* data() noexcept { return _storage.data(); }
 
  constexpr const value_type* data() const noexcept { return _storage.data(); };
 
  constexpr size_type size() const noexcept { return _storage.size(); }
 
  constexpr std::array<size_type, 2> strides() const noexcept
  {
    return {shape[1] * sizeof(T), sizeof(T)};
  }
 
  constexpr bool empty() const noexcept { return _storage.empty(); }
 
  std::array<size_type, 2> shape;
 
private:
  std::vector<T, Allocator> _storage;
};
} // namespace dolfinx