Shape.hpp

#pragma once
 
#include <algorithm>
#include <cstddef>
#include <initializer_list>
#include <ostream>
#include <vector>
#include <limits>
 
#include <stratax/core/Buffer.hpp>
#include <stratax/exceptions/Exceptions.hpp>
#include <stratax/indexing/Normalize.hpp>
 
namespace stratax::core {

class Shape
{
private:
    Buffer<std::size_t> dims_;
 
public:
    using value_type = std::size_t;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;
    using const_reference = const value_type&;
    using const_iterator = Buffer<value_type>::const_iterator;
    using const_reverse_iterator = Buffer<value_type>::const_reverse_iterator;

    Shape() noexcept = default;

    Shape(std::initializer_list<value_type> dims)
        : dims_(dims)
    {}

    Shape(const std::vector<value_type>& dims)
        : dims_(dims.size())
    {
        std::copy(dims.begin(), dims.end(), dims_.begin());
    }

    [[nodiscard]] size_type elements() const
    {
        if (empty())
        {
            return 0;
        }
        size_type prod = 1;
        for (value_type dim : dims_)
        {
            if (dim == 0)
            {
                return 0;
            }
 
            if (prod > std::numeric_limits<size_type>::max() / dim)
            {
                throw Exceptions::DimensionError(
                    "Shape element count exceeds the maximum representable size.");
            }
 
            prod *= dim;
        }
 
        return prod;
    }

    [[nodiscard]] size_type rank() const noexcept
    {
        return dims_.size();
    }

    [[nodiscard]] Shape strides() const
    {
        if (empty())
        {
            return {};
        }
 
        std::vector<value_type> stride_values(rank());
        stride_values[rank() - 1] = 1;
 
        for (size_type i = rank() - 1; i > 0; --i)
        {
            if (dims_[i] != 0 &&
                stride_values[i] > std::numeric_limits<value_type>::max() / dims_[i])
            {
                throw Exceptions::DimensionError(
                    "Stride value exceeds the maximum representable size.");
            }
 
            stride_values[i - 1] = stride_values[i] * dims_[i];
        }
 
        return Shape(stride_values);
    }

    [[nodiscard]] const_reference operator[](size_type index) const noexcept
    {
        return dims_[index];
    }

    [[nodiscard]] const_reference at(difference_type index) const
    {
        return dims_[stratax::indexing::normalize_index(index, rank())];
    }

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

    [[nodiscard]] bool operator==(const Shape& other) const noexcept
    {
        if (rank() != other.rank())
        {
            return false;
        }
        for (size_type i = 0; i < rank(); ++i)
        {
            if (dims_[i] != other.dims_[i])
            {
                return false;
            }
        }
        return true;
    }

    const_iterator begin() const noexcept
    {
        return dims_.begin();
    }

    const_iterator end() const noexcept
    {
        return dims_.end();
    }

    const_iterator cbegin() const noexcept
    {
        return dims_.cbegin();
    }

    const_iterator cend() const noexcept
    {
        return dims_.cend();
    }

    const_reverse_iterator rbegin() const noexcept
    {
        return dims_.rbegin();
    }

    const_reverse_iterator crbegin() const noexcept
    {
        return dims_.crbegin();
    }

    const_reverse_iterator rend() const noexcept
    {
        return dims_.rend();
    }

    const_reverse_iterator crend() const noexcept
    {
        return dims_.crend();
    }

    void swap(Shape& other) noexcept
    {
        dims_.swap(other.dims_);
    }
 
};

inline std::ostream& operator<<(std::ostream& os, const Shape& shape)
{
    os << "(";
 
    bool first = true;
    for (Shape::value_type dim : shape)
    {
        if (!first)
            os << ", ";
 
        os << dim;
        first = false;
    }
 
    if (shape.rank() == 1)
    {
        os << ",";
    }
 
    os << ")";
 
    return os;
}
 
}