Allow returning sets and lists from Fortran (pyccel#2102)

Add support to the Fortran to C wrapper to allow sets and lists to be
returned from Fortran. This is done by packing the elements into an
array. Fixes pyccel#1667 . Fixes pyccel#1666

Author: EmilyBourne

CHANGELOG.md

@@ -33,12 +33,12 @@ All notable changes to this project will be documented in this file.
 -   #1874 : Add C and Fortran support for the `len()` function for the `list` container.
 -   #1875 : Add C and Fortran support for the `len()` function for the `set` container.
 -   #1908 : Add C and Fortran support for the `len()` function for the `dict` container.
--   #1665 : Add C support for returning lists from functions.
+-   #1665 : Add C and Fortran support for returning lists from functions.
 -   #1689 : Add C and Fortran support for list method `append()`.
 -   #1876 : Add C support for indexing lists.
 -   #1690 : Add C support for list method `pop()`.
--   #1663 : Add C support for sets as constant arguments.
--   #1664 : Add C support for returning sets from functions.
+-   #1663 : Add C support for sets as arguments.
+-   #1664 : Add C and Fortran support for returning sets from functions.
 -   #2023 : Add support for iterating over a `set`.
 -   #1877 : Add C and Fortran Support for set method `pop()`.
 -   #1917 : Add C and Fortran support for set method `add()`.

pyccel/codegen/printing/ccode.py

@@ -1813,7 +1813,7 @@ def _print_Deallocate(self, expr):
         if isinstance(expr.variable.dtype, CustomDataType):
             Pyccel__del = expr.variable.cls_base.scope.find('__del__').name
             return f"{Pyccel__del}({variable_address});\n"
-        elif isinstance(expr.variable.class_type, (NumpyNDArrayType, HomogeneousContainerType)):
+        elif isinstance(expr.variable.class_type, (NumpyNDArrayType, HomogeneousTupleType)):
             if expr.variable.is_alias:
                 return f'free_pointer({variable_address});\n'
             else:

pyccel/codegen/wrapper/c_to_python_wrapper.py

@@ -2758,13 +2758,17 @@ def _extract_HomogeneousContainerType_FunctionDefResult(self, orig_var, is_bind_
         if isinstance(orig_var, PythonTuple):
             return self._extract_InhomogeneousTupleType_FunctionDefResult(orig_var, is_bind_c, funcdef)
 
-        if is_bind_c:
-            raise NotImplementedError("Support for returning sets from Fortran code is not yet available")
-
         name = getattr(orig_var, 'name', 'tmp')
         py_res = self.get_new_PyObject(f'{name}_obj', orig_var.dtype)
-        c_res = orig_var.clone(self.scope.get_new_name(name), is_argument = False)
-        loop_size = Variable(PythonNativeInt(), self.scope.get_new_name(f'{name}_size'))
+        if is_bind_c:
+            c_res = Variable(CStackArray(orig_var.class_type.element_type),
+                             self.scope.get_new_name(funcdef.results[0].var.name))
+            loop_size = funcdef.results[1].var.clone(self.scope.get_new_name(funcdef.results[1].var.name), is_argument = False)
+            c_results = [ObjectAddress(c_res), loop_size]
+        else:
+            c_res = orig_var.clone(self.scope.get_new_name(name), is_argument = False)
+            c_results = [c_res]
+            loop_size = Variable(PythonNativeInt(), self.scope.get_new_name(f'{name}_size'))
         idx = Variable(PythonNativeInt(), self.scope.get_new_name())
         self.scope.insert_variable(c_res)
         self.scope.insert_variable(loop_size)
@@ -2777,7 +2781,10 @@ def _extract_HomogeneousContainerType_FunctionDefResult(self, orig_var, is_bind_
 
         class_type = orig_var.class_type
         if isinstance(class_type, HomogeneousSetType):
-            element = SetPop(c_res)
+            if is_bind_c:
+                element = IndexedElement(c_res, idx)
+            else:
+                element = SetPop(c_res)
             elem_set = PySet_Add(py_res, element_extraction['py_result'])
             init = PySet_New()
         elif isinstance(class_type, HomogeneousListType):
@@ -2795,8 +2802,10 @@ def _extract_HomogeneousContainerType_FunctionDefResult(self, orig_var, is_bind_
                 *element_extraction['body'],
                 If(IfSection(PyccelEq(elem_set, PyccelUnarySub(LiteralInteger(1))),
                                          [Return([self._error_exit_code])]))]
-        body = [Assign(loop_size, PythonLen(c_res)),
-                AliasAssign(py_res, init),
-                For((idx,), PythonRange(loop_size), for_body, for_scope)]
+        body = [Assign(loop_size, PythonLen(c_res))] if not is_bind_c else []
+        body += [AliasAssign(py_res, init),
+                 For((idx,), PythonRange(loop_size), for_body, for_scope)]
+        if is_bind_c:
+            body.append(Deallocate(c_res))
 
-        return {'c_results': [c_res], 'py_result': py_res, 'body': body}
+        return {'c_results': c_results, 'py_result': py_res, 'body': body}

pyccel/codegen/wrapper/fortran_to_c_wrapper.py

@@ -13,16 +13,20 @@
 from pyccel.ast.bind_c import CLocFunc, BindCModule, BindCVariable
 from pyccel.ast.bind_c import BindCArrayVariable, BindCClassDef, DeallocatePointer
 from pyccel.ast.bind_c import BindCClassProperty
+from pyccel.ast.builtins import VariableIterator
 from pyccel.ast.core import Assign, FunctionCall, FunctionCallArgument
 from pyccel.ast.core import Allocate, EmptyNode, FunctionAddress
 from pyccel.ast.core import If, IfSection, Import, Interface, FunctionDefArgument
 from pyccel.ast.core import AsName, Module, AliasAssign, FunctionDefResult
+from pyccel.ast.core import For
 from pyccel.ast.datatypes import CustomDataType, FixedSizeNumericType
 from pyccel.ast.datatypes import HomogeneousTupleType, TupleType
+from pyccel.ast.datatypes import HomogeneousSetType, PythonNativeInt
+from pyccel.ast.datatypes import HomogeneousListType
 from pyccel.ast.internals import Slice
 from pyccel.ast.literals import LiteralInteger, Nil, LiteralTrue
 from pyccel.ast.numpytypes import NumpyNDArrayType
-from pyccel.ast.operators import PyccelIsNot, PyccelMul
+from pyccel.ast.operators import PyccelIsNot, PyccelMul, PyccelAdd
 from pyccel.ast.variable import Variable, IndexedElement, DottedVariable
 from pyccel.ast.numpyext import NumpyNDArrayType
 from pyccel.errors.errors import Errors
@@ -419,14 +423,33 @@ def _wrap_FunctionDefResult(self, expr):
 
             if not (var.is_alias or wrap_dotted):
                 # Create an array variable which can be passed to CLocFunc
-                ptr_var = var.clone(scope.get_new_name(name+'_ptr'),
+                ptr_var = Variable(NumpyNDArrayType(var.dtype, var.rank, var.order), scope.get_new_name(name+'_ptr'),
                                     memory_handling='alias')
                 scope.insert_variable(ptr_var)
 
                 # Define the additional steps necessary to define and fill ptr_var
                 alloc = Allocate(ptr_var, shape=result.shape, status='unallocated')
-                copy = Assign(ptr_var, local_var)
-                self._additional_exprs.extend([alloc, copy])
+                if isinstance(local_var.class_type, (NumpyNDArrayType, HomogeneousTupleType, CustomDataType)):
+                    copy = Assign(ptr_var, local_var)
+                    self._additional_exprs.extend([alloc, copy])
+                elif isinstance(local_var.class_type, (HomogeneousSetType, HomogeneousListType)):
+                    iterator = VariableIterator(local_var)
+                    elem = Variable(var.class_type.element_type, self.scope.get_new_name())
+                    idx = Variable(PythonNativeInt(), self.scope.get_new_name())
+                    self.scope.insert_variable(elem)
+                    assign = Assign(idx, LiteralInteger(0))
+                    for_scope = self.scope.create_new_loop_scope()
+                    for_body = [Assign(IndexedElement(ptr_var, idx), elem)]
+                    if isinstance(local_var.class_type, HomogeneousSetType):
+                        self.scope.insert_variable(idx)
+                        for_body.append(Assign(idx, PyccelAdd(idx, LiteralInteger(1))))
+                    else:
+                        iterator.set_loop_counter(idx)
+                    fill_for = For((elem,), iterator, for_body, scope = for_scope)
+                    self._additional_exprs.extend([alloc, assign, fill_for])
+                else:
+                    raise errors.report(f"Don't know how to return an object of type {local_var.class_type} to C code.",
+                            severity='fatal', symbol = var)
             else:
                 ptr_var = var
 

tests/epyccel/test_epyccel_lists.py

@@ -770,12 +770,12 @@ def list_ptr():
     assert isinstance(python_result, type(pyccel_result))
     assert python_result == pyccel_result
 
-def test_list_return(stc_language):
+def test_list_return(language):
     def list_return():
         a = [1,2,3,4,5]
         return a
 
-    epyccel_func = epyccel(list_return, language = stc_language)
+    epyccel_func = epyccel(list_return, language = language)
     pyccel_result = epyccel_func()
     python_result = list_return()
     assert python_result == pyccel_result

tests/epyccel/test_epyccel_sets.py

@@ -698,14 +698,12 @@ def set_arg(arg : 'set[int]', n : int):
     set_arg(arg_pyt, n)
     assert arg_pyc == arg_pyt
 
-def test_set_return(stc_language):
+def test_set_return(language):
     def set_return():
         a = {1,2,3,4,5}
-        b = {4,5,6}
-        c = a.union(b) # Use union to avoid #2084
-        return c
+        return a
 
-    epyccel_func = epyccel(set_return, language = stc_language)
+    epyccel_func = epyccel(set_return, language = language)
     pyccel_result = epyccel_func()
     python_result = set_return()
     assert python_result == pyccel_result