array.hpp

#ifndef ARRAY_HPP_
#define ARRAY_HPP_

#include <vector>
#include <array>
#include <type_traits>
#include <stdexcept>


/*! Simple data-by-reference strided array class
 */
template <typename Scalar, int NDim, typename Storage = std::vector<Scalar> >
class Array
{
private:
    Storage& data_;
    std::array<size_t, NDim> shape_;
    std::array<size_t, NDim> stride_;
    size_t offset_ = 0;

public:
    Array(Array<Scalar,NDim>&& array)
        : data_(array.data_), shape_(array.shape_), stride_(array.stride_), offset_(array.offset_)
        {}

    Array(Storage& data, const std::array<size_t, NDim> shape, const std::array<size_t, NDim> stride, const size_t offset=0)
        : data_(data), shape_(shape), stride_(stride), offset_(offset)
        {}

    Array(Storage& data, const std::array<size_t, NDim> shape)
        : data_(data), shape_(shape), offset_(0)
        {
            // row-major strides
            for (size_t i = NDim; i > 0; --i) {
                if (i == NDim)
                    stride_[i-1] = 1;
                else
                    stride_[i-1] = shape[i] * stride_[i];
            }
        }

    template <typename... Idx>
    Scalar& operator()(Idx... idxs) { return data_[index(idxs...)]; }

    template <typename... Idx>
    const Scalar&& operator()(Idx... idxs) const { return data_[index(idxs...)]; }

    template <size_t axis>
    Array<Scalar, NDim-1> slice(size_t pos=0) const
        {
            static_assert(axis < NDim, "invalid slice index");

            std::array<size_t, NDim-1> shape;
            std::array<size_t, NDim-1> stride;
            size_t offset;

            offset = stride[axis] * pos;

            for (size_t i = 0, j = 0; i < NDim; ++i) {
                if (i == axis)
                    continue;
                shape[j] = shape_[i];
                stride[j] = stride_[i];
                ++j;
            }

            return {data_, shape, stride, offset};
        }

    template <size_t axis>
    Array<Scalar, NDim> slice(size_t begin, size_t end) const
        {
            static_assert(axis < NDim, "invalid slice index");

            std::array<size_t, NDim> shape;
            std::array<size_t, NDim> stride;
            size_t offset;

            offset = stride[axis] * begin;

            for (size_t i = 0; i < NDim; ++i) {
                if (i == axis)
                    shape[i] = end - begin;
                else
                    shape[i] = shape_[i];
                stride[i] = stride_[i];
            }

            return {data_, shape, stride, offset};
        }

    Storage& data() { return data_; }
    
    template <size_t axis>
    const size_t shape() const { return shape_[axis]; }

    template <size_t axis>
    const size_t stride() const { return stride_[axis]; }

    size_t offset() const { return offset_; }

    template <typename... Idx>
    size_t index(Idx... idxs) const
        {
            static_assert(sizeof...(idxs) == NDim,
                          "number of indices must equal the number of dimensions");

            const std::array<size_t, NDim> m{idxs...};
            size_t idx = offset_;

            for (size_t i = 0; i < NDim; ++i) {
                idx += stride_[i] * m[i];
#ifndef NO_BOUNDS_CHECK
                if (m[i] >= shape_[i])
                    throw std::out_of_range("index out of bounds");
#endif
            }

            return idx;
        }
};

#endif