Comparison.hpp

#pragma once
 
#include <functional>
 
#include <stratax/core/dtypes/Concepts.hpp>
#include <stratax/core/ArrayTraits.hpp>
#include <stratax/core/dtypes/Types.hpp>
#include <stratax/ops/Broadcasting.hpp>
 
namespace stratax::core::comparison_detail {
 
template<Array L, Array R, typename Op>
auto comparison_op(
    const L& lhs,
    const R& rhs,
    Op op)
{
    using result_type =
        stratax::core::promote_array_t<
            L,
            R,
            stratax::dtype::bool_>;
 
    const auto result_shape =
        broadcasted_shape(lhs.shape(), rhs.shape());
 
    result_type result(result_shape);
 
    for (std::size_t i = 0; i < result.size(); ++i)
    {
        const std::size_t lhs_index =
            stratax::core::broadcast_detail::flat_operand_index(
                i,
                result_shape,
                lhs.shape());
 
        const std::size_t rhs_index =
            stratax::core::broadcast_detail::flat_operand_index(
                i,
                result_shape,
                rhs.shape());
 
        result[i] = static_cast<stratax::dtype::bool_>(
            op(lhs[lhs_index], rhs[rhs_index]));
    }
 
    return result;
}
 
template<Array A, DType Scalar, typename Op>
auto comparison_scalar_op(
    const A& lhs,
    const Scalar& rhs,
    Op op)
{
    using result_type =
        stratax::core::rebind_array_t<
            A,
            stratax::dtype::bool_>;
 
    result_type result(lhs.shape());
 
    auto out = result.begin();
 
    for (auto it = lhs.begin(); it != lhs.end(); ++it, ++out)
    {
        *out = static_cast<stratax::dtype::bool_>(
            op(*it, rhs));
    }
 
    return result;
}
 
template<DType Scalar, Array A, typename Op>
auto comparison_scalar_op(
    const Scalar& lhs,
    const A& rhs,
    Op op)
{
    using result_type =
        stratax::core::rebind_array_t<
            A,
            stratax::dtype::bool_>;
 
    result_type result(rhs.shape());
 
    auto out = result.begin();
 
    for (auto it = rhs.begin(); it != rhs.end(); ++it, ++out)
    {
        *out = static_cast<stratax::dtype::bool_>(
            op(lhs, *it));
    }
 
    return result;
}
 
template<Array L, Array R>
[[nodiscard]] bool array_equal(const L& lhs, const R& rhs)
{
    if (lhs.shape() != rhs.shape())
    {
        return false;
    }
 
    auto lhs_it = lhs.begin();
    auto rhs_it = rhs.begin();
 
    for (; lhs_it != lhs.end(); ++lhs_it, ++rhs_it)
    {
        if (*lhs_it != *rhs_it)
        {
            return false;
        }
    }
 
    return true;
}
 
} // namespace stratax::core::comparison_detail
 
template<Array L, Array R>
auto equal(const L& lhs, const R& rhs)
{
    return stratax::core::comparison_detail::comparison_op(
        lhs, rhs, std::equal_to<>{});
}
 
template<Array L, Array R>
auto not_equal(const L& lhs, const R& rhs)
{
    return stratax::core::comparison_detail::comparison_op(
        lhs, rhs, std::not_equal_to<>{});
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto less(const L& lhs, const R& rhs)
{
    return stratax::core::comparison_detail::comparison_op(
        lhs, rhs, std::less<>{});
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto less_equal(const L& lhs, const R& rhs)
{
    return stratax::core::comparison_detail::comparison_op(
        lhs, rhs, std::less_equal<>{});
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto greater(const L& lhs, const R& rhs)
{
    return stratax::core::comparison_detail::comparison_op(
        lhs, rhs, std::greater<>{});
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto greater_equal(const L& lhs, const R& rhs)
{
    return stratax::core::comparison_detail::comparison_op(
        lhs, rhs, std::greater_equal<>{});
}
 
template<Array A, DType Scalar>
auto equal(const A& lhs, const Scalar& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::equal_to<>{});
}
 
template<DType Scalar, Array A>
auto equal(const Scalar& lhs, const A& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::equal_to<>{});
}
 
template<Array A, DType Scalar>
auto not_equal(const A& lhs, const Scalar& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::not_equal_to<>{});
}
 
template<DType Scalar, Array A>
auto not_equal(const Scalar& lhs, const A& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::not_equal_to<>{});
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto less(const A& lhs, const Scalar& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::less<>{});
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto less(const Scalar& lhs, const A& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::less<>{});
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto less_equal(const A& lhs, const Scalar& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::less_equal<>{});
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto less_equal(const Scalar& lhs, const A& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::less_equal<>{});
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto greater(const A& lhs, const Scalar& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::greater<>{});
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto greater(const Scalar& lhs, const A& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::greater<>{});
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto greater_equal(const A& lhs, const Scalar& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::greater_equal<>{});
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto greater_equal(const Scalar& lhs, const A& rhs)
{
    return stratax::core::comparison_detail::comparison_scalar_op(
        lhs, rhs, std::greater_equal<>{});
}
 
template<Array L, Array R>
auto operator==(const L& lhs, const R& rhs)
{
    return equal(lhs, rhs);
}
 
template<Array L, Array R>
auto operator!=(const L& lhs, const R& rhs)
{
    return not_equal(lhs, rhs);
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto operator<(const L& lhs, const R& rhs)
{
    return less(lhs, rhs);
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto operator<=(const L& lhs, const R& rhs)
{
    return less_equal(lhs, rhs);
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto operator>(const L& lhs, const R& rhs)
{
    return greater(lhs, rhs);
}
 
template<Array L, Array R>
requires (
    Ordered<typename L::value_type> &&
    Ordered<typename R::value_type>
)
auto operator>=(const L& lhs, const R& rhs)
{
    return greater_equal(lhs, rhs);
}
 
template<Array A, DType Scalar>
auto operator==(const A& lhs, const Scalar& rhs)
{
    return equal(lhs, rhs);
}
 
template<DType Scalar, Array A>
auto operator==(const Scalar& lhs, const A& rhs)
{
    return equal(lhs, rhs);
}
 
template<Array A, DType Scalar>
auto operator!=(const A& lhs, const Scalar& rhs)
{
    return not_equal(lhs, rhs);
}
 
template<DType Scalar, Array A>
auto operator!=(const Scalar& lhs, const A& rhs)
{
    return not_equal(lhs, rhs);
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto operator<(const A& lhs, const Scalar& rhs)
{
    return less(lhs, rhs);
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto operator<(const Scalar& lhs, const A& rhs)
{
    return less(lhs, rhs);
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto operator<=(const A& lhs, const Scalar& rhs)
{
    return less_equal(lhs, rhs);
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto operator<=(const Scalar& lhs, const A& rhs)
{
    return less_equal(lhs, rhs);
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto operator>(const A& lhs, const Scalar& rhs)
{
    return greater(lhs, rhs);
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto operator>(const Scalar& lhs, const A& rhs)
{
    return greater(lhs, rhs);
}
 
template<Array A, Ordered Scalar>
requires Ordered<typename A::value_type>
auto operator>=(const A& lhs, const Scalar& rhs)
{
    return greater_equal(lhs, rhs);
}
 
template<Ordered Scalar, Array A>
requires Ordered<typename A::value_type>
auto operator>=(const Scalar& lhs, const A& rhs)
{
    return greater_equal(lhs, rhs);
}