Strides.hpp

#pragma once
 
#include "Buffer.hpp"
#include "Shape.hpp"
#include "../Exceptions.hpp"
#include "../validation/Validation.hpp"
 
#include <cstddef>
#include <stdexcept>
#include <utility>
#include <ostream>
 
namespace stratax::core {

class Strides
{
private:
    Buffer<std::size_t> buffer_;
 
public:
    using iterator = Buffer<std::size_t>::const_iterator;

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

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

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

    Strides() noexcept = default;

    explicit Strides(const Shape& shape)
        : buffer_(shape.rank())
    {
        if (shape.empty()) {
            return;
        }
 
        buffer_[shape.rank() - 1] = 1;
 
        for (std::size_t i = shape.rank() - 1; i > 0; --i) {
            buffer_[i - 1] = validation::checked_multiply(
                buffer_[i],
                shape(i),
                "Strides overflow for shape");
        }
    }

    Strides(const Strides&) = default;

    Strides(Strides&&) noexcept = default;

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

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

    ~Strides() = 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 buffer_.size();
    }

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

    const std::size_t& at(std::size_t index) const
    {
        validation::require_index(index, rank(), "Strides index out of bounds");
        return buffer_[index];
    }

    const std::size_t& front() const
    {
        return buffer_.front();
    }

    const std::size_t& back() const
    {
        return buffer_.back();
    }

    [[nodiscard]] const std::size_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 swap(Strides& other) noexcept
    {
        buffer_.swap(other.buffer_);
    }

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

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

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