binary_expr.hpp

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

namespace etl {

template <typename T, expr_or_scalar<T> LeftExpr, typename BinaryOp, expr_or_scalar<T> RightExpr>
struct binary_expr final : dim_testable<binary_expr<T, LeftExpr, BinaryOp, RightExpr>>,
                           value_testable<binary_expr<T, LeftExpr, BinaryOp, RightExpr>>,
                           iterable<binary_expr<T, LeftExpr, BinaryOp, RightExpr>> {
private:
    using this_type = binary_expr<T, LeftExpr, BinaryOp, RightExpr>; 

    LeftExpr lhs;  
    RightExpr rhs; 

    friend struct etl_traits<binary_expr>;
    friend struct optimizer<binary_expr>;
    friend struct optimizable<binary_expr>;
    friend struct transformer<binary_expr>;

public:
    using value_type        = T;                              
    using memory_type       = void;                           
    using const_memory_type = void;                           
    using iterator          = etl::iterator<const this_type>; 
    using const_iterator    = etl::iterator<const this_type>; 
    using operator_type     = BinaryOp;                       

    using left_type  = std::decay_t<LeftExpr>;  
    using right_type = std::decay_t<RightExpr>; 

    template <typename V = default_vec>
    using vec_type = typename V::template vec_type<T>;

    //Cannot be constructed with no args
    binary_expr() = delete;

    binary_expr(LeftExpr l, RightExpr r) : lhs(std::forward<LeftExpr>(l)), rhs(std::forward<RightExpr>(r)) {
        //Nothing else to init
    }

    binary_expr(const binary_expr& e) = default;

    binary_expr(binary_expr&& e) noexcept = default;

    //Expressions are invariant
    binary_expr& operator=(const binary_expr& e) = delete;
    binary_expr& operator=(binary_expr&& e) = delete;

    template <typename E>
    bool alias(const E& other) const noexcept {
        return lhs.alias(other) || rhs.alias(other);
    }

    //Apply the expression

    value_type operator[](size_t i) const {
        return BinaryOp::apply(lhs[i], rhs[i]);
    }

    value_type read_flat(size_t i) const {
        return BinaryOp::apply(lhs.read_flat(i), rhs.read_flat(i));
    }

    template <typename V = default_vec>
    ETL_STRONG_INLINE(vec_type<V>)
    load(size_t i) const {
        return BinaryOp::template load<V>(lhs.template load<V>(i), rhs.template load<V>(i));
    }

    template <typename V = default_vec>
    ETL_STRONG_INLINE(vec_type<V>)
    loadu(size_t i) const {
        return BinaryOp::template load<V>(lhs.template loadu<V>(i), rhs.template loadu<V>(i));
    }

    template <size_c... S>
    value_type operator()(S... args) const requires(sizeof...(S) == safe_dimensions<this_type>) {
        return BinaryOp::apply(lhs(args...), rhs(args...));
    }

    auto operator()(size_t i) requires sub_capable<this_type> {
        return sub(*this, i);
    }

    auto operator()(size_t i) const requires sub_capable<this_type> {
        return sub(*this, i);
    }

    auto slice(size_t first, size_t last) noexcept {
        return etl::slice(*this, first, last);
    }

    auto slice(size_t first, size_t last) const noexcept {
        return etl::slice(*this, first, last);
    }

    template <typename Y>
    decltype(auto) gpu_compute_hint(Y& y) const {
        return BinaryOp::gpu_compute_hint(lhs, rhs, y);
    }

    template <typename Y>
    decltype(auto) gpu_compute(Y& y) const {
        return BinaryOp::gpu_compute(lhs, rhs, y);
    }

    // Assignment functions

    template <typename L>
    void assign_to(L&& result) const {
        std_assign_evaluate(*this, std::forward<L>(result));
    }

    template <typename L>
    void assign_add_to(L&& result) const {
        std_add_evaluate(*this, std::forward<L>(result));
    }

    template <typename L>
    void assign_sub_to(L&& result) const {
        std_sub_evaluate(*this, std::forward<L>(result));
    }

    template <typename L>
    void assign_mul_to(L&& result) const {
        std_mul_evaluate(*this, std::forward<L>(result));
    }

    template <typename L>
    void assign_div_to(L&& result) const {
        std_div_evaluate(*this, std::forward<L>(result));
    }

    template <typename L>
    void assign_mod_to(L&& result) const {
        std_mod_evaluate(*this, std::forward<L>(result));
    }

    // Internals

    void visit(detail::evaluator_visitor& visitor) const {
        lhs.visit(visitor);
        rhs.visit(visitor);
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        lhs.ensure_cpu_up_to_date();
        rhs.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        lhs.ensure_gpu_up_to_date();
        rhs.ensure_gpu_up_to_date();
    }

    const LeftExpr& get_lhs() const {
        return lhs;
    }

    const RightExpr& get_rhs() const {
        return rhs;
    }

    friend std::ostream& operator<<(std::ostream& os, const binary_expr& expr) {
        if constexpr (BinaryOp::desc_func) {
            return os << BinaryOp::desc() << "(" << expr.lhs << ", " << expr.rhs << ")";
        } else {
            return os << "(" << expr.lhs << ' ' << BinaryOp::desc() << ' ' << expr.rhs << ")";
        }
    }
};

template <typename T, typename LE, typename BinaryOp, typename RE>
struct etl_traits<etl::binary_expr<T, LE, BinaryOp, RE>> {
    using expr_t       = etl::binary_expr<T, LE, BinaryOp, RE>; 
    using left_expr_t  = std::decay_t<LE>;                      
    using right_expr_t = std::decay_t<RE>;                      
    using value_type   = T;                                     

    static constexpr bool left_directed =
        !etl_traits<left_expr_t>::is_generator; 

    using sub_expr_t = std::conditional_t<left_directed, left_expr_t, right_expr_t>; 

    using sub_traits   = etl_traits<sub_expr_t>;   
    using left_traits  = etl_traits<left_expr_t>;  
    using right_traits = etl_traits<right_expr_t>; 

    static constexpr bool is_etl         = true;                
    static constexpr bool is_transformer = false;               
    static constexpr bool is_view        = false;               
    static constexpr bool is_magic_view  = false;               
    static constexpr bool is_fast        = sub_traits::is_fast; 
    static constexpr bool is_linear      = left_traits::is_linear && right_traits::is_linear && BinaryOp::linear; 
    static constexpr bool is_thread_safe =
        left_traits::is_thread_safe && right_traits::is_thread_safe && BinaryOp::thread_safe;     
    static constexpr bool is_value     = false;                                                   
    static constexpr bool is_direct    = false;                                                   
    static constexpr bool is_generator = left_traits::is_generator && right_traits::is_generator; 
    static constexpr bool is_temporary = left_traits::is_temporary || right_traits::is_temporary; 
    static constexpr bool is_padded    = is_linear && left_traits::is_padded && right_traits::is_padded;   
    static constexpr bool is_aligned   = is_linear && left_traits::is_aligned && right_traits::is_aligned; 
    static constexpr order storage_order =
        left_traits::is_generator ? right_traits::storage_order : left_traits::storage_order; 

    static constexpr bool gpu_computable = all_gpu_computable<LE, RE> && BinaryOp::template gpu_computable<LE, RE> && all_homogeneous<LE, RE>;

    template <vector_mode_t V>
    static constexpr bool vectorizable =
        all_homogeneous<LE, RE>&& left_traits::template vectorizable<V>&& right_traits::template vectorizable<V>&& BinaryOp::template vectorizable<V>;

    static constexpr auto& get(const expr_t& v) {
        if constexpr (left_directed) {
            return v.lhs;
        } else {
            return v.rhs;
        }
    }

    static size_t size(const expr_t& v) {
        return sub_traits::size(get(v));
    }

    static size_t dim(const expr_t& v, size_t d) {
        return sub_traits::dim(get(v), d);
    }

    static constexpr size_t size() {
        return sub_traits::size();
    }

    template <size_t D>
    static constexpr size_t dim() {
        return sub_traits::template dim<D>();
    }

    static constexpr size_t dimensions() {
        return sub_traits::dimensions();
    }

    static constexpr int complexity() noexcept {
        return BinaryOp::complexity() + left_traits::complexity() + right_traits::complexity();
    }
};

} //end of namespace etl