Add Fortran support for set initialisation (pyccel#2004)

EmilyBourne · web-flow · commit 0f25459ecdc8 · 2024-10-05T20:35:43.000+02:00
Add Fortran support for set initialisation in a similar way to lists. Fixes pyccel#1658
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -28,6 +28,8 @@ All notable changes to this project will be documented in this file.
 -   #1659 : Add the appropriate C language equivalent for declaring a Python `set` container using the STC library.
 -   #1944 : Add the appropriate C language equivalent for declaring a Python `dict` container using the STC library.
 -   #1657 : Add the appropriate Fortran language equivalent for declaring a Python `list` container using the gFTL library.
+-   #1658 : Add the appropriate Fortran language equivalent for declaring a Python `set` container using the gFTL library.
+-   #1944 : Add the appropriate Fortran language equivalent for declaring a Python `dict` container using the gFTL library.
 -   #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.
diff --git a/docs/builtin-functions.md b/docs/builtin-functions.md
@@ -92,6 +92,27 @@ Python contains a limited number of builtin functions defined [here](https://doc
 | `reverse` | No |
 | `sort` | Python-only |
 
+## Set methods
+
+| Method | Supported |
+|----------|-----------|
+| `add` | Python-only |
+| `clear` | Python-only |
+| `copy` | Python-only |
+| `difference` | No |
+| `difference_update` | No |
+| `discard` | Python-only |
+| `intersection` | No |
+| `intersection_update` | No |
+| `isdisjoint` | No |
+| `issubset` | No |
+| `issuperset` | No |
+| `pop` | Python-only |
+| `remove` | Python-only |
+| `symmetric_difference` | No |
+| `symmetric_difference_update` | No |
+| `union` | No |
+| `update` | Python-only |
 
 ## Dictionary methods
 
diff --git a/pyccel/ast/builtins.py b/pyccel/ast/builtins.py
@@ -730,8 +730,10 @@ class PythonListFunction(PyccelFunction):
     __slots__ = ('_class_type', '_shape')
     _attribute_nodes = ()
 
-    def __new__(cls, arg):
-        if isinstance(arg, PythonList):
+    def __new__(cls, arg = None):
+        if arg is None:
+            return PythonList()
+        elif isinstance(arg, PythonList):
             return arg
         elif isinstance(arg.shape[0], LiteralInteger):
             return PythonList(*[arg[i] for i in range(arg.shape[0])])
@@ -774,6 +776,11 @@ def __init__(self, *args):
         super().__init__()
         if pyccel_stage == 'syntactic':
             return
+        elif len(args) == 0:
+            self._shape = (LiteralInteger(0),)
+            self._class_type = HomogeneousSetType(GenericType())
+            return
+
         arg0 = args[0]
         is_homogeneous = arg0.class_type is not GenericType() and \
                          all(a.class_type is not GenericType() and \
@@ -829,8 +836,10 @@ class PythonSetFunction(PyccelFunction):
 
     __slots__ = ('_shape', '_class_type')
     name = 'set'
-    def __new__(cls, arg):
-        if isinstance(arg.class_type, HomogeneousSetType):
+    def __new__(cls, arg = None):
+        if arg is None:
+            return PythonSet()
+        elif isinstance(arg.class_type, HomogeneousSetType):
             return arg
         elif isinstance(arg, (PythonList, PythonSet, PythonTuple)):
             return PythonSet(*arg)
diff --git a/pyccel/codegen/pipeline.py b/pyccel/codegen/pipeline.py
@@ -335,13 +335,19 @@ def handle_error(stage):
                               verbose  = verbose)
 
             mod_obj.add_dependencies(stdlib)
-
+            modules.append(stdlib)
 
     # Iterate over the external_libs list and determine if the printer
     # requires an external lib to be included.
     for key, import_node in codegen.get_printer_imports().items():
         try:
-            deps = generate_extension_modules(key, import_node, pyccel_dirpath, language)
+            deps = generate_extension_modules(key, import_node, pyccel_dirpath,
+                                              includes     = includes,
+                                              libs         = compile_libs,
+                                              libdirs      = libdirs,
+                                              dependencies = modules,
+                                              accelerators = accelerators,
+                                              language = language)
         except NotImplementedError as error:
             errors.report(f'{error}\n'+PYCCEL_RESTRICTION_TODO,
                 severity='error',
diff --git a/pyccel/codegen/printing/fcode.py b/pyccel/codegen/printing/fcode.py
@@ -24,7 +24,7 @@
 from pyccel.ast.builtins import PythonTuple, DtypePrecisionToCastFunction
 from pyccel.ast.builtins import PythonBool, PythonList, PythonSet
 
-from pyccel.ast.core import FunctionDef
+from pyccel.ast.core import FunctionDef, FunctionDefArgument, FunctionDefResult
 from pyccel.ast.core import SeparatorComment, Comment
 from pyccel.ast.core import ConstructorCall
 from pyccel.ast.core import FunctionCallArgument
@@ -37,7 +37,7 @@
 from pyccel.ast.datatypes import PrimitiveBooleanType, PrimitiveIntegerType, PrimitiveFloatingPointType, PrimitiveComplexType
 from pyccel.ast.datatypes import SymbolicType, StringType, FixedSizeNumericType, HomogeneousContainerType
 from pyccel.ast.datatypes import HomogeneousTupleType, HomogeneousListType, HomogeneousSetType, DictType
-from pyccel.ast.datatypes import PythonNativeInt
+from pyccel.ast.datatypes import PythonNativeInt, PythonNativeBool
 from pyccel.ast.datatypes import CustomDataType, InhomogeneousTupleType, TupleType
 from pyccel.ast.datatypes import pyccel_type_to_original_type, PyccelType
 
@@ -550,22 +550,46 @@ def _build_gFTL_module(self, expr_type):
                 include = Import(LiteralString('vector/template.inc'), Module('_', (), ()))
                 element_type = expr_type.element_type
                 if isinstance(element_type, FixedSizeNumericType):
-                    macros = [MacroDefinition('T', element_type.primitive_type),
+                    imports_and_macros = [MacroDefinition('T', element_type.primitive_type),
                               MacroDefinition('T_KINDLEN(context)', KindSpecification(element_type))]
                 else:
-                    macros = [MacroDefinition('T', element_type)]
+                    imports_and_macros = [MacroDefinition('T', element_type)]
                 if isinstance(element_type, (NumpyNDArrayType, HomogeneousTupleType)):
-                    macros.append(MacroDefinition('T_rank', element_type.rank))
+                    imports_and_macros.append(MacroDefinition('T_rank', element_type.rank))
                 elif not isinstance(element_type, FixedSizeNumericType):
                     raise NotImplementedError("Support for lists of types defined in other modules is not yet implemented")
-                macros.append(MacroDefinition('Vector', expr_type))
-                macros.append(MacroDefinition('VectorIterator', IteratorType(expr_type)))
+                imports_and_macros.append(MacroDefinition('Vector', expr_type))
+                imports_and_macros.append(MacroDefinition('VectorIterator', IteratorType(expr_type)))
+            elif isinstance(expr_type, HomogeneousSetType):
+                include = Import(LiteralString('set/template.inc'), Module('_', (), ()))
+                element_type = expr_type.element_type
+                imports_and_macros = []
+                if isinstance(element_type, FixedSizeNumericType):
+                    tmpVar_x = Variable(element_type, 'x')
+                    tmpVar_y = Variable(element_type, 'y')
+                    if isinstance(element_type.primitive_type, PrimitiveComplexType):
+                        complex_tool_import = Import('pyc_tools_f90', Module('pyc_tools_f90',(),()))
+                        self.add_import(complex_tool_import)
+                        imports_and_macros.append(complex_tool_import)
+                        compare_func = FunctionDef('complex_comparison',
+                                                   [FunctionDefArgument(tmpVar_x), FunctionDefArgument(tmpVar_y)],
+                                                   [FunctionDefResult(Variable(PythonNativeBool(), 'c'))], [])
+                        lt_def = FunctionCall(compare_func, [tmpVar_x, tmpVar_y])
+                    else:
+                        lt_def = PyccelAssociativeParenthesis(PyccelLt(tmpVar_x, tmpVar_y))
+                    imports_and_macros.extend([MacroDefinition('T', element_type.primitive_type),
+                              MacroDefinition('T_KINDLEN(context)', KindSpecification(element_type)),
+                              MacroDefinition('T_LT(x,y)', lt_def)])
+                else:
+                    raise NotImplementedError("Support for sets of types which define their own < operator is not yet implemented")
+                imports_and_macros.append(MacroDefinition('Set', expr_type))
+                imports_and_macros.append(MacroDefinition('SetIterator', IteratorType(expr_type)))
             else:
                 raise NotImplementedError(f"Unkown gFTL import for type {expr_type}")
 
             typename = self._print(expr_type)
             mod_name = f'{typename}_mod'
-            module = Module(mod_name, (), (), scope = Scope(), imports = [*macros, include],
+            module = Module(mod_name, (), (), scope = Scope(), imports = [*imports_and_macros, include],
                                        is_external = True)
 
             self._generated_gFTL_extensions[expr_type] = module
@@ -1091,6 +1115,21 @@ def _print_PythonList(self, expr):
             vec_type = self._print(expr.class_type)
         return f'{vec_type}({list_arg})'
 
+    def _print_PythonSet(self, expr):
+        if len(expr.args) == 0:
+            list_arg = ''
+            assign = expr.get_direct_user_nodes(lambda a : isinstance(a, Assign))
+            if assign:
+                set_type = self._print(assign[0].lhs.class_type)
+            else:
+                raise errors.report("Can't use an empty set without assigning it to a variable as the type cannot be deduced",
+                        severity='fatal', symbol=expr)
+
+        else:
+            list_arg = self._print_PythonTuple(expr)
+            set_type = self._print(expr.class_type)
+        return f'{set_type}({list_arg})'
+
     def _print_InhomogeneousTupleVariable(self, expr):
         fs = ', '.join(self._print(f) for f in expr)
         return '[{0}]'.format(fs)
@@ -1967,7 +2006,7 @@ def _print_Deallocate(self, expr):
             Pyccel_del_args = [FunctionCallArgument(var)]
             return self._print(FunctionCall(Pyccel__del, Pyccel_del_args))
 
-        if var.is_alias or isinstance(class_type, HomogeneousListType):
+        if var.is_alias or isinstance(class_type, (HomogeneousListType, HomogeneousSetType)):
             return ''
         elif isinstance(class_type, (NumpyNDArrayType, HomogeneousTupleType, StringType)):
             var_code = self._print(var)
@@ -2011,6 +2050,9 @@ def _print_PythonNativeBool(self, expr):
     def _print_HomogeneousListType(self, expr):
         return 'Vector_'+self._print(expr.element_type)
 
+    def _print_HomogeneousSetType(self, expr):
+        return 'Set_'+self._print(expr.element_type)
+
     def _print_IteratorType(self, expr):
         iterable_type = self._print(expr.iterable_type)
         return f"{iterable_type}_Iterator"
diff --git a/pyccel/codegen/utilities.py b/pyccel/codegen/utilities.py
@@ -40,6 +40,7 @@
     "ndarrays"        : ("ndarrays", CompileObj("ndarrays.c",folder="ndarrays")),
     "pyc_math_f90"    : ("math", CompileObj("pyc_math_f90.f90",folder="math")),
     "pyc_math_c"      : ("math", CompileObj("pyc_math_c.c",folder="math")),
+    "pyc_tools_f90"   : ("tools", CompileObj("pyc_tools_f90.f90",folder="tools")),
     "cwrapper"        : ("cwrapper", CompileObj("cwrapper.c",folder="cwrapper", accelerators=('python',))),
     "numpy_f90"       : ("numpy", CompileObj("numpy_f90.f90",folder="numpy")),
     "numpy_c"         : ("numpy", CompileObj("numpy_c.c",folder="numpy")),
@@ -180,7 +181,9 @@ def copy_internal_library(lib_folder, pyccel_dirpath, extra_files = None):
     return lib_dest_path
 
 #==============================================================================
-def generate_extension_modules(import_key, import_node, pyccel_dirpath, language):
+def generate_extension_modules(import_key, import_node, pyccel_dirpath,
+                               includes, libs, libdirs, dependencies,
+                               accelerators, language):
     """
     Generate any new modules that describe extensions.
 
@@ -196,6 +199,16 @@ def generate_extension_modules(import_key, import_node, pyccel_dirpath, language
         be printed).
     pyccel_dirpath : str
         The folder where files are being saved.
+    includes : iterable of strs
+        Include directories paths.
+    libs : iterable of strs
+        Required libraries.
+    libdirs : iterable of strs
+        Paths to directories containing the required libraries.
+    dependencies : iterable of CompileObjs
+        Objects which must also be compiled in order to compile this module/program.
+    accelerators : iterable of str
+        Tool used to accelerate the code (e.g. openmp openacc).
     language : str
         The language in which code is being printed.
 
@@ -221,7 +234,9 @@ def generate_extension_modules(import_key, import_node, pyccel_dirpath, language
                 f.write(code)
 
         new_dependencies.append(CompileObj(os.path.basename(filename), folder=folder,
-                            includes=(os.path.join(pyccel_dirpath, 'gFTL'),)))
+                            includes=(os.path.join(pyccel_dirpath, 'gFTL'), *includes),
+                            libs=libs, libdirs=libdirs, dependencies=dependencies,
+                            accelerators=accelerators))
 
     if lib_name in external_libs:
         copy_internal_library(lib_name, pyccel_dirpath)
diff --git a/pyccel/stdlib/tools/pyc_tools_f90.f90 b/pyccel/stdlib/tools/pyc_tools_f90.f90
@@ -0,0 +1,51 @@
+! -------------------------------------------------------------------------------------- !
+! This file is part of Pyccel which is released under MIT License. See the LICENSE file  !
+! or go to https://github.com/pyccel/pyccel/blob/devel/LICENSE for full license details. !
+! -------------------------------------------------------------------------------------- !
+
+module pyc_tools_f90
+
+  use, intrinsic :: ISO_C_Binding, only : b1 => C_BOOL, &
+         f32 => C_FLOAT, & 
+         f64 => C_DOUBLE, &
+         c64 => C_DOUBLE_COMPLEX, &
+         c32 => C_FLOAT_COMPLEX
+
+  implicit none
+
+  interface complex_comparison
+    module procedure complex_comparison_4
+    module procedure complex_comparison_8
+  end interface complex_comparison
+
+  contains
+
+  function complex_comparison_4(x, y) result(c)
+    complex(c32) :: x
+    complex(c32) :: y
+    logical(b1) :: c
+    real(f32) :: real_x
+    real(f32) :: real_y
+
+    real_x = real(x)
+    real_y = real(y)
+
+    c = merge(real_x < real_y, aimag(x) < aimag(y), real_x /= real_y)
+
+  end function complex_comparison_4
+
+  function complex_comparison_8(x, y) result(c)
+    complex(c64) :: x
+    complex(c64) :: y
+    logical(b1) :: c
+    real(f64) :: real_x
+    real(f64) :: real_y
+
+    real_x = real(x)
+    real_y = real(y)
+
+    c = merge(real_x < real_y, aimag(x) < aimag(y), real_x /= real_y)
+
+  end function complex_comparison_8
+
+end module pyc_tools_f90
diff --git a/tests/epyccel/test_epyccel_sets.py b/tests/epyccel/test_epyccel_sets.py
diff --git a/tests/epyccel/test_epyccel_variable_annotations.py b/tests/epyccel/test_epyccel_variable_annotations.py
