Fix array to slice assignment (pyccel#1810)

Fix assigning an array to a slice and add tests. Fixes pyccel#1218

This is achieved by fixing the indexing in the case where the variable
on the left-hand side of the assignment has a higher rank than the
variable on the right hand side.

---------

Co-authored-by: Yaman Güçlü <[email protected]>

Author: EmilyBourne

CHANGELOG.md

@@ -25,6 +25,7 @@ All notable changes to this project will be documented in this file.
 -   #1795 : Fix bug when returning slices in C.
 -   #1732 : Fix multidimensional list indexing in Python.
 -   #1785 : Add missing cast when creating an array of booleans from non-boolean values.
+-   #1218 : Fix bug when assigning an array to a slice in Fortran.
 
 ### Changed
 -   #1720 : functions with the `@inline` decorator are no longer exposed to Python in the shared library.

pyccel/ast/utilities.py

@@ -494,14 +494,14 @@ def collect_loops(block, indices, new_index, language_has_vectors = False, resul
             # Loop over indexes, inserting until the expression can be evaluated
             # in the desired language
             new_level = 0
-            for index in range(-rank,0):
+            for index_depth in range(-rank, 0):
                 new_level += 1
                 # If an index exists at the same depth, reuse it if not create one
-                if rank+index >= len(indices):
-                    indices.append(new_index(PythonNativeInt(),'i'))
-                index_var = indices[rank+index]
-                new_vars = [insert_index(v, index, index_var) for v in new_vars]
-                handled_funcs = [insert_index(v, index, index_var) for v in handled_funcs]
+                if rank+index_depth >= len(indices):
+                    indices.append(new_index(PythonNativeInt(), 'i'))
+                index = indices[rank+index_depth]
+                new_vars = [insert_index(v, index_depth, index) for v in new_vars]
+                handled_funcs = [insert_index(v, index_depth, index) for v in handled_funcs]
                 if compatible_operation(*new_vars, *handled_funcs, language_has_vectors = language_has_vectors):
                     break
 
@@ -541,22 +541,34 @@ def collect_loops(block, indices, new_index, language_has_vectors = False, resul
             current_level = new_level
 
         elif isinstance(line, Assign) and isinstance(line.lhs, IndexedElement) \
-                and isinstance(line.rhs, (PythonTuple, NumpyArray)) and not language_has_vectors:
-
+                and isinstance(line.rhs, (PythonTuple, NumpyArray)):
             lhs = line.lhs
             rhs = line.rhs
-            if isinstance(rhs, NumpyArray):
-                rhs = rhs.arg
-
-            lhs_rank = lhs.rank
-
-            new_assigns = [Assign(
-                            insert_index(expr=lhs,
-                                pos       = -lhs_rank,
-                                index_var = LiteralInteger(j)),
-                            rj) # lhs[j] = rhs[j]
-                          for j, rj in enumerate(rhs)]
-            collect_loops(new_assigns, indices, new_index, language_has_vectors, result = result)
+            if lhs.rank > rhs.rank:
+                for index_depth in range(lhs.rank-rhs.rank):
+                    # If an index exists at the same depth, reuse it if not create one
+                    if index_depth >= len(indices):
+                        indices.append(new_index(PythonNativeInt(), 'i'))
+                    index = indices[index_depth]
+                    lhs = insert_index(lhs, index_depth, index)
+                collect_loops([Assign(lhs, rhs)], indices, new_index, language_has_vectors, result = result)
+
+            elif not language_has_vectors:
+                if isinstance(rhs, NumpyArray):
+                    rhs = rhs.arg
+
+                lhs_rank = lhs.rank
+
+                new_assigns = [Assign(
+                                insert_index(expr=lhs,
+                                    pos       = -lhs_rank,
+                                    index_var = LiteralInteger(j)),
+                                rj) # lhs[j] = rhs[j]
+                              for j, rj in enumerate(rhs)]
+                collect_loops(new_assigns, indices, new_index, language_has_vectors, result = result)
+
+            else:
+                result.append(line)
 
         elif isinstance(line, Assign) and isinstance(line.rhs, Concatenate):
             lhs = line.lhs

tests/epyccel/modules/arrays.py

@@ -1581,6 +1581,29 @@ def array_2d_C_slice_stride_23(a : 'int[:,:]'):
     b = a[::d, ::c]
     return np.sum(b), b[0][0], b[-1][-1], len(b), len(b[0])
 
+#==============================================================================
+# Slice assignment
+#==============================================================================
+
+def copy_to_slice_issue_1218(n : int):
+    from numpy import zeros, array
+    x = 1
+    arr = zeros((2, n))
+    arr[0:x, 0:6:2] = array([2, 5, 6])
+    return arr
+
+def copy_to_slice_1(a : 'float[:]', b : 'float[:]'):
+    a[1:-1] = b
+
+def copy_to_slice_2(a : 'float[:,:]', b : 'float[:]'):
+    a[:, 1:-1] = b
+
+def copy_to_slice_3(a : 'float[:,:]', b : 'float[:]'):
+    a[:, 0] = b
+
+def copy_to_slice_4(a : 'float[:]', b : 'float[:]'):
+    a[::2] = b
+
 #==============================================================================
 # ARITHMETIC OPERATIONS
 #==============================================================================

tests/epyccel/test_arrays.py

@@ -3488,6 +3488,62 @@ def test_array_2d_C_slice_stride_23(language):
     f2 = epyccel(f1, language = language)
     assert f1(a) == f2(a)
 
+#==============================================================================
+# TEST : Slice assignment
+#==============================================================================
+def test_copy_to_slice_issue_1218(language):
+    pyth_f = arrays.copy_to_slice_issue_1218
+    epyc_f = epyccel(pyth_f, language = language)
+
+    n = 10
+    pyth_arr = pyth_f(n)
+    epyc_arr = epyc_f(n)
+    check_array_equal(pyth_arr, epyc_arr)
+
+def test_copy_to_slice_1(language):
+    pyth_f = arrays.copy_to_slice_1
+    epyc_f = epyccel(pyth_f, language = language)
+
+    pyth_a = np.arange(10, dtype=float)
+    epyc_a = pyth_a.copy()
+    b = np.arange(20, 28, dtype=float)
+    pyth_f(pyth_a, b)
+    epyc_f(epyc_a, b)
+    check_array_equal(pyth_a, epyc_a)
+
+def test_copy_to_slice_2(language):
+    pyth_f = arrays.copy_to_slice_2
+    epyc_f = epyccel(pyth_f, language = language)
+
+    pyth_a = np.arange(20, dtype=float).reshape(2, 10)
+    epyc_a = pyth_a.copy()
+    b = np.arange(20, 28, dtype=float)
+    pyth_f(pyth_a, b)
+    epyc_f(epyc_a, b)
+    check_array_equal(pyth_a, epyc_a)
+
+def test_copy_to_slice_3(language):
+    pyth_f = arrays.copy_to_slice_3
+    epyc_f = epyccel(pyth_f, language = language)
+
+    pyth_a = np.arange(20, dtype=float).reshape(4, 5)
+    epyc_a = pyth_a.copy()
+    b = np.arange(20, 24, dtype=float)
+    pyth_f(pyth_a, b)
+    epyc_f(epyc_a, b)
+    check_array_equal(pyth_a, epyc_a)
+
+def test_copy_to_slice_4(language):
+    pyth_f = arrays.copy_to_slice_4
+    epyc_f = epyccel(pyth_f, language = language)
+
+    pyth_a = np.arange(10, dtype=float)
+    epyc_a = pyth_a.copy()
+    b = np.arange(20, 25, dtype=float)
+    pyth_f(pyth_a, b)
+    epyc_f(epyc_a, b)
+    check_array_equal(pyth_a, epyc_a)
+
 #==============================================================================
 # TEST : arithmetic operations
 #==============================================================================