Merge pull request #2633 from stfc/2125_tuple_inline_fixes

(Closes #2125) Fix bugs in SymbolTable deep copy and InlineTrans when array bounds reference formal arg

changelog

@@ -169,9 +169,14 @@
 	59) PR #2649 towards #2586. Adds support for pointer assignments of
 	the form `var =[>] this[%...]%procedure(...)`.
 
-	60) PR #2645 for #2585. Add an optional "otherwise" argument to
+	60) PR #2645 towards #2585. Add an optional "otherwise" argument to
 	the SymbolTalbe lookup method.
 
+	61) PR #2633 for #2125. Introduces the "replace_symbols_using" method
+	to guarantee that PSyIR trees after a copy have references to the
+	updated symbols (including in attributes that are not directly in the
+	tree).
+
 release 2.5.0 14th of February 2024
 
 	1) PR #2199 for #2189. Fix bugs with missing maps in enter data

doc/developer_guide/psyir.rst

@@ -257,6 +257,24 @@ relationship.
 Methods like ``node.detach()``, ``node.copy()`` and ``node.pop_all_children()``
 can be used to move or replicate existing children into different nodes. 
 
+Tree Copying
+============
+
+The ability to create a deep-copy of a PSyIR tree is used heavily in PSyclone,
+primarily by the PSyIR backends. (This is because those backends often need
+to modify the tree while ensuring that the one provided by the caller remains
+unchanged.) As mentioned in the previous section, the ``node.copy()`` method
+provides this functionality:
+
+.. automethod:: psyclone.psyir.nodes.Node.copy
+
+As part of this copy operation, all Symbols referred to in the new tree must
+also be replaced with their equivalents from the symbol tables in the new tree.
+Since these symbol tables are associated with instances of ``ScopingNode``, it
+is ``ScopingNode._refine_copy`` which handles this:
+
+.. automethod:: psyclone.psyir.nodes.ScopingNode._refine_copy
+
 .. _update_signals_label:
 
 Dynamic Tree Updates

doc/developer_guide/psyir_symbols.rst

@@ -117,10 +117,6 @@ explicitly listed may be assumed to be unsupported):
 |                      |PRIVATE             |                    |
 +----------------------+--------------------+--------------------+
 
-.. warning:: Checking for equality between Type objects is currently
-	     only implemented for ``ScalarType``. This will be
-	     completed in #1799.
-
 It was decided to include datatype intrinsic as an attribute of ScalarType
 rather than subclassing. So, for example, a 4-byte real scalar is
 defined like this:
@@ -220,6 +216,42 @@ the `__contains__` method has no mechanism to pass a `scope_limit`
 optional argument. This would probably require a separate `setter` and
 `getter` to specify whether to check ancestors or not.
 
+Copying Symbols and Symbol Tables
+=================================
+
+Since Symbols can contain PSyIR nodes and other Symbols (e.g. as part
+of the definition of their precision or initial value), creating copies
+is not entirely straightforward. Every `Symbol` has the `copy` method:
+
+.. automethod:: psyclone.psyir.symbols.Symbol.copy
+
+This ensures that the precision and initial-value PSyIR sub-trees are
+copied appropriately while any Symbols referred to inside those nodes remain
+unchanged (and therefore can still be used in the same scope).
+
+However, when performing a *deep* copy of a PSyIR tree, all Symbols will
+need to be replaced with their equivalents in the new tree. The
+`SymbolTable.deep_copy()` method:
+
+.. automethod:: psyclone.psyir.symbols.SymbolTable.deep_copy
+
+handles this by first creating shallow copies of all Symbols in the
+table and then ensuring that each is updated to refer to Symbols in
+the new scope. This is achieved with the `replace_symbols_using`
+method:
+
+.. automethod:: psyclone.psyir.symbols.Symbol.replace_symbols_using
+
+All PSyIR `Node` classes also implement this method and call it when a
+`copy` operation is performed on the tree. The implementation in `Node`
+walks down the PSyIR tree and updates any Symbols using those in the supplied
+table. As the PSyIR supports nested scopes,  each `ScopingNode` is associated
+with a new symbol table. Therefore, the implementation within this class
+is slightly different:
+
+.. automethod:: psyclone.psyir.nodes.ScopingNode.replace_symbols_using
+
+
 Specialising Symbols
 ====================
 

src/psyclone/domain/lfric/lfric_constants.py

@@ -515,6 +515,8 @@ def precision_for_type(self, data_type):
         '''
         for module_info in self.DATA_TYPE_MAP.values():
             if module_info["type"] == data_type:
+                # TODO #2659 - this method should probably just return a name
+                #              rather than create a symbol itself.
                 # pylint: disable=import-outside-toplevel
                 from psyclone.domain.lfric.lfric_types import LFRicTypes
                 return LFRicTypes(module_info["kind"].upper())

src/psyclone/domain/lfric/lfric_extract_driver_creator.py

@@ -75,10 +75,7 @@ class LFRicExtractDriverCreator(BaseDriverCreator):
     1. The corresponding :py:class:`psyclone.psyGen.Invoke` statement that
        contains the kernel(s) is copied. This way we avoid affecting the tree
        of the caller. We need the invoke since it contains the symbol table.
-    2. We remove all halo exchange nodes. For now, the extract transformation
-       will not work when distributed memory is enabled, but since this
-       restriction is expected to be lifted, the code to handle this is
-       already added.
+    2. We remove all halo exchange nodes.
     3. We lower each kernel (child of the invoke) that was requested to
        be extracted, all others are removed. This is required since the kernel
        extraction will not contain the required data for the other kernels to
@@ -363,7 +360,8 @@ def _add_all_kernel_symbols(self, sched, symbol_table, proxy_name_mapping,
             new_symbol = symbol_table.new_symbol(root_name=reference.name,
                                                  tag=reference.name,
                                                  symbol_type=DataSymbol,
-                                                 datatype=datatype)
+                                                 datatype=datatype.copy())
+            new_symbol.replace_symbols_using(symbol_table)
             reference.symbol = new_symbol
 
         # Now handle all derived type. The name of a derived type is

src/psyclone/domain/lfric/lfric_types.py

@@ -160,11 +160,14 @@ def _create_precision_from_const_module():
         # Generate LFRic module symbols from definitions
         for module_info in modules:
             module_name = module_info.name.lower()
-            # Create the module (using a PSyIR ContainerSymbol)
+            # Create the module (using a PSyIR ContainerSymbol).
+            # TODO #2659 - this ContainerSymbol should be added to
+            # a SymbolTable!
             LFRicTypes._name_to_class[module_name] = \
                 ContainerSymbol(module_info.name)
             # Create the variables specified by the module (using
-            # PSyIR DataSymbols)
+            # PSyIR DataSymbols). TODO #2659 - these DataSymbols should
+            # be added to a SymbolTable!
             for module_var in module_info.vars:
                 var_name = module_var.upper()
                 interface = ImportInterface(LFRicTypes(module_name))

src/psyclone/psyir/backend/visitor.py

@@ -298,7 +298,9 @@ def _visit(self, node):
             # set so we ignore it and continue on down to any children.
             results = []
             for child in node.children:
-                results.append(self._visit(child))
+                result = self._visit(child)
+                if result is not None:
+                    results.append(result)
             return "".join(results)
 
         raise VisitorError(

src/psyclone/psyir/nodes/literal.py

@@ -39,8 +39,6 @@
 
 ''' This module contains the Literal node implementation.'''
 
-from __future__ import absolute_import
-
 import re
 
 from psyclone.psyir.nodes.datanode import DataNode
@@ -169,5 +167,22 @@ def node_str(self, colour=True):
         :returns: description of this PSyIR node.
         :rtype: str
         '''
-        return f"{self.coloured_name(colour)}"\
-               f"[value:'{self._value}', {self.datatype}]"
+        return (f"{self.coloured_name(colour)}"
+                f"[value:'{self._value}', {self.datatype}]")
+
+    def replace_symbols_using(self, table):
+        '''
+        Replace any Symbols referred to by this object with those in the
+        supplied SymbolTable with matching names. If there
+        is no match for a given Symbol then it is left unchanged.
+
+        :param table: the symbol table in which to look up replacement symbols.
+        :type table: :py:class:`psyclone.psyir.symbols.SymbolTable`
+
+        '''
+        self.datatype.replace_symbols_using(table)
+        super().replace_symbols_using(table)
+
+
+# For AutoAPI documentation generation
+__all__ = ['Literal']

src/psyclone/psyir/nodes/loop.py

@@ -125,9 +125,11 @@ def __eq__(self, other):
     @property
     def loop_type(self):
         '''
-        :returns: the type of this loop.
-        :rtype: str
+        :returns: the type of this loop, if set.
+        :rtype: Optional[str]
         '''
+        if not self._variable:
+            return None
         return self._loop_type_inference_rules.get(self.variable.name, None)
 
     @classmethod
@@ -412,11 +414,32 @@ def variable(self, var):
         self._check_variable(var)
         self._variable = var
 
+    def replace_symbols_using(self, table):
+        '''
+        Replace the Symbol referred to by this object's `variable` property
+        with that in the supplied SymbolTable with a matching name. If there
+        is no match then it is left unchanged.
+
+        :param table: symbol table in which to look up the replacement symbol.
+        :type table: :py:class:`psyclone.psyir.symbols.SymbolTable`
+
+        '''
+        if self._variable:
+            try:
+                new_sym = table.lookup(self._variable.name)
+                self.variable = new_sym
+            except KeyError:
+                pass
+        super().replace_symbols_using(table)
+
     def __str__(self):
         # Give Loop sub-classes a specialised name
         name = self.__class__.__name__
         result = name + "["
-        result += f"variable:'{self.variable.name}'"
+        if self._variable:
+            result += f"variable:'{self.variable.name}'"
+        else:
+            result += "variable:None"
         if self.loop_type:
             result += f", loop_type:'{self.loop_type}'"
         result += "]\n"

src/psyclone/psyir/nodes/node.py

@@ -1732,6 +1732,21 @@ def path_from(self, ancestor):
         result_list.reverse()
         return result_list
 
+    def replace_symbols_using(self, table):
+        '''
+        Replace any Symbols referred to by this object with those in the
+        supplied SymbolTable with matching names. If there
+        is no match for a given Symbol then it is left unchanged.
+
+        This base implementation simply propagates the call to any child Nodes.
+
+        :param table: the symbol table in which to look up replacement symbols.
+        :type table: :py:class:`psyclone.psyir.symbols.SymbolTable`
+
+        '''
+        for child in self.children:
+            child.replace_symbols_using(table)
+
 
 # For automatic documentation generation
 # TODO #913 the 'colored' routine shouldn't be in this module.

src/psyclone/psyir/nodes/reference.py

@@ -250,6 +250,24 @@ def next_access(self):
             return all_accesses[index+1].node
         return None
 
+    def replace_symbols_using(self, table):
+        '''
+        Update any Symbols referenced by this Node with those in the
+        supplied table with matching names. If there is no match for a given
+        Symbol then it is left unchanged.
+
+        :param table: the symbol table in which to look up replacement symbols.
+        :type table: :py:class:`psyclone.psyir.symbols.SymbolTable`
+
+        '''
+        try:
+            self.symbol = table.lookup(self.symbol.name)
+        except KeyError:
+            pass
+
+        # Walk on down the tree.
+        super().replace_symbols_using(table)
+
 
 # For AutoAPI documentation generation
 __all__ = ['Reference']

src/psyclone/psyir/nodes/scoping_node.py

@@ -37,8 +37,7 @@
 ''' This module contains the ScopingNode implementation.'''
 
 from psyclone.psyir.nodes.node import Node
-from psyclone.psyir.nodes.reference import Reference
-from psyclone.psyir.symbols import SymbolTable
+from psyclone.psyir.symbols import SymbolError, SymbolTable
 
 
 class ScopingNode(Node):
@@ -54,7 +53,7 @@ class ScopingNode(Node):
     :param parent: the parent node of this node in the PSyIR.
     :type parent: Optional[:py:class:`psyclone.psyir.nodes.Node`]
     :param symbol_table: attach the given symbol table to the new node.
-    :type symbol_table: \
+    :type symbol_table:
             Optional[:py:class:`psyclone.psyir.symbols.SymbolTable`]
 
     '''
@@ -90,35 +89,75 @@ def _refine_copy(self, other):
         ''' Refine the object attributes when a shallow copy is not the most
         appropriate operation during a call to the copy() method.
 
+        This method creates a deep copy of the SymbolTable associated with
+        the `other` scoping node and then calls `replace_symbols_using` to
+        update all Symbols referenced in the tree below this node.
+
+        .. warning::
+
+          Since `replace_symbols_using` only uses symbol *names*, this won't
+          get the correct symbol if the PSyIR has symbols shadowed in nested
+          scopes, e.g.:
+
+          .. code-block::
+
+              subroutine test
+                integer :: a
+                integer :: b = 1
+                if condition then
+                  ! PSyIR declares a shadowed, locally-scoped a'
+                  a' = 1
+                  if condition2 then
+                    ! PSyIR declares a shadowed, locally-scoped b'
+                    b' = 2
+                    a = a' + b'
+
+          Here, the final assignment will end up being `a' = a' + b'` and
+          thus the semantics of the code are changed. TODO #2666.
+
         :param other: object we are copying from.
         :type other: :py:class:`psyclone.psyir.node.Node`
 
         '''
         super(ScopingNode, self)._refine_copy(other)
         self._symbol_table = other.symbol_table.deep_copy()
-        # pylint: disable=protected-access
+        # pylint: disable-next=protected-access
         self._symbol_table._node = self  # Associate to self
 
-        # Update of children references to point to the equivalent symbols in
-        # the new symbol table attached to self.
-        # TODO #1377 Unfortunately Loop nodes currently store the associated
-        # loop variable in a `_variable` property rather than as a child so we
-        # must handle those separately. Also, in the LFRic API a Loop does not
-        # initially have the `_variable` property set which means that calling
-        # the `variable` getter causes an error (because it checks the
-        # internal-consistency of the Loop node). We therefore have to check
-        # the value of the 'private' `_variable` for now.
-        # We have to import Loop here to avoid a circular dependency.
-        # pylint: disable=import-outside-toplevel
-        from psyclone.psyir.nodes.loop import Loop
-        for node in self.walk((Reference, Loop)):
-            if isinstance(node, Reference):
-                if node.symbol in other.symbol_table.symbols:
-                    node.symbol = self.symbol_table.lookup(node.symbol.name)
-            if isinstance(node, Loop) and node._variable:
-                if node.variable in other.symbol_table.symbols:
-                    node.variable = self.symbol_table.lookup(
-                        node.variable.name)
+        # Now we've updated the symbol table, walk back down the tree and
+        # update any Symbols.
+        # Since this will get called from every ScopingNode in the tree, there
+        # could be a lot of repeated walks back down the tree. If this becomes
+        # a performance issue we could keep track of the depth of the recursive
+        # call to _refine_copy and only do this call when that depth is zero.
+        self.replace_symbols_using(self._symbol_table)
+
+    def replace_symbols_using(self, table):
+        '''
+        Update any Symbols referenced by this Node (and its descendants) with
+        those in the supplied table with matching names. If there is no match
+        for a given Symbol then it is left unchanged.
+
+        Since this is a ScopingNode, it is associated with a symbol table.
+        Therefore, if the supplied table is the one for the scope containing
+        this node (if any), the one passed to the child nodes is updated to be
+        the one associated with this node.
+
+        :param table: the symbol table in which to look up replacement symbols.
+        :type table: :py:class:`psyclone.psyir.symbols.SymbolTable`
+
+        '''
+        next_table = table
+        if self.parent:
+            try:
+                # If this node is not within a scope we get a SymbolError
+                if table is self.parent.scope.symbol_table:
+                    next_table = self.symbol_table
+            except SymbolError:
+                pass
+
+        for child in self.children:
+            child.replace_symbols_using(next_table)
 
     @property
     def symbol_table(self):