Bitwise.hpp

#pragma once
 
#include <stratax/core/dtypes/Concepts.hpp>
#include <stratax/core/validation/Validation.hpp>
#include <stratax/ops/Broadcasting.hpp>
 
#include <functional>
#include <climits>
#include <cstdint>
#include <type_traits>
 
namespace stratax::core::bitwise_detail {

template<typename Value, Integral Count>
constexpr bool valid_shift_count(const Count& count) noexcept
{
    if constexpr (std::is_signed_v<std::remove_cvref_t<Count>>)
    {
        if (count < 0)
        {
            return false;
        }
    }
 
    using value_type = std::remove_cvref_t<Value>;
 
    return static_cast<std::uintmax_t>(count) <
        sizeof(value_type) * CHAR_BIT;
}

template<typename Value, Integral Count>
void require_valid_shift_count(const Count& count)
{
    if (!valid_shift_count<Value>(count))
    {
        throw Exceptions::StrataxError(
            "Shift count must be non-negative and less than the bit width of the shifted value.");
    }
}

template<Array A, Integral Count, typename Op>
requires Integral<typename A::value_type>
auto shift_scalar_op(
    const A& lhs,
    const Count& rhs,
    Op op)
{
    using value_type = typename A::value_type;
 
    require_valid_shift_count<value_type>(rhs);
 
    using result_type =
        stratax::core::rebind_array_t<A, value_type>;
 
    result_type result(lhs.shape());
 
    auto out = result.begin();
 
    for (auto it = lhs.begin(); it != lhs.end(); ++it, ++out)
    {
        *out = static_cast<value_type>(
            op(*it, rhs));
    }
 
    return result;
}

template<Array L, Array R, typename Op>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
auto shift_array_op(
    const L& lhs,
    const R& rhs,
    Op op)
{
    using value_type = typename L::value_type;
 
    using result_type =
        stratax::core::promote_array_t<
            L,
            R,
            value_type>;
 
    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());
 
        const auto count = rhs[rhs_index];
 
        require_valid_shift_count<value_type>(count);
 
        result[i] = static_cast<value_type>(
            op(lhs[lhs_index], count));
    }
 
    return result;
}

template<Integral Scalar, Array A, typename Op>
requires Integral<typename A::value_type>
auto scalar_shift_array_op(
    const Scalar& lhs,
    const A& rhs,
    Op op)
{
    using value_type = std::remove_cvref_t<Scalar>;
 
    using result_type =
        stratax::core::rebind_array_t<
            A,
            value_type>;
 
    result_type result(rhs.shape());
 
    auto out = result.begin();
 
    for (auto it = rhs.begin(); it != rhs.end(); ++it, ++out)
    {
        require_valid_shift_count<value_type>(*it);
 
        *out = static_cast<value_type>(
            op(lhs, *it));
    }
 
    return result;
}

template<Array L, Array R, typename Op>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
L& compound_bitwise_op(
    L& lhs,
    const R& rhs,
    Op op)
{
    const auto result_shape =
        broadcasted_shape(lhs.shape(), rhs.shape());
 
    if (result_shape != lhs.shape())
    {
        throw Exceptions::BroadcastError(
            "Compound bitwise assignment cannot change the left-hand shape.");
    }
 
    for (std::size_t i = 0; i < lhs.size(); ++i)
    {
        const std::size_t rhs_index =
            stratax::core::broadcast_detail::flat_operand_index(
                i,
                lhs.shape(),
                rhs.shape());
 
        lhs[i] = static_cast<typename L::value_type>(
            op(lhs[i], rhs[rhs_index]));
    }
 
    return lhs;
}

template<Array A, Integral Scalar, typename Op>
requires Integral<typename A::value_type>
A& compound_scalar_bitwise_op(
    A& lhs,
    const Scalar& rhs,
    Op op)
{
    for (std::size_t i = 0; i < lhs.size(); ++i)
    {
        lhs[i] = static_cast<typename A::value_type>(
            op(lhs[i], rhs));
    }
 
    return lhs;
}

template<Array L, Array R, typename Op>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
L& compound_shift_op(
    L& lhs,
    const R& rhs,
    Op op)
{
    const auto result_shape =
        broadcasted_shape(lhs.shape(), rhs.shape());
 
    if (result_shape != lhs.shape())
    {
        throw Exceptions::BroadcastError(
            "Compound shift assignment cannot change the left-hand shape.");
    }
 
    using value_type = typename L::value_type;
 
    for (std::size_t i = 0; i < lhs.size(); ++i)
    {
        const std::size_t rhs_index =
            stratax::core::broadcast_detail::flat_operand_index(
                i,
                lhs.shape(),
                rhs.shape());
 
        require_valid_shift_count<value_type>(rhs[rhs_index]);
    }
 
    for (std::size_t i = 0; i < lhs.size(); ++i)
    {
        const std::size_t rhs_index =
            stratax::core::broadcast_detail::flat_operand_index(
                i,
                lhs.shape(),
                rhs.shape());
 
        const auto count = rhs[rhs_index];
 
        lhs[i] = static_cast<value_type>(
            op(lhs[i], count));
    }
 
    return lhs;
}

template<Array L, Array R, typename Op>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
auto binary_bitwise_op(
    const L& lhs,
    const R& rhs,
    Op op)
{
    using result_value_type =
        stratax::core::promote_t<
            typename L::value_type,
            typename R::value_type>;
 
    return stratax::core::broadcasted_op<result_value_type>(
        lhs,
        rhs,
        op);
}

template<Array A, Integral Scalar, typename Op>
auto binary_scalar_bitwise_op(
    const A& arr,
    const Scalar& scalar,
    Op op)
{
    using result_value_type =
        promote_t<typename A::value_type, Scalar>;
 
    using result_type =
        rebind_array_t<A, result_value_type>;
 
    result_type result(arr.shape());
 
    for (std::size_t i = 0; i < arr.size(); ++i)
    {
        result[i] = static_cast<result_value_type>(
            op(arr[i], scalar));
    }
 
    return result;
}

template<Integral Scalar, Array A, typename Op>
auto binary_scalar_bitwise_op(
    const Scalar& scalar,
    const A& arr,
    Op op)
{
    using result_value_type =
        promote_t<Scalar, typename A::value_type>;
 
    using result_type =
        rebind_array_t<A, result_value_type>;
 
    result_type result(arr.shape());
 
    for (std::size_t i = 0; i < arr.size(); ++i)
    {
        result[i] = static_cast<result_value_type>(
            op(scalar, arr[i]));
    }
 
    return result;
}

template<Array A, Integral Count, typename Op>
requires Integral<typename A::value_type>
A& compound_scalar_shift_op(
    A& lhs,
    const Count& rhs,
    Op op)
{
    using value_type = typename A::value_type;
 
    require_valid_shift_count<value_type>(rhs);
 
    for (std::size_t i = 0; i < lhs.size(); ++i)
    {
        lhs[i] = static_cast<value_type>(
            op(lhs[i], rhs));
    }
 
    return lhs;
}
 
} // namespace stratax::core::bitwise_detail
 
// Unary

template<Array A>
requires Integral<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 L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
auto operator&(const L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::binary_bitwise_op(
        lhs, rhs, std::bit_and<>{});
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
auto operator|(const L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::binary_bitwise_op(
        lhs, rhs, std::bit_or<>{});
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
auto operator^(const L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::binary_bitwise_op(
        lhs, rhs, std::bit_xor<>{});
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
auto operator<<(const L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::shift_array_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value << count;
        });
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
auto operator>>(const L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::shift_array_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value >> count;
        });
}
 
// Array-scalar

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

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

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

template<Array A, Integral Scalar>
requires Integral<typename A::value_type>
auto operator<<(const A& lhs, const Scalar& rhs)
{
    return stratax::core::bitwise_detail::shift_scalar_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value << count;
        });
}

template<Array A, Integral Scalar>
requires Integral<typename A::value_type>
auto operator>>(const A& lhs, const Scalar& rhs)
{
    return stratax::core::bitwise_detail::shift_scalar_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value >> count;
        });
}
 
// Scalar-array (reverse)

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

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

template<Integral Scalar, Array A>
requires Integral<typename A::value_type>
auto operator^(const Scalar& lhs, const A& rhs)
{
    return stratax::core::bitwise_detail::binary_scalar_bitwise_op(
        lhs, rhs, std::bit_xor<>{});
}
 
template<Integral Scalar, Array A>
requires Integral<typename A::value_type>
auto operator<<(const Scalar& lhs, const A& rhs)
{
    return stratax::core::bitwise_detail::scalar_shift_array_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value << count;
        });
}
 
template<Integral Scalar, Array A>
requires Integral<typename A::value_type>
auto operator>>(const Scalar& lhs, const A& rhs)
{
    return stratax::core::bitwise_detail::scalar_shift_array_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value >> count;
        });
}
 
// In-place array-array

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
L& operator&=(L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::compound_bitwise_op(
        lhs, rhs, std::bit_and<>{});
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
L& operator|=(L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::compound_bitwise_op(
        lhs, rhs, std::bit_or<>{});
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
L& operator^=(L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::compound_bitwise_op(
        lhs, rhs, std::bit_xor<>{});
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
L& operator<<=(L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::compound_shift_op(
        lhs,
        rhs,
        [](auto value, auto count) {
            return value << count;
        });
}

template<Array L, Array R>
requires (
    Integral<typename L::value_type> &&
    Integral<typename R::value_type>
)
L& operator>>=(L& lhs, const R& rhs)
{
    return stratax::core::bitwise_detail::compound_shift_op(
        lhs,
        rhs,
        [](auto value, auto count) {
            return value >> count;
        });
}
 
// In-place array-scalar
 
template<Array A, Integral Scalar>
requires Integral<typename A::value_type>
A& operator&=(A& lhs, const Scalar& rhs)
{
    return stratax::core::bitwise_detail::compound_scalar_bitwise_op(
        lhs, rhs, std::bit_and<>{});
}
 
template<Array A, Integral Scalar>
requires Integral<typename A::value_type>
A& operator|=(A& lhs, const Scalar& rhs)
{
    return stratax::core::bitwise_detail::compound_scalar_bitwise_op(
        lhs, rhs, std::bit_or<>{});
}
 
template<Array A, Integral Scalar>
requires Integral<typename A::value_type>
A& operator^=(A& lhs, const Scalar& rhs)
{
    return stratax::core::bitwise_detail::compound_scalar_bitwise_op(
        lhs, rhs, std::bit_xor<>{});
}
 
template<Array A, Integral Scalar>
requires Integral<typename A::value_type>
A& operator<<=(A& lhs, const Scalar& rhs)
{
    return stratax::core::bitwise_detail::compound_scalar_shift_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value << count;
        });
}
 
template<Array A, Integral Scalar>
requires Integral<typename A::value_type>
A& operator>>=(A& lhs, const Scalar& rhs)
{
    return stratax::core::bitwise_detail::compound_scalar_shift_op(
        lhs,
        rhs,
        [](auto value, auto count)
        {
            return value >> count;
        });
}