rep_transformers.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, typename D>
struct rep_transformer {
    using derived_t  = D;          
    using sub_type   = T;          
    using value_type = value_t<T>; 

    static constexpr bool gpu_computable = false;

protected:
    sub_type sub; 

public:
    explicit rep_transformer(sub_type expr) : sub(expr) {}

    template <typename... Sizes>
    value_type operator()(Sizes... args) const {
        return as_derived().selected_only(make_index_range<derived_t::dim_start, derived_t::dim_end>(), args...);
    }

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

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(visitor));
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

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

private:
    const derived_t& as_derived() const noexcept {
        return *static_cast<const derived_t*>(this);
    }
};

template <typename T, size_t... D>
struct rep_r_transformer : rep_transformer<T, rep_r_transformer<T, D...>> {
    using this_type  = rep_r_transformer<T, D...>;     
    using base_type  = rep_transformer<T, this_type>;  
    using sub_type   = typename base_type::sub_type;   
    using value_type = typename base_type::value_type; 

    static constexpr bool gpu_computable = false;

private:
    static constexpr size_t sub_d      = decay_traits<sub_type>::dimensions(); 
    static constexpr size_t dimensions = sizeof...(D) + sub_d;                 
    static constexpr size_t dim_start  = 0;                                    
    static constexpr size_t dim_end    = sub_d;                                

    friend struct rep_transformer<T, rep_r_transformer>;
    friend struct etl_traits<rep_r_transformer>;

public:
    explicit rep_r_transformer(sub_type expr) : base_type(expr) {}

    value_type operator[](size_t i) const {
        return this->sub[i / (D * ...)];
    }

    value_type read_flat(size_t i) const noexcept {
        return this->sub.read_flat(i / (D * ...));
    }

    template <typename... Sizes, size_t... I>
    value_type selected_only(const std::index_sequence<I...>& /*seq*/, Sizes... sizes) const {
        return this->sub(cpp::nth_value<I>(sizes...)...);
    }

    friend std::ostream& operator<<(std::ostream& os, const rep_r_transformer& transformer) {
        return os << "rep_r[" << concat_sizes(D...) << "](" << transformer.sub << ")";
    }
};

template <typename T, size_t... D>
struct rep_l_transformer : rep_transformer<T, rep_l_transformer<T, D...>> {
    using this_type  = rep_l_transformer<T, D...>;     
    using base_type  = rep_transformer<T, this_type>;  
    using sub_type   = typename base_type::sub_type;   
    using value_type = typename base_type::value_type; 

    static constexpr bool gpu_computable = false;

private:
    static constexpr size_t sub_d      = decay_traits<sub_type>::dimensions(); 
    static constexpr size_t dimensions = sizeof...(D) + sub_d;                 
    static constexpr size_t dim_start  = sizeof...(D);                         
    static constexpr size_t dim_end    = dimensions;                           

    friend struct rep_transformer<T, rep_l_transformer>;
    friend struct etl_traits<rep_l_transformer>;

public:
    explicit rep_l_transformer(sub_type expr) : base_type(expr) {}

    value_type operator[](size_t i) const {
        return this->sub[i % etl::size(this->sub)];
    }

    value_type read_flat(size_t i) const noexcept {
        return this->sub.read_flat(i % etl::size(this->sub));
    }

    template <typename... Sizes, size_t... I>
    value_type selected_only(const std::index_sequence<I...>& /*seq*/, Sizes... sizes) const {
        return this->sub(cpp::nth_value<I>(sizes...)...);
    }

    friend std::ostream& operator<<(std::ostream& os, const rep_l_transformer& transformer) {
        return os << "rep_l[" << concat_sizes(D...) << "](" << transformer.sub << ")";
    }
};

template <typename T, size_t D>
struct dyn_rep_r_transformer : rep_transformer<T, dyn_rep_r_transformer<T, D>> {
    using this_type  = dyn_rep_r_transformer<T, D>;    
    using base_type  = rep_transformer<T, this_type>;  
    using sub_type   = typename base_type::sub_type;   
    using value_type = typename base_type::value_type; 

    static constexpr bool gpu_computable = false;

private:
    static constexpr size_t sub_d      = decay_traits<sub_type>::dimensions(); 
    static constexpr size_t dimensions = D + sub_d;                            
    static constexpr size_t dim_start  = 0;                                    
    static constexpr size_t dim_end    = sub_d;                                

    std::array<size_t, D> reps; 
    size_t m;                   

    friend struct rep_transformer<T, dyn_rep_r_transformer>;
    friend struct etl_traits<dyn_rep_r_transformer>;

public:
    dyn_rep_r_transformer(sub_type expr, std::array<size_t, D> reps_a) : base_type(expr), reps(reps_a) {
        m = std::accumulate(reps.begin(), reps.end(), 1UL, [](size_t a, size_t b) { return a * b; });
    }

    value_type operator[](size_t i) const {
        return this->sub[i / m];
    }

    value_type read_flat(size_t i) const noexcept {
        return this->sub.read_flat(i / m);
    }

    template <typename... Sizes, size_t... I>
    value_type selected_only(const std::index_sequence<I...>& /*seq*/, Sizes... sizes) const {
        return this->sub(cpp::nth_value<I>(sizes...)...);
    }
};

template <typename T, size_t D>
struct dyn_rep_l_transformer : rep_transformer<T, dyn_rep_l_transformer<T, D>> {
    using this_type  = dyn_rep_l_transformer<T, D>;    
    using base_type  = rep_transformer<T, this_type>;  
    using sub_type   = typename base_type::sub_type;   
    using value_type = typename base_type::value_type; 

    static constexpr bool gpu_computable = false;

private:
    static constexpr size_t sub_d      = decay_traits<sub_type>::dimensions(); 
    static constexpr size_t dimensions = D + sub_d;                            
    static constexpr size_t dim_start  = D;                                    
    static constexpr size_t dim_end    = dimensions;                           

    std::array<size_t, D> reps; 
    size_t m;                   

    friend struct rep_transformer<T, dyn_rep_l_transformer>;
    friend struct etl_traits<dyn_rep_l_transformer>;

public:
    dyn_rep_l_transformer(sub_type expr, std::array<size_t, D> reps_a) : base_type(expr), reps(reps_a) {
        m = std::accumulate(reps.begin(), reps.end(), 1UL, [](size_t a, size_t b) { return a * b; });
    }

    value_type operator[](size_t i) const {
        return this->sub[i % etl::size(this->sub)];
    }

    value_type read_flat(size_t i) const {
        return this->sub.read_flat(i % etl::size(this->sub));
    }

    template <typename... Sizes, size_t... I>
    value_type selected_only(const std::index_sequence<I...>& /*seq*/, Sizes... sizes) const {
        return this->sub(cpp::nth_value<I>(sizes...)...);
    }
};

template <typename T, size_t... D>
struct etl_traits<rep_r_transformer<T, D...>> {
    using expr_t     = etl::rep_r_transformer<T, D...>;             
    using sub_expr_t = std::decay_t<T>;                             
    using value_type = typename etl_traits<sub_expr_t>::value_type; 

    static constexpr bool is_etl         = true;                                   
    static constexpr bool is_transformer = true;                                   
    static constexpr bool is_view        = false;                                  
    static constexpr bool is_magic_view  = false;                                  
    static constexpr bool is_fast        = etl_traits<sub_expr_t>::is_fast;        
    static constexpr bool is_linear      = false;                                  
    static constexpr bool is_thread_safe = etl_traits<sub_expr_t>::is_thread_safe; 
    static constexpr bool is_value       = false;                                  
    static constexpr bool is_direct      = false;                                  
    static constexpr bool is_generator   = false;                                  
    static constexpr bool is_padded      = false;                                  
    static constexpr bool is_aligned     = false;                                  
    static constexpr bool is_temporary   = etl_traits<sub_expr_t>::is_temporary;   
    static constexpr bool gpu_computable = false;                                  
    static constexpr order storage_order = etl_traits<sub_expr_t>::storage_order;  

    static constexpr size_t sub_d = etl_traits<sub_expr_t>::dimensions(); 

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    static size_t size(const expr_t& v) {
        return (D * ...) * etl_traits<sub_expr_t>::size(v.sub);
    }

    static size_t dim(const expr_t& v, size_t d) {
        if (d < sub_d) {
            return etl_traits<sub_expr_t>::dim(v.sub, d);
        } else {
            return dyn_nth_size<D...>(d - sub_d);
        }
    }

    static constexpr size_t size() {
        return (D * ...) * etl_traits<sub_expr_t>::size();
    }

    template <size_t D2>
    static constexpr size_t dim() {
        if constexpr (D2 < sub_d) {
            return etl_traits<sub_expr_t>::template dim<D2>();
        } else {
            return nth_size<D2 - sub_d, 0, D...>();
        }
    }

    static constexpr size_t dimensions() {
        return sizeof...(D) + etl_traits<sub_expr_t>::dimensions();
    }

    static constexpr int complexity() noexcept {
        return 3;
    }
};

template <typename T, size_t... D>
struct etl_traits<rep_l_transformer<T, D...>> {
    using expr_t     = etl::rep_l_transformer<T, D...>;             
    using sub_expr_t = std::decay_t<T>;                             
    using value_type = typename etl_traits<sub_expr_t>::value_type; 

    static constexpr bool is_etl         = true;                                   
    static constexpr bool is_transformer = true;                                   
    static constexpr bool is_view        = false;                                  
    static constexpr bool is_magic_view  = false;                                  
    static constexpr bool is_fast        = etl_traits<sub_expr_t>::is_fast;        
    static constexpr bool is_linear      = false;                                  
    static constexpr bool is_thread_safe = etl_traits<sub_expr_t>::is_thread_safe; 
    static constexpr bool is_value       = false;                                  
    static constexpr bool is_direct      = false;                                  
    static constexpr bool is_generator   = false;                                  
    static constexpr bool is_padded      = false;                                  
    static constexpr bool is_aligned     = false;                                  
    static constexpr bool is_temporary   = etl_traits<sub_expr_t>::is_temporary;   
    static constexpr bool gpu_computable = false;                                  
    static constexpr order storage_order = etl_traits<sub_expr_t>::storage_order;  

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    static size_t size(const expr_t& v) {
        return (D * ...) * etl_traits<sub_expr_t>::size(v.sub);
    }

    static size_t dim(const expr_t& v, size_t d) {
        if (d >= sizeof...(D)) {
            return etl_traits<sub_expr_t>::dim(v.sub, d - sizeof...(D));
        } else {
            return dyn_nth_size<D...>(d);
        }
    }

    static constexpr size_t size() {
        return (D * ...) * etl_traits<sub_expr_t>::size();
    }

    template <size_t D2>
    static constexpr size_t dim() {
        if constexpr (D2 >= sizeof...(D)) {
            return etl_traits<sub_expr_t>::template dim<D2 - sizeof...(D)>();
        } else {
            return nth_size<D2, 0, D...>();
        }
    }

    static constexpr size_t dimensions() {
        return sizeof...(D) + etl_traits<sub_expr_t>::dimensions();
    }

    static constexpr int complexity() noexcept {
        return 3;
    }
};

template <typename T, size_t D>
struct etl_traits<dyn_rep_r_transformer<T, D>> {
    using expr_t     = etl::dyn_rep_r_transformer<T, D>;            
    using sub_expr_t = std::decay_t<T>;                             
    using value_type = typename etl_traits<sub_expr_t>::value_type; 

    static constexpr bool is_etl         = true;                                   
    static constexpr bool is_transformer = true;                                   
    static constexpr bool is_view        = false;                                  
    static constexpr bool is_magic_view  = false;                                  
    static constexpr bool is_fast        = false;                                  
    static constexpr bool is_linear      = false;                                  
    static constexpr bool is_thread_safe = etl_traits<sub_expr_t>::is_thread_safe; 
    static constexpr bool is_value       = false;                                  
    static constexpr bool is_direct      = false;                                  
    static constexpr bool is_generator   = false;                                  
    static constexpr bool is_padded      = false;                                  
    static constexpr bool is_aligned     = false;                                  
    static constexpr bool is_temporary   = etl_traits<sub_expr_t>::is_temporary;   
    static constexpr bool gpu_computable = false;                                  
    static constexpr order storage_order = etl_traits<sub_expr_t>::storage_order;  

    static constexpr size_t sub_d = etl_traits<sub_expr_t>::dimensions(); 

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    static size_t size(const expr_t& v) {
        return v.m * etl_traits<sub_expr_t>::size(v.sub);
    }

    static size_t dim(const expr_t& v, size_t d) {
        return d < sub_d ? etl_traits<sub_expr_t>::dim(v.sub, d) : v.reps[d - sub_d];
    }

    static constexpr size_t dimensions() {
        return D + etl_traits<sub_expr_t>::dimensions();
    }

    static constexpr int complexity() noexcept {
        return 3;
    }
};

template <typename T, size_t D>
struct etl_traits<dyn_rep_l_transformer<T, D>> {
    using expr_t     = etl::dyn_rep_l_transformer<T, D>;            
    using sub_expr_t = std::decay_t<T>;                             
    using value_type = typename etl_traits<sub_expr_t>::value_type; 

    static constexpr bool is_etl         = true;                                   
    static constexpr bool is_transformer = true;                                   
    static constexpr bool is_view        = false;                                  
    static constexpr bool is_magic_view  = false;                                  
    static constexpr bool is_fast        = false;                                  
    static constexpr bool is_linear      = false;                                  
    static constexpr bool is_thread_safe = etl_traits<sub_expr_t>::is_thread_safe; 
    static constexpr bool is_value       = false;                                  
    static constexpr bool is_direct      = false;                                  
    static constexpr bool is_generator   = false;                                  
    static constexpr bool is_padded      = false;                                  
    static constexpr bool is_aligned     = false;                                  
    static constexpr bool is_temporary   = etl_traits<sub_expr_t>::is_temporary;   
    static constexpr bool gpu_computable = false;                                  
    static constexpr order storage_order = etl_traits<sub_expr_t>::storage_order;  

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    static size_t size(const expr_t& v) {
        return v.m * etl_traits<sub_expr_t>::size(v.sub);
    }

    static size_t dim(const expr_t& v, size_t d) {
        return d >= D ? etl_traits<sub_expr_t>::dim(v.sub, d - D) : v.reps[d];
    }

    static constexpr size_t dimensions() {
        return D + etl_traits<sub_expr_t>::dimensions();
    }

    static constexpr int complexity() noexcept {
        return 3;
    }
};

} //end of namespace etl