Allow iterating over a set (pyccel#2030)

Add support for iterating over a set. Fixes pyccel#2023

**Commit Summary**
- Add class type information for `zip`
- Complete definition of `IteratorType` following documentation
- Add printing support for iterating over a set
- Add iterator tests for sets

Author: EmilyBourne

CHANGELOG.md

@@ -35,6 +35,7 @@ All notable changes to this project will be documented in this file.
 -   #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()`.
+-   #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()`.
 -   #1918 : Add support for set method `clear()`.

pyccel/ast/builtins.py

@@ -1150,7 +1150,7 @@ class PythonZip(PyccelFunction):
     *args : tuple of TypedAstNode
         The arguments passed to the function.
     """
-    __slots__ = ('_length',)
+    __slots__ = ('_length', '_class_type')
     name = 'zip'
 
     def __init__(self, *args):
@@ -1168,6 +1168,7 @@ def __init__(self, *args):
                 self._length = min(lengths)
             else:
                 self._length = self.args[0].shape[0]
+            self._class_type = InhomogeneousTupleType(*[a.class_type for a in args])
 
     @property
     def length(self):

pyccel/ast/low_level_tools.py

@@ -42,6 +42,39 @@ def iterable_type(self):
         """
         return self._iterable_type
 
+    def __str__(self):
+        return f'Iter[{self._iterable_type}]'
+
+    @property
+    def datatype(self):
+        """
+        The datatype of the object.
+
+        The datatype of the object.
+        """
+        return self
+
+    @property
+    def rank(self):
+        """
+        Number of dimensions of the object.
+
+        Number of dimensions of the object. If the object is a scalar then
+        this is equal to 0.
+        """
+        return 0
+
+    @property
+    def order(self):
+        """
+        The data layout ordering in memory.
+
+        Indicates whether the data is stored in row-major ('C') or column-major
+        ('F') format. This is only relevant if rank > 1. When it is not relevant
+        this function returns None.
+        """
+        return None
+
 #------------------------------------------------------------------------------
 class PairType(PyccelType, metaclass=ArgumentSingleton):
     """

pyccel/codegen/printing/ccode.py

@@ -42,6 +42,8 @@
 from pyccel.ast.literals  import LiteralString, LiteralInteger, Literal
 from pyccel.ast.literals  import Nil
 
+from pyccel.ast.low_level_tools import IteratorType
+
 from pyccel.ast.mathext  import math_constants
 
 from pyccel.ast.numpyext import NumpyFull, NumpyArray
@@ -2326,45 +2328,44 @@ def _print_AliasAssign(self, expr):
     def _print_For(self, expr):
         self.set_scope(expr.scope)
 
-        indices = expr.iterable.loop_counters
-        index = indices[0] if indices else expr.target
-        if expr.iterable.num_loop_counters_required:
-            self.scope.insert_variable(index)
-
-        target   = index
-        iterable = expr.iterable.get_range()
-
-        if not isinstance(iterable, PythonRange):
-            # Only iterable currently supported is PythonRange
-            errors.report(PYCCEL_RESTRICTION_TODO, symbol=expr,
-                severity='fatal')
-
-        counter    = self._print(target)
-        body       = self._print(expr.body)
-
-        additional_assign = CodeBlock(expr.iterable.get_assigns(expr.target))
-        body = self._print(additional_assign) + body
-
-        start = self._print(iterable.start)
-        stop  = self._print(iterable.stop )
-        step  = self._print(iterable.step )
+        iterable = expr.iterable
+        iterable_type = iterable.iterable.class_type
+        indices = iterable.loop_counters
+
+        if isinstance(iterable_type, (DictType, HomogeneousSetType, HomogeneousListType)):
+            counter = Variable(IteratorType(iterable_type), indices[0].name)
+            c_type = self.get_c_type(iterable_type)
+            iterable_code = self._print(iterable.iterable)
+            for_code = f'c_foreach ({self._print(counter)}, {c_type}, {iterable_code})'
+            additional_assign = CodeBlock([Assign(expr.target, DottedVariable(VoidType(), 'ref',
+                memory_handling='alias', lhs = counter))])
+        else:
+            index = indices[0] if indices else expr.target
+            if iterable.num_loop_counters_required:
+                self.scope.insert_variable(index)
+
+            target   = index
+            counter  = self._print(target)
+            iterable = expr.iterable.get_range()
+            additional_assign = CodeBlock(expr.iterable.get_assigns(expr.target))
+
+            step = iterable.step
+            start_code = self._print(iterable.start)
+            stop_code  = self._print(iterable.stop )
+            step_code  = self._print(iterable.step )
+
+            # testing if the step is a value or an expression
+            if is_literal_integer(step):
+                op = '>' if int(step) < 0 else '<'
+                stop_condition = f'{counter} {op} {stop_code}'
+            else:
+                stop_condition = f'({step_code} > 0) ? ({counter} < {stop_code}) : ({counter} > {stop_code})'
+            for_code = f'for ({counter} = {start_code}; {stop_condition}; {counter} += {step_code})\n'
 
-        test_step = iterable.step
-        if isinstance(test_step, PyccelUnarySub):
-            test_step = iterable.step.args[0]
+        body = self._print(additional_assign) + self._print(expr.body)
 
         self.exit_scope()
-        # testing if the step is a value or an expression
-        if isinstance(test_step, Literal):
-            op = '>' if isinstance(iterable.step, PyccelUnarySub) else '<'
-            return ('for ({counter} = {start}; {counter} {op} {stop}; {counter} += '
-                        '{step})\n{{\n{body}}}\n').format(counter=counter, start=start, op=op,
-                                                          stop=stop, step=step, body=body)
-        else:
-            return (
-                'for ({counter} = {start}; ({step} > 0) ? ({counter} < {stop}) : ({counter} > {stop}); {counter} += '
-                '{step})\n{{\n{body}}}\n').format(counter=counter, start=start,
-                                                  stop=stop, step=step, body=body)
+        return for_code + '{\n' + body + '}\n'
 
     def _print_FunctionalFor(self, expr):
         loops = ''.join(self._print(i) for i in expr.loops)

pyccel/codegen/printing/fcode.py

@@ -1865,8 +1865,10 @@ def _print_Declare(self, expr):
         rankstr   = ''
 
         # ... print datatype
-        if isinstance(expr_type, CustomDataType):
-            name   = expr_type.name
+        if isinstance(expr_type, (CustomDataType, IteratorType, HomogeneousListType, HomogeneousSetType, DictType)):
+            name   = self._print(expr_type)
+            if isinstance(expr_type, (HomogeneousContainerType, DictType)):
+                self.add_import(self._build_gFTL_module(expr_type))
 
             if var.is_argument:
                 sig = 'class'
@@ -1894,10 +1896,6 @@ def _print_Declare(self, expr):
                     raise NotImplementedError("Fortran rank string undetermined")
                 rankstr = f'({rankstr})'
 
-        elif isinstance(expr_type, (HomogeneousListType, HomogeneousSetType, DictType)):
-            self.add_import(self._build_gFTL_module(expr_type))
-            typename = self._print(expr_type)
-            dtype_str = f'type({typename})'
         elif isinstance(dtype, StringType):
             dtype_str = self._print(dtype)
 
@@ -2231,6 +2229,9 @@ def _print_IteratorType(self, expr):
         iterable_type = self._print(expr.iterable_type)
         return f"{iterable_type}_Iterator"
 
+    def _print_CustomDataType(self, expr):
+        return expr.name
+
     def _print_DataType(self, expr):
         return self._print(expr.name)
 
@@ -2568,27 +2569,42 @@ def _print_FunctionalFor(self, expr):
     def _print_For(self, expr):
         self.set_scope(expr.scope)
 
+        iterable = expr.iterable
+        iterable_type = iterable.iterable.class_type
         indices = expr.iterable.loop_counters
-        index = indices[0] if indices else expr.target
-        if expr.iterable.num_loop_counters_required:
-            self.scope.insert_variable(index)
 
-        target   = index
-        my_range = expr.iterable.get_range()
-
-        if not isinstance(my_range, PythonRange):
-            # Only iterable currently supported is PythonRange
-            errors.report(PYCCEL_RESTRICTION_TODO, symbol=expr,
-                severity='fatal')
-
-        tar        = self._print(target)
-        range_code = self._print(my_range)
-
-        prolog = 'do {0} = {1}\n'.format(tar, range_code)
-        epilog = 'end do\n'
-
-        additional_assign = CodeBlock(expr.iterable.get_assigns(expr.target))
-        prolog += self._print(additional_assign)
+        if isinstance(iterable_type, (DictType, HomogeneousSetType)):
+            if isinstance(iterable.iterable, Variable):
+                suggested_name = iterable.iterable.name + '_'
+            else:
+                suggested_name = ''
+                errors.report("Iterating over a temporary object. This may cause compilation issues or cause calculations to be carried out twice",
+                        severity='warning', symbol=expr)
+            iterable = self._print(iterable.iterable)
+            iterator = self.scope.get_temporary_variable(IteratorType(iterable_type),
+                    name = suggested_name + 'iter')
+            last = self.scope.get_temporary_variable(IteratorType(iterable_type),
+                    name = suggested_name + 'last')
+            target = self._print(expr.target)
+            prolog = (f'{iterator} = {iterable} % begin()\n'
+                      f'{last} = {iterable} % end()\n'
+                      f'do while ({iterator} /= {last})\n'
+                      f'{target} = {iterator} % of()\n')
+            epilog = f'call {iterator} % next()\nend do\n'
+        else:
+            index = indices[0] if indices else expr.target
+            if iterable.num_loop_counters_required:
+                self.scope.insert_variable(index)
+
+            my_range = iterable.get_range()
+
+            target     = self._print(index)
+            range_code = self._print(my_range)
+
+            prolog = f'do {target} = {range_code}\n'
+            epilog = 'end do\n'
+
+            prolog += self._print(CodeBlock(iterable.get_assigns(expr.target)))
 
         body = self._print(expr.body)
 
@@ -2598,9 +2614,7 @@ def _print_For(self, expr):
 
         self.exit_scope()
 
-        return ('{prolog}'
-                '{body}'
-                '{epilog}').format(prolog=prolog, body=body, epilog=epilog)
+        return prolog + body + epilog
 
     # .....................................................
     #               Print OpenMP AnnotatedComment

tests/epyccel/test_epyccel_sets.py

@@ -532,3 +532,33 @@ def set_ptr():
     pyccel_result = epyccel_func()
     python_result = set_ptr()
     assert python_result == pyccel_result
+
+def test_set_iter(language):
+    def set_sum_int():
+        a =  {1,2,3,4,5,6,7,8,9,12}
+        sum_a = 0
+        for ai in a:
+            sum_a += ai
+        return sum_a
+
+    epyccel_func = epyccel(set_sum_int, language = language)
+    pyccel_result = epyccel_func()
+    python_result = set_sum_int()
+    assert python_result == pyccel_result
+    assert isinstance(python_result, type(pyccel_result))
+
+def test_set_iter_prod(language):
+    def set_iter_prod():
+        from itertools import product
+        a = {1,2,3,4,5,6,7,8,9,12}
+        b = {2.0, 4.0, 9.0, 2.5, 8.3}
+        assemble = 0.0
+        for ai, bi in product(a,b):
+            assemble += ai*bi
+        return assemble
+
+    epyccel_func = epyccel(set_iter_prod, language = language)
+    pyccel_result = epyccel_func()
+    python_result = set_iter_prod()
+    assert python_result == pyccel_result
+    assert isinstance(python_result, type(pyccel_result))