Extend the C min and max functions to more than two variables (pyccel#2041)

This PR aims to extend the C support for min and max functions to more
than two variables.
It leverages the `stc/common.h` functionalities to define two variadic
functions for min and max in
`stdlib/STC_EXTENSIONS/Common_extensions.h`.
This PR also fixed the import scheme for extension headers:
Previously, stc headers and extension headers were included separately
this is problematic because at the end of `stc` headers, we find a call
to `#include "priv/template2.h"`
which undefines what was defined by the header unless `imore` is
defined.
This prohibits the use of these in the extension header.

The proposed solution combines the `stc` header with its respective
extension and defines `imore`.

---------

Co-authored-by: Emily Bourne <[email protected]>

Author: jalalium

CHANGELOG.md

@@ -50,6 +50,7 @@ All notable changes to this project will be documented in this file.
 -   #738 : Add support for homogeneous tuples with scalar elements as arguments.
 -   Add a warning about containers in lists.
 -   #2016 : Add support for translating arithmetic magic methods (methods cannot yet be used from Python).
+-   #1980 : Extend The C support for min and max to more than two variables
 -   \[INTERNALS\] Add abstract class `SetMethod` to handle calls to various set methods.
 -   \[INTERNALS\] Added `container_rank` property to `ast.datatypes.PyccelType` objects.
 -   \[INTERNALS\] Add a `__call__` method to `FunctionDef` to create `FunctionCall` instances.

pyccel/ast/core.py

@@ -3711,6 +3711,30 @@ def define_target(self, new_target):
         else:
             self._target[new_target] = None
 
+    def remove_target(self, target_to_remove):
+        """
+        Remove a target from the imports.
+
+        Remove a target from the imports.
+        I.e., if `imp` is an Import defined as:
+        >>> from numpy import ones, cos
+
+        and we call `imp.remove_target('cos')`
+        then it becomes:
+        >>> from numpy import ones
+
+        Parameters
+        ----------
+        target_to_remove : str | AsName | iterable[str | AsName]
+            The import target(s) to remove.
+        """
+        assert pyccel_stage != "syntactic"
+        if iterable(target_to_remove):
+            for t in target_to_remove:
+                self._target.pop(t, None)
+        else:
+            self._target.pop(target_to_remove, None)
+
     def find_module_target(self, new_target):
         """
         Find the specified target amongst the targets of the Import.

pyccel/codegen/printing/ccode.py

@@ -14,8 +14,9 @@
 
 from pyccel.ast.bind_c    import BindCPointer
 
-from pyccel.ast.builtins  import PythonRange, PythonComplex
+from pyccel.ast.builtins  import PythonRange, PythonComplex, PythonMin
 from pyccel.ast.builtins  import PythonPrint, PythonType, VariableIterator
+
 from pyccel.ast.builtins  import PythonList, PythonTuple, PythonSet, PythonDict, PythonLen
 
 from pyccel.ast.builtin_methods.dict_methods  import DictItems
@@ -248,8 +249,14 @@
                  'stdbool',
                  'assert']}
 
-import_header_guard_prefix = {'Set_extensions'  : '_TOOLS_SET',
-                              'List_extensions' : '_TOOLS_LIST'}
+import_header_guard_prefix = {'Set_extensions'    : '_TOOLS_SET',
+                              'List_extensions'   : '_TOOLS_LIST',
+                              'Common_extensions' : '_TOOLS_COMMON'}
+
+
+stc_header_mapping = {'List_extensions': 'stc/vec',
+                      'Set_extensions': 'stc/hset',
+                      'Common_extensions': 'stc/common'}
 
 class CCodePrinter(CodePrinter):
     """
@@ -675,6 +682,32 @@ def init_stc_container(self, expr, assignment_var):
         init = f'{container_name} = c_init({dtype}, {keyraw});\n'
         return init
 
+    def invalidate_stc_headers(self, imports):
+        """
+        Invalidate STC headers when STC extension headers are present.
+
+        This function iterates over the list of imports and removes any targets
+        from STC headers if the target is present in their corresponding
+        STC extension headers.
+        The STC extension headers take care of including the standard
+        headers.
+
+        Parameters
+        ----------
+        imports : list of Import
+            The list of Import objects representing the header files to include.
+
+        Returns
+        -------
+        None
+            The function modifies the `imports` list in-place.
+        """
+        for imp in imports:
+            if imp.source in stc_header_mapping:
+                for imp2 in imports:
+                    if imp2.source == stc_header_mapping[imp.source]:
+                        imp2.remove_target(imp.target)
+
     def rename_imported_methods(self, expr):
         """
         Rename class methods from user-defined imports.
@@ -711,12 +744,13 @@ def _print_PythonAbs(self, expr):
             func = "labs"
         return "{}({})".format(func, self._print(expr.arg))
 
-    def _print_PythonMin(self, expr):
+    def _print_PythonMinMax(self, expr):
         arg = expr.args[0]
         if arg.dtype.primitive_type is PrimitiveFloatingPointType() and len(arg) == 2:
             self.add_import(c_imports['math'])
-            return "fmin({}, {})".format(self._print(arg[0]),
-                                         self._print(arg[1]))
+            arg1 = self._print(arg[0])
+            arg2 = self._print(arg[1])
+            return f"f{expr.name}({arg1}, {arg2})"
         elif arg.dtype.primitive_type is PrimitiveIntegerType() and len(arg) == 2:
             if isinstance(arg[0], Variable):
                 arg1 = self._print(arg[0])
@@ -734,38 +768,22 @@ def _print_PythonMin(self, expr):
                 self._additional_code += self._print(assign2)
                 arg2 = self._print(arg2_temp)
 
-            return f"({arg1} < {arg2} ? {arg1} : {arg2})"
+            op = '<' if isinstance(expr, PythonMin) else '>'
+            return f"({arg1} {op} {arg2} ? {arg1} : {arg2})"
+        elif len(arg) > 2 and isinstance(arg.dtype.primitive_type, (PrimitiveFloatingPointType, PrimitiveIntegerType)):
+            key = self.get_declare_type(arg[0])
+            self.add_import(Import('stc/common', AsName(VariableTypeAnnotation(arg.dtype), key)))
+            self.add_import(Import('Common_extensions', AsName(VariableTypeAnnotation(arg.dtype), key)))
+            return  f'{key}_{expr.name}({len(arg)}, {", ".join(self._print(a) for a in arg)})'
         else:
-            return errors.report("min in C is only supported for 2 scalar arguments", symbol=expr,
+            return errors.report(f"{expr.name} in C does not support arguments of type {arg.dtype}", symbol=expr,
                     severity='fatal')
 
-    def _print_PythonMax(self, expr):
-        arg = expr.args[0]
-        if arg.dtype.primitive_type is PrimitiveFloatingPointType() and len(arg) == 2:
-            self.add_import(c_imports['math'])
-            return "fmax({}, {})".format(self._print(arg[0]),
-                                         self._print(arg[1]))
-        elif arg.dtype.primitive_type is PrimitiveIntegerType() and len(arg) == 2:
-            if isinstance(arg[0], Variable):
-                arg1 = self._print(arg[0])
-            else:
-                arg1_temp = self.scope.get_temporary_variable(PythonNativeInt())
-                assign1 = Assign(arg1_temp, arg[0])
-                self._additional_code += self._print(assign1)
-                arg1 = self._print(arg1_temp)
-
-            if isinstance(arg[1], Variable):
-                arg2 = self._print(arg[1])
-            else:
-                arg2_temp = self.scope.get_temporary_variable(PythonNativeInt())
-                assign2 = Assign(arg2_temp, arg[1])
-                self._additional_code += self._print(assign2)
-                arg2 = self._print(arg2_temp)
+    def _print_PythonMin(self, expr):
+        return self._print_PythonMinMax(expr)
 
-            return f"({arg1} > {arg2} ? {arg1} : {arg2})"
-        else:
-            return errors.report("max in C is only supported for 2 scalar arguments", symbol=expr,
-                    severity='fatal')
+    def _print_PythonMax(self, expr):
+        return self._print_PythonMinMax(expr)
 
     def _print_SysExit(self, expr):
         code = ""
@@ -857,6 +875,7 @@ def _print_ModuleHeader(self, expr):
 
         # Print imports last to be sure that all additional_imports have been collected
         imports = [*expr.module.imports, *self._additional_imports.values()]
+        self.invalidate_stc_headers(imports)
         imports = ''.join(self._print(i) for i in imports)
 
         self._in_header = False
@@ -1033,7 +1052,20 @@ def _print_Import(self, expr):
             source = source.name[-1].python_value
         else:
             source = self._print(source)
-        if source.startswith('stc/') or source in import_header_guard_prefix:
+        if source == 'Common_extensions':
+            code = ''
+            for t in expr.target:
+                element_decl = f'#define i_key {t.local_alias}\n'
+                header_guard_prefix = import_header_guard_prefix.get(source, '')
+                header_guard = f'{header_guard_prefix}_{t.local_alias.upper()}'
+                code += ''.join((f'#ifndef {header_guard}\n',
+                     f'#define {header_guard}\n',
+                     element_decl,
+                     f'#include <{stc_header_mapping[source]}.h>\n', 
+                     f'#include <{source}.h>\n',
+                     f'#endif // {header_guard}\n\n'))
+            return code
+        elif source.startswith('stc/') or source in import_header_guard_prefix:
             code = ''
             for t in expr.target:
                 class_type = t.object.class_type
@@ -1055,6 +1087,8 @@ def _print_Import(self, expr):
                         f'#define {header_guard}\n',
                         f'#define i_type {container_type}\n',
                         element_decl,
+                        '#define i_more\n' if source in import_header_guard_prefix else '',
+                        f'#include <{stc_header_mapping[source]}.h>\n' if source in import_header_guard_prefix else '', 
                         f'#include <{source}.h>\n',
                         f'#endif // {header_guard}\n\n'))
             return code
@@ -2630,6 +2664,7 @@ def _print_Program(self, expr):
         decs = ''.join(self._print(Declare(v)) for v in variables)
 
         imports = [*expr.imports, *self._additional_imports.values()]
+        self.invalidate_stc_headers(imports)
         imports = ''.join(self._print(i) for i in imports)
 
         self.exit_scope()

pyccel/codegen/utilities.py

@@ -49,6 +49,7 @@
     "numpy_c"         : ("numpy", CompileObj("numpy_c.c",folder="numpy")),
     "Set_extensions"  : ("STC_Extensions", CompileObj("Set_Extensions.h", folder="STC_Extensions", has_target_file = False)),
     "List_extensions" : ("STC_Extensions", CompileObj("List_Extensions.h", folder="STC_Extensions", has_target_file = False)),
+    "Common_extensions" : ("STC_Extensions", CompileObj("Common_Extensions.h", folder="STC_Extensions", has_target_file = False)),
     "gFTL_functions/Set_extensions"  : ("gFTL_functions", CompileObj("Set_Extensions.inc", folder="gFTL_functions", has_target_file = False)),
 }
 internal_libs["cwrapper_ndarrays"] = ("cwrapper_ndarrays", CompileObj("cwrapper_ndarrays.c",folder="cwrapper_ndarrays",

pyccel/stdlib/STC_Extensions/Common_extensions.h

@@ -0,0 +1,31 @@
+#include <stdarg.h>
+
+static inline i_key c_JOIN(i_key,_min)(size_t count, ...) {
+    va_list args;
+    va_start(args, count);
+    i_key min_value = va_arg(args, i_key);
+    for (size_t i = 1; i < count; ++i) {
+        i_key value = va_arg(args, i_key);
+        if (value < min_value) {
+            min_value = value;
+        }
+    }
+    va_end(args);
+    return min_value;
+}
+
+static inline i_key c_JOIN(i_key, _max)(size_t count, ...) {
+    va_list args;
+    va_start(args, count);
+    i_key max_value = va_arg(args, i_key);
+    for (size_t i = 1; i < count; ++i) {
+        i_key value = va_arg(args, i_key);
+        if (value > max_value) {
+            max_value = value;
+        }
+    }
+    va_end(args);
+    return max_value;
+}
+
+#undef i_key

pyccel/stdlib/STC_Extensions/List_extensions.h

@@ -2,8 +2,6 @@
 #define LIST_EXTENSIONS_H
 
 
-#define _c_MEMB(name) c_JOIN(i_type, name)
-
 // This function represents a call to the .pop() method.
 // i_type: Class type (e.g., hset_int64_t).
 // i_key: Data type of the elements in the set (e.g., int64_t).
@@ -24,4 +22,5 @@ static inline i_key _c_MEMB(_pull_elem)(i_type* self, intptr_t pop_idx) {
 
 #undef i_type
 #undef i_key
+#include <stc/priv/template2.h>
 #endif

pyccel/stdlib/STC_Extensions/Set_extensions.h

@@ -2,8 +2,6 @@
 #define SET_EXTENSIONS_H
 #include <stdarg.h>
 
-#define _c_MEMB(name) c_JOIN(i_type, name)
-
 // This function represents a call to the .pop() method.
 // i_type: Class type (e.g., hset_int64_t).
 // i_key: Data type of the elements in the set (e.g., int64_t).
@@ -44,4 +42,5 @@ static inline i_type _c_MEMB(_union)(i_type* self, int n, ...) {
 
 #undef i_type
 #undef i_key
+#include <stc/priv/template2.h>
 #endif

tests/epyccel/test_builtins.py

@@ -109,15 +109,6 @@ def f(x : 'float', y : 'float'):
 
     assert np.isclose(epyc_f(*float_args), f(*float_args), rtol=RTOL, atol=ATOL)
 
-@pytest.mark.parametrize( 'language', (
-        pytest.param("fortran", marks = pytest.mark.fortran),
-        pytest.param("c", marks = [
-            pytest.mark.skip(reason="min not implemented in C for more than 2 args"),
-            pytest.mark.c]
-        ),
-        pytest.param("python", marks = pytest.mark.python)
-    )
-)
 def test_min_3_args(language):
     @template('T', [int, float])
     def f(x : 'T', y : 'T', z : 'T'):
@@ -228,19 +219,10 @@ def f(x : 'float', y : 'float'):
 
     assert np.isclose(epyc_f(*float_args), f(*float_args), rtol=RTOL, atol=ATOL)
 
-@pytest.mark.parametrize( 'language', (
-        pytest.param("fortran", marks = pytest.mark.fortran),
-        pytest.param("c", marks = [
-            pytest.mark.skip(reason="max not implemented in C for more than 2 args"),
-            pytest.mark.c]
-        ),
-        pytest.param("python", marks = pytest.mark.python)
-    )
-)
 def test_max_3_args(language):
     @template('T', [int, float])
     def f(x : 'T', y : 'T', z : 'T'):
-        return min(x, y, z)
+        return max(x, y, z)
 
     epyc_f = epyccel(f, language=language)