Repository: pyccel/pyccel (python)

Issue:
Assigning an array to a slice does't work in Fortran
**Describe the bug**
A wrong Fortran code is generated when trying to assign an array to a slice


**To Reproduce**
For example this python code:
```python
import numpy as np

def func(n:int):
    x = 1
    arr = np.zeros((2, n))
    arr[0:x,0:6:2] = np.array([2,5,6])
    return arr
```

Generate this Fortran code:
```bash/fortran/c
module error_0001


  use, intrinsic :: ISO_C_Binding, only : f64 => C_DOUBLE , i64 => &
        C_INT64_T
  implicit none

  contains

  !........................................
  subroutine func(n, arr) 

    implicit none

    real(f64), allocatable, intent(out) :: arr(:,:)
    integer(i64), value :: n
    integer(i64) :: x

    x = 1_i64
    allocate(arr(0:n - 1_i64, 0:1_i64))
    arr = 0.0_f64
    arr(0_i64:5_i64:2_i64, 0_i64:x - 1_i64) = [2_i64, 5_i64, 6_i64]
    return

  end subroutine func
  !........................................

end module error_0001
```

Which fail in the compilation stage with this error message:
```
22 |     arr(0_i64:5_i64:2_i64, 0_i64:x - 1_i64) = [2_i64, 5_i64, 6_i64]
      |    1
Error: Incompatible ranks 2 and 1 in assignment at (1)
```

**Language**
Fortran



Pull request description:
Fix array to slice assignment
Fix assigning an array to a slice and add tests. Fixes #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.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 28409842..9216ed8f 100644
--- a/CHANGELOG.md
+++ b/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.
diff --git a/pyccel/ast/utilities.py b/pyccel/ast/utilities.py
index 5d495d77..eb14caeb 100644
--- a/pyccel/ast/utilities.py
+++ b/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


Test patch:
diff --git a/tests/epyccel/modules/arrays.py b/tests/epyccel/modules/arrays.py
index 4fda55fb..9152252b 100644
--- a/tests/epyccel/modules/arrays.py
+++ b/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
 #==============================================================================
diff --git a/tests/epyccel/test_arrays.py b/tests/epyccel/test_arrays.py
index 5320bda4..6578511a 100644
--- a/tests/epyccel/test_arrays.py
+++ b/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
 #==============================================================================