Tensor.hpp

// TODO: make normalize flat offset more explicit
 
#pragma once
 
#include <array>
#include <cstddef>
#include <type_traits>
#include <vector>
 
#include <stratax/core/dtypes/Concepts.hpp>
#include <stratax/core/ArrayBase.hpp>
#include <stratax/core/Shape.hpp>
#include <stratax/indexing/Indexing.hpp>
 
namespace stratax::container {

template<typename T>
requires DType<T>
class Tensor : public core::ArrayBase<T>
{
public:
    using value_type = typename core::ArrayBase<T>::value_type;
    using size_type = typename core::ArrayBase<T>::size_type;
    using difference_type = typename core::ArrayBase<T>::difference_type;
    using reference = typename core::ArrayBase<T>::reference;
    using const_reference = typename core::ArrayBase<T>::const_reference;
    using pointer = typename core::ArrayBase<T>::pointer;
    using const_pointer = typename core::ArrayBase<T>::const_pointer;
    using iterator = typename core::ArrayBase<T>::iterator;
    using const_iterator = typename core::ArrayBase<T>::const_iterator;
    using reverse_iterator = typename core::ArrayBase<T>::reverse_iterator;
    using const_reverse_iterator = typename core::ArrayBase<T>::const_reverse_iterator;
 
protected:
    using core::ArrayBase<T>::normalized_flat_offset;
 
public:
    using core::ArrayBase<T>::at;

    Tensor() : Tensor(core::Shape{0}) {}

    explicit Tensor(const core::Shape& shape)
        : core::ArrayBase<T>(shape)
    {}

    Tensor(const core::Shape& shape, const_reference value)
        : core::ArrayBase<T>(shape, value)
    {}

    template<typename... Rest>
    requires ((std::is_integral_v<Rest>) && ...)
    reference operator()(size_type first, Rest... rest)
    {
        std::array<size_type, sizeof...(Rest) + 1> indices{
            first,
            static_cast<size_type>(rest)...
        };
 
        return (*this)[indexing::offset(this->strides(), indices)];
    }

    template<typename... Rest>
    requires ((std::is_integral_v<Rest>) && ...)
    const_reference operator()(size_type first, Rest... rest) const
    {
        std::array<size_type, sizeof...(Rest) + 1> indices{
            first,
            static_cast<size_type>(rest)...
        };
 
        return (*this)[indexing::offset(this->strides(), indices)];
    }

    reference operator()(const std::vector<size_type>& indices) {return (*this)[indexing::offset(this->strides(), indices)];}

    const_reference operator()(const std::vector<size_type>& indices) const {return (*this)[indexing::offset(this->strides(), indices)];}

    template<typename... Rest>
    requires ((std::is_integral_v<Rest>) && ...)
    reference at(difference_type first, Rest... rest)
    {
        std::array<difference_type, sizeof...(Rest) + 1> raw_indices{
            first,
            static_cast<difference_type>(rest)...
        };
 
        return (*this)[normalized_flat_offset(
            raw_indices,
            "Tensor multi-index rank must match tensor rank.",
            "Tensor multi-index component is out of bounds.")];
    }

    template<typename... Rest>
    requires ((std::is_integral_v<Rest>) && ...)
    const_reference at(difference_type first, Rest... rest) const
    {
        std::array<difference_type, sizeof...(Rest) + 1> raw_indices{
            first,
            static_cast<difference_type>(rest)...
        };
 
        return (*this)[normalized_flat_offset(
            raw_indices,
            "Tensor multi-index rank must match tensor rank.",
            "Tensor multi-index component is out of bounds.")];
    }

    reference at(const std::vector<difference_type>& raw_indices)
    {
        return (*this)[normalized_flat_offset(
            raw_indices,
            "Tensor multi-index rank must match tensor rank.",
            "Tensor multi-index component is out of bounds.")];
    }

    const_reference at(const std::vector<difference_type>& raw_indices) const
    {
        return (*this)[normalized_flat_offset(
            raw_indices,
            "Tensor multi-index rank must match tensor rank.",
            "Tensor multi-index component is out of bounds.")];
    }

    void swap(Tensor& other) noexcept {core::ArrayBase<T>::swap(other);}

    friend void swap(Tensor& lhs, Tensor& rhs) noexcept {lhs.swap(rhs);}
};
 
} // namespace stratax::container