Tensor.hpp

#pragma once
 
#include <stratax/core/containers/Buffer.hpp>
#include <stratax/core/Concepts.hpp>
#include <stratax/core/containers/Shape.hpp>
#include <stratax/core/Exceptions.hpp>
#include <stratax/core/validation/Validation.hpp>
#include <stratax/core/containers/Strides.hpp>
#include <stratax/core/ops/Indexing.hpp>
 
#include <array>
#include <cstddef>
#include <initializer_list>
#include <stdexcept>
#include <type_traits>
#include <utility>
 
namespace stratax::container {

template<typename T>
requires Numeric<T>
class Tensor
{
private:
    core::Shape shape_;
    core::Strides strides_;
    core::Buffer<T> buffer_;
 
public:
    using value_type = T;

    using iterator = typename core::Buffer<T>::iterator;

    using const_iterator = typename core::Buffer<T>::const_iterator;

    using reverse_iterator = typename core::Buffer<T>::reverse_iterator;

    using const_reverse_iterator = typename core::Buffer<T>::const_reverse_iterator;

    template<typename U>
    using rebind = Tensor<U>;

    Tensor() noexcept = default;

    explicit Tensor(const core::Shape& shape)
        : shape_(shape),
          strides_(shape_),
          buffer_(shape_.elements())
    {
    }

    Tensor(const core::Shape& shape, const T& value)
        : shape_(shape),
          strides_(shape_),
          buffer_(shape_.elements(), value)
    {
    }

    Tensor(const Tensor&) = default;

    Tensor(Tensor&&) noexcept = default;

    Tensor& operator=(const Tensor&) = default;

    Tensor& operator=(Tensor&&) noexcept = default;

    ~Tensor() = default;

    [[nodiscard]] std::size_t size() const noexcept
    {
        return buffer_.size();
    }

    [[nodiscard]] bool empty() const noexcept
    {
        return buffer_.empty();
    }

    [[nodiscard]] std::size_t rank() const noexcept
    {
        return shape_.rank();
    }

    const core::Shape& shape() const noexcept
    {
        return shape_;
    }

    const core::Strides& strides() const noexcept
    {
        return strides_;
    }

    T& operator()(std::size_t index) noexcept
    {
        return buffer_[index];
    }

    const T& operator()(std::size_t index) const noexcept
    {
        return buffer_[index];
    }

    T& operator[](std::size_t index) noexcept
    {
        return buffer_[index];
    }

    const T& operator[](std::size_t index) const noexcept
    {
        return buffer_[index];
    }

    template<typename... Rest>
    T& operator()(std::size_t first, std::size_t second, Rest... rest)
    {
        std::array<std::size_t, sizeof...(Rest) + 2> indices{
            first,
            second,
            static_cast<std::size_t>(rest)...
        };
 
        return buffer_[offset(shape_, strides_, indices)];
    }

    template<typename... Rest>
    const T& operator()(std::size_t first, std::size_t second, Rest... rest) const
    {
        std::array<std::size_t, sizeof...(Rest) + 2> indices{
            first,
            second,
            static_cast<std::size_t>(rest)...
        };
 
        return buffer_[offset(shape_, strides_, indices)];
    }

    T& operator()(const std::vector<std::size_t>& indices)
    {
        return buffer_[offset(shape_, strides_, indices)];
    }

    const T& operator()(const std::vector<std::size_t>& indices) const
    {
        return buffer_[offset(shape_, strides_, indices)];
    }

    T& at(std::ptrdiff_t index)
    {
        const std::size_t normalized =
            core::validation::normalize_index(index, size(), "Tensor flat index out of bounds.");
        return buffer_[normalized];
    }

    const T& at(std::ptrdiff_t index) const
    {
        const std::size_t normalized =
            core::validation::normalize_index(index, size(), "Tensor flat index out of bounds.");
        return buffer_[normalized];
    }

    template<typename... Rest>
    requires ((std::is_integral_v<Rest>) && ...)
    T& at(std::ptrdiff_t first, std::ptrdiff_t second, Rest... rest)
    {
        std::array<std::ptrdiff_t, sizeof...(Rest) + 2> raw_indices{
            first,
            second,
            static_cast<std::ptrdiff_t>(rest)...
        };
 
        std::array<std::size_t, sizeof...(Rest) + 2> indices{};
        if (indices.size() != rank())
        {
            throw Exceptions::IndexError("Tensor multi-index rank must match tensor rank.");
        }
 
        for (std::size_t i = 0; i < indices.size(); ++i)
        {
            indices[i] = core::validation::normalize_index(
                raw_indices[i],
                shape()(i),
                "Tensor multi-index component is out of bounds.");
        }
 
        try
        {
            return buffer_[offset(shape_, strides_, indices)];
        }
        catch (const Exceptions::DimensionError&)
        {
            throw Exceptions::IndexError("Tensor multi-index rank must match tensor rank.");
        }
        catch (const Exceptions::IndexError&)
        {
            throw Exceptions::IndexError("Tensor multi-index component is out of bounds.");
        }
    }

    template<typename... Rest>
    requires ((std::is_integral_v<Rest>) && ...)
    const T& at(std::ptrdiff_t first, std::ptrdiff_t second, Rest... rest) const
    {
        std::array<std::ptrdiff_t, sizeof...(Rest) + 2> raw_indices{
            first,
            second,
            static_cast<std::ptrdiff_t>(rest)...
        };
 
        std::array<std::size_t, sizeof...(Rest) + 2> indices{};
        if (indices.size() != rank())
        {
            throw Exceptions::IndexError("Tensor multi-index rank must match tensor rank.");
        }
 
        for (std::size_t i = 0; i < indices.size(); ++i)
        {
            indices[i] = core::validation::normalize_index(
                raw_indices[i],
                shape()(i),
                "Tensor multi-index component is out of bounds.");
        }
 
        try
        {
            return buffer_[offset(shape_, strides_, indices)];
        }
        catch (const Exceptions::DimensionError&)
        {
            throw Exceptions::IndexError("Tensor multi-index rank must match tensor rank.");
        }
        catch (const Exceptions::IndexError&)
        {
            throw Exceptions::IndexError("Tensor multi-index component is out of bounds.");
        }
    }

    T& front()
    {
        return buffer_.front();
    }

    const T& front() const
    {
        return buffer_.front();
    }

    T& back()
    {
        return buffer_.back();
    }

    const T& back() const
    {
        return buffer_.back();
    }

    [[nodiscard]] T* data() noexcept
    {
        return buffer_.data();
    }

    [[nodiscard]] const T* data() const noexcept
    {
        return buffer_.data();
    }

    [[nodiscard]] iterator begin() noexcept
    {
        return buffer_.begin();
    }

    [[nodiscard]] const_iterator begin() const noexcept
    {
        return buffer_.begin();
    }

    [[nodiscard]] const_iterator cbegin() const noexcept
    {
        return buffer_.cbegin();
    }

    [[nodiscard]] iterator end() noexcept
    {
        return buffer_.end();
    }

    [[nodiscard]] const_iterator end() const noexcept
    {
        return buffer_.end();
    }

    [[nodiscard]] const_iterator cend() const noexcept
    {
        return buffer_.cend();
    }

    [[nodiscard]] reverse_iterator rbegin() noexcept
    {
        return buffer_.rbegin();
    }

    [[nodiscard]] const_reverse_iterator rbegin() const noexcept
    {
        return buffer_.rbegin();
    }

    [[nodiscard]] const_reverse_iterator crbegin() const noexcept
    {
        return buffer_.crbegin();
    }

    [[nodiscard]] reverse_iterator rend() noexcept
    {
        return buffer_.rend();
    }

    [[nodiscard]] const_reverse_iterator rend() const noexcept
    {
        return buffer_.rend();
    }

    [[nodiscard]] const_reverse_iterator crend() const noexcept
    {
        return buffer_.crend();
    }

    void fill(const T& value)
    {
        buffer_.fill(value);
    }

    void swap(Tensor& other) noexcept
    {
        shape_.swap(other.shape_);
        strides_.swap(other.strides_);
        buffer_.swap(other.buffer_);
    }
};
 
}