xaxis_slice_iterator.hpp

/***************************************************************************
 * Copyright (c) Johan Mabille, Sylvain Corlay and Wolf Vollprecht          *
 * Copyright (c) QuantStack                                                 *
 *                                                                          *
 * Distributed under the terms of the BSD 3-Clause License.                 *
 *                                                                          *
 * The full license is in the file LICENSE, distributed with this software. *
 ****************************************************************************/
 
#ifndef XTENSOR_AXIS_SLICE_ITERATOR_HPP
#define XTENSOR_AXIS_SLICE_ITERATOR_HPP
 
#include "../views/xstrided_view.hpp"
 
namespace xt
{

    template <class CT>
    class xaxis_slice_iterator
    {
    public:
 
        using self_type = xaxis_slice_iterator<CT>;
 
        using xexpression_type = std::decay_t<CT>;
        using size_type = typename xexpression_type::size_type;
        using difference_type = typename xexpression_type::difference_type;
        using shape_type = typename xexpression_type::shape_type;
        using strides_type = typename xexpression_type::strides_type;
        using value_type = xstrided_view<CT, shape_type>;
        using reference = std::remove_reference_t<xtl::apply_cv_t<CT, value_type>>;
        using pointer = xtl::xclosure_pointer<std::remove_reference_t<xtl::apply_cv_t<CT, value_type>>>;
 
        using iterator_category = std::forward_iterator_tag;
 
        template <class CTA>
        xaxis_slice_iterator(CTA&& e, size_type axis);
        template <class CTA>
        xaxis_slice_iterator(CTA&& e, size_type axis, size_type index, size_type offset);
 
        self_type& operator++();
        self_type operator++(int);
 
        reference operator*() const;
        pointer operator->() const;
 
        bool equal(const self_type& rhs) const;
 
    private:
 
        using storing_type = xtl::ptr_closure_type_t<CT>;
        mutable storing_type p_expression;
        size_type m_index;
        size_type m_offset;
        size_type m_axis_stride;
        size_type m_lower_shape;
        size_type m_upper_shape;
        size_type m_iter_size;
        bool m_is_target_axis;
        value_type m_sv;
 
        template <class T, class CTA>
        T get_storage_init(CTA&& e) const;
    };
 
    template <class CT>
    bool operator==(const xaxis_slice_iterator<CT>& lhs, const xaxis_slice_iterator<CT>& rhs);
 
    template <class CT>
    bool operator!=(const xaxis_slice_iterator<CT>& lhs, const xaxis_slice_iterator<CT>& rhs);
 
    template <class E>
    auto xaxis_slice_begin(E&& e);
 
    template <class E>
    auto xaxis_slice_begin(E&& e, typename std::decay_t<E>::size_type axis);
 
    template <class E>
    auto xaxis_slice_end(E&& e);
 
    template <class E>
    auto xaxis_slice_end(E&& e, typename std::decay_t<E>::size_type axis);
 
    /***************************************
     * xaxis_slice_iterator implementation *
     ***************************************/
 
    template <class CT>
    template <class T, class CTA>
    T xaxis_slice_iterator<CT>::get_storage_init(CTA&& e) const
    {
        if constexpr (xtl::pointer_concept<T>)
        {
            return &e;
        }
        else
        {
            return e;
        }
    }


    template <class CT>
    template <class CTA>
    inline xaxis_slice_iterator<CT>::xaxis_slice_iterator(CTA&& e, size_type axis)
        : xaxis_slice_iterator(std::forward<CTA>(e), axis, 0, e.data_offset())
    {
    }

    template <class CT>
    template <class CTA>
    inline xaxis_slice_iterator<CT>::xaxis_slice_iterator(CTA&& e, size_type axis, size_type index, size_type offset)
        : p_expression(get_storage_init<storing_type>(std::forward<CTA>(e)))
        , m_index(index)
        , m_offset(offset)
        , m_axis_stride(static_cast<size_type>(e.strides()[axis]) * (e.shape()[axis] - 1u))
        , m_lower_shape(0)
        , m_upper_shape(0)
        , m_iter_size(0)
        , m_is_target_axis(false)
        , m_sv(strided_view(
              std::forward<CT>(e),
              std::forward<shape_type>({e.shape()[axis]}),
              std::forward<strides_type>({e.strides()[axis]}),
              offset,
              e.layout()
          ))
    {
        if (e.layout() == layout_type::row_major)
        {
            m_is_target_axis = axis == e.dimension() - 1;
            m_lower_shape = std::accumulate(
                e.shape().begin() + axis + 1,
                e.shape().end(),
                size_t(1),
                std::multiplies<>()
            );
            m_iter_size = std::accumulate(e.shape().begin() + 1, e.shape().end(), size_t(1), std::multiplies<>());
        }
        else
        {
            m_is_target_axis = axis == 0;
            m_lower_shape = std::accumulate(
                e.shape().begin(),
                e.shape().begin() + axis,
                size_t(1),
                std::multiplies<>()
            );
            m_iter_size = std::accumulate(e.shape().begin(), e.shape().end() - 1, size_t(1), std::multiplies<>());
        }
        m_upper_shape = m_lower_shape + m_axis_stride;
    }



    template <class CT>
    inline auto xaxis_slice_iterator<CT>::operator++() -> self_type&
    {
        ++m_index;
        ++m_offset;
        auto index_compare = (m_offset % m_iter_size);
        if (m_is_target_axis || (m_upper_shape >= index_compare && index_compare >= m_lower_shape))
        {
            m_offset += m_axis_stride;
        }
        m_sv.set_offset(m_offset);
        return *this;
    }

    template <class CT>
    inline auto xaxis_slice_iterator<CT>::operator++(int) -> self_type
    {
        self_type tmp(*this);
        ++(*this);
        return tmp;
    }



    template <class CT>
    inline auto xaxis_slice_iterator<CT>::operator*() const -> reference
    {
        return m_sv;
    }

    template <class CT>
    inline auto xaxis_slice_iterator<CT>::operator->() const -> pointer
    {
        return xtl::closure_pointer(operator*());
    }

 
    /*
     * @name Comparisons
     */

    template <class CT>
    inline bool xaxis_slice_iterator<CT>::equal(const self_type& rhs) const
    {
        return p_expression == rhs.p_expression && m_index == rhs.m_index;
    }

    template <class CT>
    inline bool operator==(const xaxis_slice_iterator<CT>& lhs, const xaxis_slice_iterator<CT>& rhs)
    {
        return lhs.equal(rhs);
    }

    template <class CT>
    inline bool operator!=(const xaxis_slice_iterator<CT>& lhs, const xaxis_slice_iterator<CT>& rhs)
    {
        return !(lhs == rhs);
    }



    template <class E>
    inline auto axis_slice_begin(E&& e)
    {
        using return_type = xaxis_slice_iterator<xtl::closure_type_t<E>>;
        return return_type(std::forward<E>(e), 0);
    }

    template <class E>
    inline auto axis_slice_begin(E&& e, typename std::decay_t<E>::size_type axis)
    {
        using return_type = xaxis_slice_iterator<xtl::closure_type_t<E>>;
        return return_type(std::forward<E>(e), axis, 0, e.data_offset());
    }

    template <class E>
    inline auto axis_slice_end(E&& e)
    {
        using return_type = xaxis_slice_iterator<xtl::closure_type_t<E>>;
        return return_type(
            std::forward<E>(e),
            0,
            std::accumulate(e.shape().begin() + 1, e.shape().end(), size_t(1), std::multiplies<>()),
            e.size()
        );
    }

    template <class E>
    inline auto axis_slice_end(E&& e, typename std::decay_t<E>::size_type axis)
    {
        using return_type = xaxis_slice_iterator<xtl::closure_type_t<E>>;
        auto index_sum = std::accumulate(
            e.shape().begin(),
            e.shape().begin() + axis,
            size_t(1),
            std::multiplies<>()
        );
        return return_type(
            std::forward<E>(e),
            axis,
            std::accumulate(e.shape().begin() + axis + 1, e.shape().end(), index_sum, std::multiplies<>()),
            e.size() + axis
        );
    }

}
 
#endif