Bitwise.hpp

#pragma once
 
#include <stratax/core/Concepts.hpp>
#include <stratax/core/validation/Validation.hpp>
 
#include <functional>

template<Array A>
requires Integer<typename A::value_type>
void require_same_bitwise_shape(const A& lhs, const A& rhs)
{
    stratax::core::validation::require_same_shape(
        lhs,
        rhs,
        "Bitwise operands must have the same shape.");
}

template<Array A, typename Op>
requires Integer<typename A::value_type>
A binary_bitwise_op(const A& lhs, const A& rhs, Op op)
{
    require_same_bitwise_shape(lhs, rhs);
 
    A result(lhs.shape());
 
    auto out = result.begin();
    auto it_lhs = lhs.begin();
    auto it_rhs = rhs.begin();
 
    for (; it_lhs != lhs.end(); ++it_lhs, ++it_rhs, ++out)
    {
        *out = static_cast<typename A::value_type>(op(*it_lhs, *it_rhs));
    }
 
    return result;
}

template<Array A, Integer Scalar, typename Op>
requires Integer<typename A::value_type>
A binary_scalar_bitwise_op(const A& lhs, const Scalar& rhs, Op op)
{
    A result(lhs.shape());
 
    auto out = result.begin();
    for (auto it = lhs.begin(); it != lhs.end(); ++it, ++out)
    {
        *out = static_cast<typename A::value_type>(op(*it, rhs));
    }
 
    return result;
}

template<Integer Scalar, Array A, typename Op>
requires Integer<typename A::value_type>
A binary_scalar_bitwise_op(const Scalar& lhs, const A& rhs, Op op)
{
    A result(rhs.shape());
 
    auto out = result.begin();
    for (auto it = rhs.begin(); it != rhs.end(); ++it, ++out)
    {
        *out = static_cast<typename A::value_type>(op(lhs, *it));
    }
 
    return result;
}
 
// Unary

template<Array A>
requires Integer<typename A::value_type>
A operator~(const A& value)
{
    A result(value.shape());
 
    auto out = result.begin();
    for (auto it = value.begin(); it != value.end(); ++it, ++out)
    {
        *out = static_cast<typename A::value_type>(std::bit_not<>{}(*it));
    }
 
    return result;
}
 
// Array-array

template<Array A>
requires Integer<typename A::value_type>
A operator&(const A& lhs, const A& rhs)
{
    return binary_bitwise_op(lhs, rhs, std::bit_and<>{});
}

template<Array A>
requires Integer<typename A::value_type>
A operator|(const A& lhs, const A& rhs)
{
    return binary_bitwise_op(lhs, rhs, std::bit_or<>{});
}

template<Array A>
requires Integer<typename A::value_type>
A operator^(const A& lhs, const A& rhs)
{
    return binary_bitwise_op(lhs, rhs, std::bit_xor<>{});
}

template<Array A>
requires Integer<typename A::value_type>
A operator<<(const A& lhs, const A& rhs)
{
    return binary_bitwise_op(lhs, rhs, [](auto left, auto right) { return left << right; });
}

template<Array A>
requires Integer<typename A::value_type>
A operator>>(const A& lhs, const A& rhs)
{
    return binary_bitwise_op(lhs, rhs, [](auto left, auto right) { return left >> right; });
}
 
// Array-scalar

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A operator&(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, std::bit_and<>{});
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A operator|(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, std::bit_or<>{});
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A operator^(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, std::bit_xor<>{});
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A operator<<(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, [](auto left, auto right) { return left << right; });
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A operator>>(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, [](auto left, auto right) { return left >> right; });
}
 
// Scalar-array (reverse)

template<Integer Scalar, Array A>
requires Integer<typename A::value_type>
A operator&(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, std::bit_and<>{});
}

template<Integer Scalar, Array A>
requires Integer<typename A::value_type>
A operator|(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, std::bit_or<>{});
}

template<Integer Scalar, Array A>
requires Integer<typename A::value_type>
A operator^(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, std::bit_xor<>{});
}

template<Integer Scalar, Array A>
requires Integer<typename A::value_type>
A operator<<(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, [](auto left, auto right) { return left << right; });
}

template<Integer Scalar, Array A>
requires Integer<typename A::value_type>
A operator>>(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_bitwise_op(lhs, rhs, [](auto left, auto right) { return left >> right; });
}
 
// In-place array-array

template<Array A>
requires Integer<typename A::value_type>
A& operator&=(A& lhs, const A& rhs)
{
    lhs = lhs & rhs;
    return lhs;
}

template<Array A>
requires Integer<typename A::value_type>
A& operator|=(A& lhs, const A& rhs)
{
    lhs = lhs | rhs;
    return lhs;
}

template<Array A>
requires Integer<typename A::value_type>
A& operator^=(A& lhs, const A& rhs)
{
    lhs = lhs ^ rhs;
    return lhs;
}

template<Array A>
requires Integer<typename A::value_type>
A& operator<<=(A& lhs, const A& rhs)
{
    lhs = lhs << rhs;
    return lhs;
}

template<Array A>
requires Integer<typename A::value_type>
A& operator>>=(A& lhs, const A& rhs)
{
    lhs = lhs >> rhs;
    return lhs;
}
 
// In-place array-scalar

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A& operator&=(A& lhs, const Scalar& rhs)
{
    lhs = lhs & rhs;
    return lhs;
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A& operator|=(A& lhs, const Scalar& rhs)
{
    lhs = lhs | rhs;
    return lhs;
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A& operator^=(A& lhs, const Scalar& rhs)
{
    lhs = lhs ^ rhs;
    return lhs;
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A& operator<<=(A& lhs, const Scalar& rhs)
{
    lhs = lhs << rhs;
    return lhs;
}

template<Array A, Integer Scalar>
requires Integer<typename A::value_type>
A& operator>>=(A& lhs, const Scalar& rhs)
{
    lhs = lhs >> rhs;
    return lhs;
}