Respect custom strides in scalar compound assignment

Author: CAOShurong

include/xtensor/core/xassign.hpp

@@ -252,6 +252,43 @@ namespace xt
 
     namespace detail
     {
+        template <class E>
+        inline bool is_dense_storage(const E& e)
+        {
+            if constexpr (strided_expression<E>)
+            {
+                const auto& shape = e.shape();
+                const auto& strides = e.strides();
+                using size_type = typename E::size_type;
+                size_type remaining = e.size();
+                size_type stride = 1;
+                // Custom strides can cover dense storage in any axis order.
+                while (remaining > 1)
+                {
+                    bool found = false;
+                    for (std::size_t axis = 0; axis < shape.size(); ++axis)
+                    {
+                        if (shape[axis] > 1 && std::cmp_equal(strides[axis], stride))
+                        {
+                            remaining /= shape[axis];
+                            stride *= shape[axis];
+                            found = true;
+                            break;
+                        }
+                    }
+                    if (!found)
+                    {
+                        return false;
+                    }
+                }
+                return true;
+            }
+            else
+            {
+                return false;
+            }
+        }
+
         template <class E1, class E2>
         constexpr bool linear_static_layout()
         {
@@ -528,11 +565,22 @@ namespace xt
     {
         E1& d = e1.derived_cast();
         using size_type = typename E1::size_type;
-        auto dst = d.storage().begin();
-        for (size_type i = d.size(); i > 0; --i)
+        auto assign = [&](auto dst)
         {
-            *dst = f(*dst, e2);
-            ++dst;
+            const size_type n = d.size();
+            for (size_type i = 0; i < n; ++i)
+            {
+                *dst = f(*dst, e2);
+                ++dst;
+            }
+        };
+        if (E1::contiguous_layout || d.is_contiguous() || detail::is_dense_storage(d))
+        {
+            assign(d.storage().begin());
+        }
+        else
+        {
+            assign(d.begin());
         }
     }
 

test/test_xadaptor_semantic.cpp

@@ -7,6 +7,9 @@
  * The full license is in the file LICENSE, distributed with this software. *
  ****************************************************************************/
 
+#include "xtensor/containers/xadapt.hpp"
+#include "xtensor/core/xnoalias.hpp"
+
 #include "test_common.hpp"
 #include "test_xsemantic.hpp"
 
@@ -47,6 +50,159 @@ namespace xt
 
 #define ADAPTOR_SEMANTIC_TYPES xarray_dynamic, xtensor_dynamic
 
+    TEST_CASE_TEMPLATE("scalar_assignment_with_strides", TypeParam, ADAPTOR_SEMANTIC_TYPES)
+    {
+        using shape_type = typename TypeParam::shape_type;
+        using strides_type = typename TypeParam::strides_type;
+        shape_type shape = {2, 1, 3};
+        strides_type strides;
+        std::vector<int> data;
+        std::vector<int> expected;
+        std::vector<int> doubled;
+
+        SUBCASE("padded rows")
+        {
+            strides = {4, 0, 1};
+            data = {2, 4, 6, -99, 8, 10, 12, -99};
+            expected = {4, 6, 8, -99, 10, 12, 14, -99};
+            doubled = {4, 8, 12, -99, 16, 20, 24, -99};
+        }
+        SUBCASE("padded columns")
+        {
+            strides = {1, 0, 3};
+            data = {2, 4, -99, 6, 8, -99, 10, 12, -99};
+            expected = {4, 6, -99, 8, 10, -99, 12, 14, -99};
+            doubled = {4, 8, -99, 12, 16, -99, 20, 24, -99};
+        }
+
+        auto original = data;
+        auto a = adapt(data.data(), 6, no_ownership(), shape, strides);
+        a += 2;
+        EXPECT_EQ(data, expected);
+        noalias(a) -= 2;
+        EXPECT_EQ(data, original);
+        a *= 2;
+        EXPECT_EQ(data, doubled);
+        noalias(a) /= 2;
+        EXPECT_EQ(data, original);
+        EXPECT_EQ(a.shape(), shape);
+        EXPECT_EQ(a.strides(), strides);
+    }
+
+    TEST_CASE_TEMPLATE("scalar_assignment_empty_strided", TypeParam, ADAPTOR_SEMANTIC_TYPES)
+    {
+        using shape_type = typename TypeParam::shape_type;
+        using strides_type = typename TypeParam::strides_type;
+        std::vector<int> data = {-99};
+        shape_type shape = {2, 0, 3};
+        strides_type strides = {4, 0, 1};
+        auto a = adapt(data.data(), 0, no_ownership(), shape, strides);
+        a += 2;
+        EXPECT_EQ(data, std::vector<int>{-99});
+    }
+
+    TEST_CASE_TEMPLATE("scalar_assignment_aliased_scalar", TypeParam, ADAPTOR_SEMANTIC_TYPES)
+    {
+        using shape_type = typename TypeParam::shape_type;
+        using strides_type = typename TypeParam::strides_type;
+        shape_type shape = {2, 1, 2};
+        strides_type strides;
+        std::vector<int> data;
+        std::vector<int> expected;
+        std::size_t scalar_index = 0;
+        std::size_t size = 4;
+
+        SUBCASE("contiguous")
+        {
+            strides = {2, 0, 1};
+            data = {1, 2, 3, 4};
+            expected = {2, 4, 5, 6};
+        }
+        SUBCASE("contiguous columns")
+        {
+            strides = {1, 0, 2};
+            data = {1, 2, 3, 4};
+            expected = {3, 4, 7, 8};
+            scalar_index = 1;
+        }
+        SUBCASE("contiguous axis permutation")
+        {
+            shape = {2, 3, 2};
+            strides = {1, 4, 2};
+            data = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
+            expected = {8, 9, 10, 11, 12, 13, 14, 22, 23, 24, 25, 26};
+            scalar_index = 6;
+            size = 12;
+        }
+        SUBCASE("padded")
+        {
+            strides = {3, 0, 1};
+            data = {1, 2, -99, 3, 4, -99};
+            expected = {2, 4, -99, 5, 6, -99};
+        }
+
+        auto a = adapt(data.data(), size, no_ownership(), shape, strides);
+        a += data[scalar_index];
+        EXPECT_EQ(data, expected);
+    }
+
+    TEST_CASE_TEMPLATE("scalar_assignment_noncopyable_functor", TypeParam, ADAPTOR_SEMANTIC_TYPES)
+    {
+        struct counted_add
+        {
+            counted_add() = default;
+            counted_add(const counted_add&) = delete;
+
+            int operator()(int lhs, int rhs)
+            {
+                return lhs + rhs + ++calls;
+            }
+
+            int calls = 0;
+        };
+
+        using shape_type = typename TypeParam::shape_type;
+        using strides_type = typename TypeParam::strides_type;
+        shape_type shape = {2, 1, 2};
+        strides_type strides = {3, 0, 1};
+        std::vector<int> data = {1, 2, -99, 3, 4, -99};
+        std::vector<int> expected;
+        std::size_t size = 0;
+
+        SUBCASE("padded")
+        {
+            size = 4;
+            expected = {4, 6, -99, 8, 10, -99};
+        }
+        SUBCASE("contiguous columns")
+        {
+            strides = {1, 0, 2};
+            size = 4;
+            data = {1, 2, 3, 4};
+            expected = {4, 6, 8, 10};
+        }
+        SUBCASE("contiguous axis permutation")
+        {
+            shape = {2, 3, 2};
+            strides = {1, 4, 2};
+            size = 12;
+            data = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
+            expected = {4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26};
+        }
+        SUBCASE("empty")
+        {
+            shape = {2, 0, 2};
+            size = 0;
+            expected = data;
+        }
+
+        auto a = adapt(data.data(), size, no_ownership(), shape, strides);
+        counted_add f;
+        scalar_computed_assign(a, 2, f);
+        EXPECT_EQ(data, expected);
+        EXPECT_EQ(f.calls, static_cast<int>(size));
+    }
+
     TEST_CASE_TEMPLATE("xsimd_info", TypeParam, ADAPTOR_SEMANTIC_TYPES)
     {
 #if defined(XTENSOR_USE_XSIMD)