Shape.hpp

#pragma once
 
#include <vector>
 
#include "Buffer.hpp"
#include "../Concepts.hpp"
#include "../Exceptions.hpp"
#include "../validation/Validation.hpp"
 
#include <ostream>
 
namespace stratax::core {

class Shape
{
private:
    Buffer<std::size_t> dims_;
 
    void validate_dimensions() const
    {
        for (std::size_t dim : dims_)
        {
            validation::nonnegative_shape_dimension(
                dim,
                "Shape dimensions cannot be negative.");
        }
    }
 
public:
    struct allow_zero_t {};

    static constexpr allow_zero_t allow_zero{};

    using iterator = Buffer<std::size_t>::iterator;

    using const_iterator = Buffer<std::size_t>::const_iterator;

    using reverse_iterator = Buffer<std::size_t>::reverse_iterator;

    using const_reverse_iterator = Buffer<std::size_t>::const_reverse_iterator;

    Shape() noexcept = default;

    template<Integral... Dims>
    requires (sizeof...(Dims) > 0)
    Shape(Dims... dims)
        : dims_{static_cast<std::size_t>(dims)...}
    {
        validate_dimensions();
    }

    Shape(std::initializer_list<std::size_t> list, allow_zero_t allow_zero) : dims_(list)
    {
        (void)allow_zero;
    }

    Shape(const Buffer<std::size_t>& dims)
    : dims_(dims)
    {
        validate_dimensions();
    }

    Shape(const Buffer<std::size_t>& dims, allow_zero_t allow_zero)
    : dims_(dims)
    {
        (void)allow_zero;
    }

    Shape(Buffer<std::size_t>&& dims)
    : dims_(std::move(dims))
    {
        validate_dimensions();
    }

    Shape(const std::vector<std::size_t>& dims)
        : dims_(dims.size())
    {
        for (std::size_t i = 0; i < dims.size(); ++i)
        {
            dims_[i] = dims[i];
        }
    }

    ~Shape() = default;

    [[nodiscard]]
    std::size_t elements() const
    {
        if (empty())
        {
            return 0;
        }
        std::size_t prod = 1;
        for (std::size_t dim : dims_)
        {
            prod = validation::checked_multiply(prod, dim, "Shape elements overflow");
        }
        return prod;
    }

    [[nodiscard]] std::size_t rank() const
    {
        return dims_.size();
    }

    const std::size_t& operator()(std::size_t index) const
    {
        validation::require_index(index, rank(), "Shape dimension index out of bounds");
        return dims_[index];
    }

    const std::size_t& operator[](std::ptrdiff_t index) const
    {
        return dims_[validation::normalize_index(index, rank(), "Shape dimension index out of bounds")];
    }

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

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

    [[nodiscard]] bool operator!=(const Shape& other) const noexcept
    {
        return !(*this == other);
    }

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

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

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

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

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

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

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

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

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

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

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

    [[nodiscard]] 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 (std::size_t dim : shape)
    {
        if (!first)
            os << ", ";
 
        os << dim;
        first = false;
    }
 
    if (shape.rank() == 1)
    {
        os << ",";
    }
 
    os << ")";
 
    return os;
}
 
}