Fix xsort calls over axis

Author: AntoinePrv

include/xtensor/xsort.hpp

@@ -46,35 +46,39 @@ namespace xt
             return stride != 0 ? stride : static_cast<std::ptrdiff_t>(shape);
         }
 
-        template <class E, class F>
-        inline void call_over_leading_axis(E& ev, F&& fct)
+        template <class E>
+        inline std::ptrdiff_t get_secondary_stride(const E& ev)
         {
-            std::size_t n_iters = 1;
-            std::ptrdiff_t secondary_stride;
+            if (ev.layout() == layout_type::row_major)
+            {
+                return adjust_secondary_stride(ev.strides()[ev.dimension() - 2], *(ev.shape().end() - 1));
+            }
+
+            return adjust_secondary_stride(ev.strides()[1], *(ev.shape().begin()));
+        }
 
+        template <class E>
+        inline std::size_t leading_axis_n_iters(const E& ev)
+        {
             if (ev.layout() == layout_type::row_major)
             {
-                n_iters = std::accumulate(
+                return std::accumulate(
                     ev.shape().begin(),
                     ev.shape().end() - 1,
                     std::size_t(1),
                     std::multiplies<>()
                 );
-                secondary_stride = adjust_secondary_stride(
-                    ev.strides()[ev.dimension() - 2],
-                    *(ev.shape().end() - 1)
-                );
-            }
-            else
-            {
-                n_iters = std::accumulate(
-                    ev.shape().begin() + 1,
-                    ev.shape().end(),
-                    std::size_t(1),
-                    std::multiplies<>()
-                );
-                secondary_stride = adjust_secondary_stride(ev.strides()[1], *(ev.shape().begin()));
             }
+            return std::accumulate(ev.shape().begin() + 1, ev.shape().end(), std::size_t(1), std::multiplies<>());
+        }
+
+        template <class E, class F>
+        inline void call_over_leading_axis(E& ev, F&& fct)
+        {
+            XTENSOR_ASSERT(ev.dimension() >= 2);
+
+            const std::size_t n_iters = leading_axis_n_iters(ev);
+            const std::ptrdiff_t secondary_stride = get_secondary_stride(ev);
 
             const auto begin = ev.data();
             const auto end = begin + n_iters * secondary_stride;
@@ -87,37 +91,13 @@ namespace xt
         template <class E1, class E2, class F>
         inline void call_over_leading_axis(E1& e1, E2& e2, F&& fct)
         {
-            std::size_t n_iters = 1;
-            std::ptrdiff_t secondary_stride1, secondary_stride2;
+            XTENSOR_ASSERT(e1.dimension() >= 2);
+            XTENSOR_ASSERT(e1.dimension() == e2.dimension());
 
-            if (e1.layout() == layout_type::row_major)
-            {
-                n_iters = std::accumulate(
-                    e1.shape().begin(),
-                    e1.shape().end() - 1,
-                    std::size_t(1),
-                    std::multiplies<>()
-                );
-                secondary_stride1 = adjust_secondary_stride(
-                    e1.strides()[e1.dimension() - 2],
-                    *(e1.shape().end() - 1)
-                );
-                secondary_stride2 = adjust_secondary_stride(
-                    e2.strides()[e2.dimension() - 2],
-                    *(e2.shape().end() - 2)
-                );
-            }
-            else
-            {
-                n_iters = std::accumulate(
-                    e1.shape().begin() + 1,
-                    e1.shape().end(),
-                    std::size_t(1),
-                    std::multiplies<>()
-                );
-                secondary_stride1 = adjust_secondary_stride(e1.strides()[1], *(e1.shape().begin()));
-                secondary_stride2 = adjust_secondary_stride(e2.strides()[1], *(e2.shape().begin()));
-            }
+            const std::size_t n_iters = leading_axis_n_iters(e1);
+            std::ptrdiff_t const secondary_stride1 = get_secondary_stride(e1);
+            std::ptrdiff_t const secondary_stride2 = get_secondary_stride(e2);
+            XTENSOR_ASSERT(secondary_stride1 == secondary_stride2);
 
             const auto begin1 = e1.data();
             const auto end1 = begin1 + n_iters * secondary_stride1;
@@ -192,7 +172,7 @@ namespace xt
             }
 
             dynamic_shape<std::size_t> permutation, reverse_permutation;
-            std::tie(permutation, reverse_permutation) = get_permutations(e.dimension(), axis, e.layout());
+            std::tie(permutation, reverse_permutation) = get_permutations(e.dimension(), ax, e.layout());
             R res = transpose(e, permutation);
             detail::call_over_leading_axis(res, std::forward<F>(lambda));
             res = transpose(res, reverse_permutation);

test/test_xsort.cpp

@@ -426,15 +426,16 @@ namespace xt
 
         SUBCASE("complex")
         {
-            xt::xarray<double> a = {
+            const xt::xarray<double> data = {
                 1014., 1017., 1019., 1020., 1023., 1026., 1026., 1028., 1030., 1032., 1039., 1047., 1071.,
                 927.,  932.,  935.,  943.,  944.,  944.,  945.,  948.,  952.,  962.,  968.,  968.,  969.,
                 969.,  974.,  981.,  993.,  994.,  994.,  1003., 1007., 1008., 1008., 1012., 1013., 1014.,
                 1080., 1085., 1088., 1111., 1112., 1117., 1119., 1128., 1130., 1209., 1309., 1426.};
-            xt::xtensor_fixed<std::size_t, xt::xshape<4>> kth = {17, 32, 18, 33};
+            const xt::xtensor_fixed<std::size_t, xt::xshape<4>> kth = {17, 32, 18, 33};
 
             SUBCASE("1D")
             {
+                const auto& a = data;
                 const auto argpart = xt::argpartition(a, kth);
                 for (std::size_t k : kth)
                 {
@@ -445,6 +446,7 @@ namespace xt
 
             SUBCASE("2D")
             {
+                auto a = data;
                 a.reshape({1, a.size()});
                 const auto argpart = xt::argpartition(a, kth, 1);
                 for (std::size_t k : kth)