Repository: stfc/PSyclone

Issue:
sympy comparison fails if variable is called lambda
A comparison using a reference to a variable called lambda causes sympy.equal to fail.

Discussion:

Patch:
diff --git a/changelog b/changelog
index 9dfe20ad2..ab34883a5 100644
--- a/changelog
+++ b/changelog
@@ -575,6 +575,10 @@
 	193) PR #2260 for #2245. Splits LFRicLoopBounds class from
 	dynamo0.3.py.
 
+	194 PR #2241 for #2215. Adds support for Fortran names being the
+	same as Python keywords when using sympy within PSyclone, e.g. in
+	comparisons between expressions.
+
 release 2.3.1 17th of June 2022
 
 	1) PR #1747 for #1720. Adds support for If blocks to PSyAD.
diff --git a/src/psyclone/psyir/backend/sympy_writer.py b/src/psyclone/psyir/backend/sympy_writer.py
index d2f8e3664..faf4c6a35 100644
--- a/src/psyclone/psyir/backend/sympy_writer.py
+++ b/src/psyclone/psyir/backend/sympy_writer.py
@@ -38,6 +38,8 @@
 '''PSyIR backend to create expressions that are handled by SymPy.
 '''
 
+import keyword
+
 from sympy import Function, Symbol
 from sympy.parsing.sympy_parser import parse_expr
 
@@ -91,6 +93,11 @@ class SymPyWriter(FortranWriter):
     # is set to True as the modifications will persist after the Writer!
     _DISABLE_LOWERING = True
 
+    # A list of all reserved Python keywords (Fortran variables that are the
+    # same as a reserved name must be renamed, otherwise parsing will fail).
+    # This class attribute will get initialised in __init__:
+    _RESERVED_NAMES = set()
+
     def __init__(self):
         super().__init__()
 
@@ -107,6 +114,17 @@ class SymPyWriter(FortranWriter):
         self._lower_bound = "sympy_lower"
         self._upper_bound = "sympy_upper"
 
+        if not SymPyWriter._RESERVED_NAMES:
+            # Get the list of all reserved Python words from the
+            # keyword module:
+            for reserved in keyword.kwlist:
+                # Some Python keywords are capitalised (True, False, None).
+                # Since all symbols in PSyclone are lowercase, these can
+                # never clash when using SymPy. So only take keywords that
+                # are in lowercase:
+                if reserved.lower() == reserved:
+                    SymPyWriter._RESERVED_NAMES.add(reserved)
+
         # This dictionary will be supplied when parsing a string by SymPy
         # and defines which symbols in the parsed expressions are scalars
         # (SymPy symbols) or arrays (SymPy functions).
@@ -244,6 +262,16 @@ class SymPyWriter(FortranWriter):
         it is important to provide all expressions at once for the symbol
         table to avoid name clashes in any expression.
 
+        This function also handles reserved names like 'lambda' which might
+        occur in one of the Fortran expressions. These reserved names must be
+        renamed (otherwise SymPy parsing, which uses ``eval`` internally,
+        fails). A symbol table is used which is initially filled with all
+        reserved names. Then for each reference accounted, a unique name is
+        generated using this symbol table - so if the reference is a reserved
+        name, a new name will be created (e.g. ``lambda`` might become
+        ``lambda_1``). The SymPyWriter (e.g. in ``reference_node``) will
+        use the renamed value when creating the string representation.
+
         :param list_of_expressions: the list of expressions from which all
             references are taken and added to a symbol table to avoid
             renaming any symbols (so that only member names will be renamed).
@@ -252,8 +280,12 @@ class SymPyWriter(FortranWriter):
         '''
         # Create a new symbol table, so previous symbol will not affect this
         # new conversion (i.e. this avoids name clashes with a previous
-        # conversion).
+        # conversion). First add all reserved names so that these names will
+        # automatically be renamed. The symbol table is used later to also
+        # create guaranteed unique names for lower and upper bounds.
         self._symbol_table = SymbolTable()
+        for reserved in SymPyWriter._RESERVED_NAMES:
+            self._symbol_table.new_symbol(reserved)
 
         # Find each reference in each of the expression, and declare this name
         # as either a SymPy Symbol (scalar reference), or a SymPy Function
@@ -261,25 +293,28 @@ class SymPyWriter(FortranWriter):
         for expr in list_of_expressions:
             for ref in expr.walk(Reference):
                 name = ref.name
-                if name in self._symbol_table:
-                    # The name has already been declared, ignore it now
+                # The reserved Python keywords do not have tags, so they
+                # will not be found.
+                if name in self._symbol_table.tags_dict:
                     continue
 
-                # Add the new name to the symbol table to mark it
-                # as done
-                self._symbol_table.find_or_create(name)
-
+                # Any symbol from the list of expressions to be handled
+                # will be created with a tag, so if the same symbol is
+                # used more than once, the previous test will prevent
+                # calling new_symbol again. If the name is a Python
+                # reserved symbol, a new unique name will be created by
+                # the symbol table.
+                unique_sym = self._symbol_table.new_symbol(name, tag=name)
                 # Test if an array or an array expression is used:
                 if not ref.is_array:
-                    # A simple scalar, create a SymPy symbol
-                    self._sympy_type_map[name] = Symbol(name)
+                    self._sympy_type_map[unique_sym.name] = Symbol(name)
                     continue
 
                 # A Fortran array is used which has not been seen before.
                 # Declare a new SymPy function for it. This SymPy function
                 # will convert array expressions back into the original
                 # Fortran code.
-                self._sympy_type_map[name] = \
+                self._sympy_type_map[unique_sym.name] = \
                     self._create_sympy_array_function(name)
 
         # Now all symbols have been added to the symbol table, create
@@ -396,6 +431,21 @@ class SymPyWriter(FortranWriter):
         # list member
         return result[0]
 
+    # -------------------------------------------------------------------------
+    def arrayreference_node(self, node):
+        '''The implementation of the method handling a
+        ArrayOfStructureReference is generic enough to also handle
+        non-structure arrays. So just use it.
+
+        :param node: a ArrayReference PSyIR node.
+        :type node: :py:class:`psyclone.psyir.nodes.ArrayReference`
+
+        :returns: the code as string.
+        :rtype: str
+
+        '''
+        return self.arrayofstructuresreference_node(node)
+
     # -------------------------------------------------------------------------
     def structurereference_node(self, node):
         '''The implementation of the method handling a
@@ -555,10 +605,13 @@ class SymPyWriter(FortranWriter):
         :rtype: str
 
         '''
+        # Support renaming a symbol (e.g. if it is a reserved Python name).
+        # Look up with the name as tag, which will return the symbol with
+        # a unique name (e.g. lambda --> lambda_1):
+        name = self._symbol_table.lookup_with_tag(node.name).name
         if not node.is_array:
-            # This reference is not an array, handle its conversion to
-            # string in the FortranWriter base class
-            return super().reference_node(node)
+            # This reference is not an array, just return the name
+            return name
 
         # Now this must be an array expression without parenthesis. Add
         # the triple-array indices to represent `lower:upper:1` for each
@@ -567,7 +620,7 @@ class SymPyWriter(FortranWriter):
         result = [f"{self.lower_bound},"
                   f"{self.upper_bound},1"]*len(shape)
 
-        return (f"{node.name}{self.array_parenthesis[0]}"
+        return (f"{name}{self.array_parenthesis[0]}"
                 f"{','.join(result)}{self.array_parenthesis[1]}")
 
     # ------------------------------------------------------------------------
diff --git a/src/psyclone/psyir/tools/dependency_tools.py b/src/psyclone/psyir/tools/dependency_tools.py
index 53e6cf9d4..761bea476 100644
--- a/src/psyclone/psyir/tools/dependency_tools.py
+++ b/src/psyclone/psyir/tools/dependency_tools.py
@@ -383,12 +383,19 @@ class DependencyTools():
             return None
 
         symbol_map = sympy_writer.type_map
-        # If the subscripts do not even depend on the specified variable,
-        # any dependency distance is possible (e.g. `do i ... a(j)=a(j)+1`)
-        if var_name not in symbol_map:
+        # Find the SymPy symbol that has the same name as the var name. We
+        # cannot use the dictionary key, since a symbol might be renamed
+        # (e.g. if a Fortran variable 'lambda' is used, it will be
+        # renamed to lambda_1, and the type map will have lambda_1 as key
+        # and the SymPy symbol for 'lambda' as value).
+        for var in symbol_map.values():
+            if str(var) == var_name:
+                break
+        else:
+            # If the subscripts do not even depend on the specified variable,
+            # any dependency distance is possible (e.g. `do i ... a(j)=a(j)+1`)
             return None
 
-        var = symbol_map[var_name]
         # Create a unique 'd_x' variable name if 'x' is the variable.
         d_var_name = "d_"+var_name
         idx = 1
@@ -423,7 +430,7 @@ class DependencyTools():
                 # evaluated here. We then also need to check if `i+di` (i.e.
                 # the iteration to which the dependency is) is a valid
                 # iteration. E.g. in case of a(i^2)=a(i^2) --> di=0 or di=-2*i
-                # --> # i+di = -i < 0 for i>0. Since this is not a valid loop
+                # --> i+di = -i < 0 for i>0. Since this is not a valid loop
                 # iteration that means no dependencies.
                 return None
 


Test patch:
diff --git a/src/psyclone/tests/core/symbolic_maths_test.py b/src/psyclone/tests/core/symbolic_maths_test.py
index 38291e7aa..f5ea86a66 100644
--- a/src/psyclone/tests/core/symbolic_maths_test.py
+++ b/src/psyclone/tests/core/symbolic_maths_test.py
@@ -356,6 +356,8 @@ def test_symbolic_math_functions_with_constants(fortran_reader, expressions):
 
 
 @pytest.mark.parametrize("expressions", [("field(1+i)", "field(i+1)"),
+                                         ("lambda", "lambda"),
+                                         ("lambda(1+i)", "lambda(i+1)"),
                                          ("a%field(b+1)", "a%field(1+b)"),
                                          ("a%b%c(a_b+1)", "a%b%c(1+a_b)"),
                                          ("a%field(field+1)",
@@ -366,14 +368,16 @@ def test_symbolic_math_functions_with_constants(fortran_reader, expressions):
 def test_symbolic_math_use_reserved_names(fortran_reader, expressions):
     '''Test that reserved names are handled as expected. The SymPy parser
     uses 'eval' internally, so if a Fortran variable name should be the
-    same as a SymPy function (e.g. 'field'), parsing will fail.
+    same as a SymPy function (e.g. 'field'), parsing will fail. Similarly,
+    a Python reserved name (like 'lambda') would cause a parsing error.
 
     '''
     # A dummy program to easily create the PSyIR for the
     # expressions we need. We just take the RHS of the assignments
     source = f'''program test_prog
                  use some_mod
-                 integer :: field(10), i
+                 integer :: field(10)
+                 integer :: i, x
                  type(my_mod_type) :: a, b
                  x = {expressions[0]}
                  x = {expressions[1]}
@@ -422,6 +426,7 @@ def test_symbolic_math_use_range(fortran_reader, expressions):
 
 
 @pytest.mark.parametrize("expr,expected", [
+    ("lambda + 1", "lambda + 1"),
     ("1.0", "1.0"),
     ("a", "a"),
     ("a*b+c", "a * b + c"),
diff --git a/src/psyclone/tests/psyir/backend/sympy_writer_test.py b/src/psyclone/tests/psyir/backend/sympy_writer_test.py
index 7f8eec0b0..7c52737db 100644
--- a/src/psyclone/tests/psyir/backend/sympy_writer_test.py
+++ b/src/psyclone/tests/psyir/backend/sympy_writer_test.py
@@ -512,3 +512,45 @@ def test_sympy_writer_user_types(fortran_reader, fortran_writer,
     sympy_reader = SymPyReader(sympy_writer)
     new_psyir = sympy_reader.psyir_from_expression(sympy_exp, symbol_table)
     assert fortran_writer(new_psyir) == fortran_expr
+
+
+@pytest.mark.parametrize("expression", ["def", "if", "raise", "del",
+                                        "import", "return", "elif", "in",
+                                        "try", "and", "else", "is", "while",
+                                        "as", "except", "lambda", "with",
+                                        "assert", "finally", "nonlocal",
+                                        "yield", "break", "for", "not",
+                                        "class", "from", "or", "continue",
+                                        "global", "pass"])
+def test_sym_writer_reserved_names(fortran_reader, expression):
+    '''Test that reserved names are properly renamed. In this example,
+    all reserved named will get a ``_1`` appended.
+    '''
+    # A dummy program to easily create the PSyIR for the
+    # expressions we need. We just take the RHS of the assignments
+    source = f'''program test_prog
+                use some_mod
+                integer :: x
+                x = {expression}
+                end program test_prog '''
+    psyir = fortran_reader.psyir_from_source(source)
+    # psyir is a FileContainer, its first child the program, and its
+    # first child the assignment, of which we take the right hand side
+    psyir_expr = psyir.children[0].children[0].rhs
+
+    # Make sure that the symbols are renamed in the string representation
+    # of the PSyIR - the symbol table will make the symbols unique by
+    # appending ``_1``
+    sympy_writer = SymPyWriter()
+    assert sympy_writer._to_str(psyir_expr) == f"{expression}_1"
+
+    # The SymPy representation will contain e.g. the symbol 'lambda_1' after
+    # renaming, but when the expression is converted to a string, it should
+    # use 'lambda' as string representation for the symbol 'lambda_1'.
+    # Explicit:
+    #  >>> lambda_1 = sp.Symbol("lambda")
+    #  >>> print(lambda_1)
+    # lambda
+
+    sympy_exp = sympy_writer(psyir_expr)
+    assert str(sympy_exp) == expression
diff --git a/src/psyclone/tests/psyir/tools/dependency_tools_test.py b/src/psyclone/tests/psyir/tools/dependency_tools_test.py
index 943684d66..4860522dd 100644
--- a/src/psyclone/tests/psyir/tools/dependency_tools_test.py
+++ b/src/psyclone/tests/psyir/tools/dependency_tools_test.py
@@ -220,7 +220,7 @@ def test_arrays_parallelise(fortran_reader):
     loops = psyir.children[0].children[0:4]
     dep_tools = DependencyTools()
 
-    # Write to array that does not depend on parallel loop variable
+    # Write to array that does not depend on the parallel loop variable
     # Test that right default variable name (outer loop jj) is used.
     parallel = dep_tools.can_loop_be_parallelised(loops[0])
     assert parallel is False
@@ -230,12 +230,12 @@ def test_arrays_parallelise(fortran_reader):
     assert msg.code == DTCode.ERROR_WRITE_WRITE_RACE
     assert msg.var_names == ["mask(jk,jk)"]
 
-    # Write to array that does not depend on parallel loop variable
+    # Write to array that does not depend on the parallel loop variable
     parallel = dep_tools.can_loop_be_parallelised(loops[1])
     assert parallel is True
     assert dep_tools.get_all_messages() == []
 
-    # Use parallel loop variable in more than one dimension
+    # Use the parallel loop variable in more than one dimension
     parallel = dep_tools.can_loop_be_parallelised(loops[2])
     assert parallel is True
 
@@ -779,3 +779,72 @@ def test_da_array_expression(fortran_reader):
     dep_tools = DependencyTools()
     result = dep_tools.can_loop_be_parallelised(loops[0])
     assert result is False
+
+
+# -----------------------------------------------------------------------------
+def test_reserved_words(fortran_reader):
+    '''Tests that using a reserved Python word ('lambda' here') as a loop
+    variable, which will be renamed when converting to SymPy, works as
+    expected. Also make sure that name clashes are handled as expected by
+    declaring local symbols lambda_1 (which will clash with the renamed
+    lambda) and lambda_1_1 (which will clash with the renamed lambda_1).
+    Otherwise this test is identical to test_arrays_parallelise.
+    '''
+    source = '''program test
+                integer ji, lambda, jk
+                integer, parameter :: jpi=5, jpj=10
+                real, dimension(jpi,jpi) :: mask, umask, lambda_1, lambda_1_1
+                do lambda = 1, jpj   ! loop 0
+                   do ji = 1, jpi
+                      mask(jk, jk) = -1.0d0
+                    end do
+                end do
+                do lambda = 1, jpj   ! loop 1
+                   do ji = 1, jpi
+                      mask(ji, lambda+lambda_1+lambda_1_1) = umask(jk, jk)
+                    end do
+                end do
+                do lambda = 1, jpj   ! loop 2
+                   do ji = 1, jpi
+                      mask(lambda, lambda) = umask(ji, lambda)
+                    end do
+                end do
+                do lambda = 1, jpj   ! loop 3
+                   do ji = 1, jpi
+                      mask(ji, lambda) = mask(ji, lambda+1)
+                    end do
+                end do
+                end program test'''
+
+    psyir = fortran_reader.psyir_from_source(source)
+    loops = psyir.children[0].children[0:4]
+    dep_tools = DependencyTools()
+
+    # Write to array that does not depend on the parallel loop variable
+    # Test that right default variable name (outer loop lambda) is used.
+    parallel = dep_tools.can_loop_be_parallelised(loops[0])
+    assert parallel is False
+    msg = dep_tools.get_all_messages()[0]
+    assert ("The write access to 'mask(jk,jk)' causes a write-write race "
+            "condition" in str(msg))
+    assert msg.code == DTCode.ERROR_WRITE_WRITE_RACE
+    assert msg.var_names == ["mask(jk,jk)"]
+
+    # Write to array that does not depend on the parallel loop variable
+    parallel = dep_tools.can_loop_be_parallelised(loops[1])
+    assert parallel is True
+    assert dep_tools.get_all_messages() == []
+
+    # Use the parallel loop variable in more than one dimension
+    parallel = dep_tools.can_loop_be_parallelised(loops[2])
+    assert parallel is True
+
+    # Use a stencil access (with write), which prevents parallelisation
+    parallel = dep_tools.can_loop_be_parallelised(loops[3])
+    assert parallel is False
+    msg = dep_tools.get_all_messages()[0]
+    assert ("The write access to 'mask(ji,lambda)' and to "
+            "'mask(ji,lambda + 1)' are dependent and cannot be parallelised."
+            in str(msg))
+    assert msg.code == DTCode.ERROR_DEPENDENCY
+    assert msg.var_names == ["mask(ji,lambda)", "mask(ji,lambda + 1)"]