index_mapper.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_INDEX_MAPPER_HPP
#define XTENSOR_INDEX_MAPPER_HPP
 
#include "xview.hpp"
 
namespace xt
{
 
    template <class UndefinedView>
    struct index_mapper;

    enum class access_t
    {
        SAFE,   
        UNSAFE  
    };

    template <class UnderlyingContainer, class... Slices>
    class index_mapper<xt::xview<UnderlyingContainer, Slices...>>
    {
    public:

        using view_type = xt::xview<UnderlyingContainer, Slices...>;

        using reference = typename xt::xview<UnderlyingContainer, Slices...>::reference;

        using const_reference = typename xt::xview<UnderlyingContainer, Slices...>::const_reference;

        static constexpr size_t n_slices = sizeof...(Slices);

        static constexpr size_t nb_integral_slices = (std::is_integral_v<Slices> + ...);

        static constexpr size_t nb_new_axis_slices = (xt::detail::is_newaxis_v<Slices> + ...);

        template <std::integral... Indices>
        static constexpr size_t n_indices_full_v = size_t(sizeof...(Indices) + nb_integral_slices);

        template <std::integral... Indices>
        reference map(UnderlyingContainer& container, const view_type& view, const Indices... indices) const;

        template <std::integral... Indices>
        const_reference
        cmap(const UnderlyingContainer& container, const view_type& view, const Indices... indices) const;

        template <std::integral... Indices>
        reference map_at(UnderlyingContainer& container, const view_type& view, const Indices... indices) const;

        template <std::integral... Indices>
        const_reference
        cmap_at(const UnderlyingContainer& container, const view_type& view, const Indices... indices) const;

        size_t dimension(const UnderlyingContainer& container) const;
 
    private:

        template <bool IS_CONST>
        using conditional_reference = std::conditional_t<IS_CONST, const_reference, reference>;

        template <size_t I>
        using slice_type = std::tuple_element_t<I, std::tuple<Slices...>>;

        template <size_t I>
        static consteval bool is_slice_integral();

        template <size_t I>
        static consteval bool is_slice_new_axis();

        template <size_t first, size_t bound, size_t... indices>
        struct indices_sequence_helper
        {
            // we add the current axis
            using not_new_axis_type = typename indices_sequence_helper<first + 1, bound, indices..., first>::type;
 
            // we skip the current axis
            using new_axis_type = typename indices_sequence_helper<first + 1, bound, indices...>::type;
 
            // NOTE: is_slice_new_axis works even if first >= sizeof...(Slices)
            using type = std::conditional_t<is_slice_new_axis<first>(), new_axis_type, not_new_axis_type>;
        };

        template <size_t bound, size_t... indices>
        struct indices_sequence_helper<bound, bound, indices...>
        {
            using type = std::index_sequence<indices...>;
        };

        template <size_t bound>
        using indices_sequence = indices_sequence_helper<0, bound>::type;

        template <size_t I, std::integral Index>
        size_t map_ith_index(const view_type& view, const Index i) const;

        template <bool IS_CONST, access_t ACCESS, std::integral FirstIndice, std::integral... OtherIndices>
        conditional_reference<IS_CONST> map_main(
            std::bool_constant<IS_CONST> /* is_const */,
            std::conditional_t<IS_CONST, const UnderlyingContainer&, UnderlyingContainer&> container,
            std::integral_constant<access_t, ACCESS> /* access */,
            const view_type& view,
            const FirstIndice firstIndice,
            const OtherIndices... otherIndices
        ) const;

        template <bool IS_CONST, access_t ACCESS>
        conditional_reference<IS_CONST> map_main(
            std::bool_constant<IS_CONST> /* is_const */,
            std::conditional_t<IS_CONST, const UnderlyingContainer&, UnderlyingContainer&> container,
            std::integral_constant<access_t, ACCESS> /* access */,
            const view_type& view
        ) const;

        template <bool IS_CONST, access_t ACCESS, size_t n_indices, size_t... Is>
        conditional_reference<IS_CONST> map_all_indices(
            std::bool_constant<IS_CONST> /* is_const */,
            std::conditional_t<IS_CONST, const UnderlyingContainer&, UnderlyingContainer&> container,
            std::integral_constant<access_t, ACCESS> /* access */,
            const view_type& view,
            std::index_sequence<Is...> /* is_seq */,
            const std::array<size_t, n_indices>& indices
        ) const;

        template <std::integral... Indices>
        std::array<size_t, n_indices_full_v<Indices...>> get_indices_full(const Indices... indices) const;
    };
 
    /*******************************
     * index_mapper implementation *
     *******************************/
 
    template <class UnderlyingContainer, class... Slices>
    template <size_t I>
    consteval bool index_mapper<xt::xview<UnderlyingContainer, Slices...>>::is_slice_integral()
    {
        if constexpr (I < sizeof...(Slices))
        {
            return std::is_integral_v<slice_type<I>>;
        }
        else
        {
            return false;
        }
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <size_t I>
    consteval bool index_mapper<xt::xview<UnderlyingContainer, Slices...>>::is_slice_new_axis()
    {
        if constexpr (I < sizeof...(Slices))
        {
            return xt::detail::is_newaxis_v<slice_type<I>>;
        }
        else
        {
            return false;
        }
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <std::integral... Indices>
    auto
    index_mapper<xt::xview<UnderlyingContainer, Slices...>>::get_indices_full(const Indices... indices) const
        -> std::array<size_t, n_indices_full_v<Indices...>>
    {
        constexpr size_t n_indices_full = n_indices_full_v<Indices...>;
 
        std::array<size_t, sizeof...(indices)> args{size_t(indices)...};
        std::array<size_t, n_indices_full> args_full;
 
        const auto fill_args_full = [&args_full, &args]<size_t... Is>(std::index_sequence<Is...>)
        {
            auto it = std::cbegin(args);
 
            ((args_full[Is] = (is_slice_integral<Is>()) ? size_t(0) : *it++), ...);
        };
 
        fill_args_full(std::make_index_sequence<n_indices_full>{});
 
        return args_full;
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <std::integral... Indices>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::map(
        UnderlyingContainer& container,
        const view_type& view,
        const Indices... indices
    ) const -> reference
    {
        return map_main(
            std::false_type{},
            container,
            std::integral_constant<access_t, access_t::UNSAFE>{},
            view,
            indices...
        );
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <std::integral... Indices>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::cmap(
        const UnderlyingContainer& container,
        const view_type& view,
        const Indices... indices
    ) const -> const_reference
    {
        return map_main(
            std::true_type{},
            container,
            std::integral_constant<access_t, access_t::UNSAFE>{},
            view,
            indices...
        );
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <std::integral... Indices>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::map_at(
        UnderlyingContainer& container,
        const view_type& view,
        const Indices... indices
    ) const -> reference
    {
        return map_main(
            std::false_type{},
            container,
            std::integral_constant<access_t, access_t::SAFE>{},
            view,
            indices...
        );
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <std::integral... Indices>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::cmap_at(
        const UnderlyingContainer& container,
        const view_type& view,
        const Indices... indices
    ) const -> const_reference
    {
        return map_main(
            std::true_type{},
            container,
            std::integral_constant<access_t, access_t::SAFE>{},
            view,
            indices...
        );
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <bool IS_CONST, access_t ACCESS, std::integral FirstIndice, std::integral... OtherIndices>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::map_main(
        std::bool_constant<IS_CONST> is_const,
        std::conditional_t<IS_CONST, const UnderlyingContainer&, UnderlyingContainer&> container,
        std::integral_constant<access_t, ACCESS> access,
        const view_type& view,
        const FirstIndice firstIndice,
        const OtherIndices... otherIndices
    ) const -> conditional_reference<IS_CONST>
    {
        constexpr size_t n_indices_full = n_indices_full_v<FirstIndice, OtherIndices...>;
 
        constexpr size_t underlying_n_dimensions = xt::static_dimension<
            typename std::decay_t<UnderlyingContainer>::shape_type>::value;
 
        // If there is too many indices, we need to drop the first ones.
        // If the number of dimensions of the underlying container is known at compile time we can drop them
        // at compile time Else a runtime-test is requires, which, breaks vectorization.
        // I don't know if we can do it in another way.
 
        if constexpr (underlying_n_dimensions != size_t(-1))
        {
            // the number of dimensions of the underlying container is known at compile time.
            constexpr size_t n_dimensions = underlying_n_dimensions - nb_integral_slices + nb_new_axis_slices;
 
            // we can perform compile time checks
            if constexpr (1 + sizeof...(OtherIndices) > n_dimensions)
            {
                return map_main(is_const, container, access, view, otherIndices...);
            }
            else
            {
                return map_all_indices(
                    is_const,
                    container,
                    access,
                    view,
                    indices_sequence<n_indices_full>{},
                    get_indices_full(firstIndice, otherIndices...)
                );
            }
        }
        else
        {
            // we need execution time checks
            if (1 + sizeof...(OtherIndices) > dimension(container))
            {
                return map_main(is_const, container, access, view, otherIndices...);
            }
            else
            {
                return map_all_indices(
                    is_const,
                    container,
                    access,
                    view,
                    indices_sequence<n_indices_full>{},
                    get_indices_full(firstIndice, otherIndices...)
                );
            }
        }
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <bool IS_CONST, access_t ACCESS>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::map_main(
        std::bool_constant<IS_CONST> is_const,
        std::conditional_t<IS_CONST, const UnderlyingContainer&, UnderlyingContainer&> container,
        std::integral_constant<access_t, ACCESS> access,
        const view_type& view
    ) const -> conditional_reference<IS_CONST>
    {
        constexpr size_t n_indices_full = n_indices_full_v<>;
 
        return map_all_indices(
            is_const,
            container,
            access,
            view,
            indices_sequence<n_indices_full>{},
            get_indices_full()
        );
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <bool IS_CONST, access_t ACCESS, size_t n_indices, size_t... Is>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::map_all_indices(
        std::bool_constant<IS_CONST> /* is_const */,
        std::conditional_t<IS_CONST, const UnderlyingContainer&, UnderlyingContainer&> container,
        std::integral_constant<access_t, ACCESS> /* access */,
        const view_type& view,
        std::index_sequence<Is...> /* is_seq */,
        const std::array<size_t, n_indices>& indices
    ) const -> conditional_reference<IS_CONST>
    {
        if constexpr (ACCESS == access_t::SAFE)
        {
            return container.at(map_ith_index<Is>(view, indices[Is])...);
        }
        else
        {
            return container(map_ith_index<Is>(view, indices[Is])...);
        }
    }
 
    template <class UnderlyingContainer, class... Slices>
    template <size_t I, std::integral Index>
    auto
    index_mapper<xt::xview<UnderlyingContainer, Slices...>>::map_ith_index(const view_type& view, const Index i) const
        -> size_t
    {
        if constexpr (I < sizeof...(Slices))
        {
            // if the slice is explicitly specified, use it
            using current_slice = std::tuple_element_t<I, std::tuple<Slices...>>;
 
            static_assert(not xt::detail::is_newaxis_v<current_slice>);
 
            const auto& slice = std::get<I>(view.slices());
 
            if constexpr (std::is_integral_v<current_slice>)
            {
                assert(i == 0);
                return size_t(slice);
            }
            else
            {
                assert(i < slice.size());
                return size_t(slice(i));
            }
        }
        else
        {
            // else assume xt::all
            return i;
        }
    }
 
    template <class UnderlyingContainer, class... Slices>
    auto index_mapper<xt::xview<UnderlyingContainer, Slices...>>::dimension(const UnderlyingContainer& container
    ) const -> size_t
    {
        return container.dimension() - nb_integral_slices + nb_new_axis_slices;
    }
 
}  // namespace xt
 
#endif  // XTENSOR_INDEX_MAPPER_HPP