Merge pull request #281 from ecmwf-ifs/naml-fix-inline-nested-assoc

Fixing nested associate scope-parentage tracking after inlining

loki/expression/expr_visitors.py

@@ -346,6 +346,11 @@ def visit_Scope(self, o, **kwargs):
         # entry in the scope's table
         self._update_symbol_table_with_decls_and_imports(o)
 
+        # Attach parent scope if it is new before passing self down to children
+        parent_scope = kwargs.get('scope', o.parent)
+        if o.parent is not parent_scope and o is not parent_scope:
+            o._reset_parent(parent=parent_scope)
+
         # Then recurse to all children
         kwargs['scope'] = o
         children = tuple(self.visit(i, **kwargs) for i in o.children)

loki/scope.py

@@ -297,7 +297,7 @@ def rescope_symbols(self):
         to a scope in the scope hierarchy
         """
         from loki.expression import AttachScopes  # pylint: disable=import-outside-toplevel,cyclic-import
-        AttachScopes().visit(self)
+        AttachScopes().visit(self, scope=self)
 
     def make_complete(self, **frontend_args):
         """

loki/transform/transform_inline.py

@@ -498,6 +498,10 @@ def inline_subroutine_calls(routine, calls, callee, allowed_aliases=None):
     # Replace calls to child procedure with the child's body
     routine.body = Transformer(call_map).visit(routine.body)
 
+    # We need this to ensure that symbols, as well as nested scopes
+    # are correctly attached to each other (eg. nested associates).
+    routine.rescope_symbols()
+
 
 def inline_internal_procedures(routine, allowed_aliases=None):
     """

tests/test_transform_inline.py

@@ -19,7 +19,8 @@
 from loki.transform import (
     inline_elemental_functions, inline_constant_parameters,
     replace_selected_kind, inline_member_procedures,
-    inline_marked_subroutines, InlineTransformation
+    inline_marked_subroutines, InlineTransformation,
+    ResolveAssociatesTransformer
 )
 from loki.expression import symbols as sym
 
@@ -856,6 +857,72 @@ def test_inline_marked_routine_with_optionals(frontend, remove_imports):
     assert len(imports) == 0 if remove_imports else 1
 
 
+@pytest.mark.parametrize('frontend', available_frontends(
+    xfail=[(OMNI, 'OMNI has no sense of humour!')])
+)
+def test_inline_marked_subroutines_with_associates(frontend):
+    """ Test subroutine inlining via marker pragmas with nested associates. """
+
+    fcode_outer = """
+subroutine test_pragma_inline_associates(never)
+  use peter_pan, only: neverland
+  implicit none
+  type(neverland), intent(inout) :: never
+
+  associate(going=>never%going_to)
+
+  associate(up=>give_you%up)
+
+  !$loki inline
+  call dave(going, up)
+
+  end associate
+
+  end associate
+end subroutine test_pragma_inline_associates
+    """
+
+    fcode_inner = """
+subroutine dave(going)
+  use your_imagination, only: astley
+  implicit none
+  type(astley), intent(inout) :: going
+
+  associate(give_you=>going%give_you)
+
+  associate(up=>give_you%up)
+
+  call rick_is(up)
+
+  end associate
+
+  end associate
+end subroutine dave
+    """
+
+    outer = Subroutine.from_source(fcode_outer, frontend=frontend)
+    inner = Subroutine.from_source(fcode_inner, frontend=frontend)
+    outer.enrich(inner)
+
+    assert FindNodes(CallStatement).visit(outer.body)[0].routine == inner
+
+    inline_marked_subroutines(routine=outer, remove_imports=True)
+
+    # Ensure that all associates are perfectly nested afterwards
+    assocs = FindNodes(Associate).visit(outer.body)
+    assert len(assocs) == 4
+    assert assocs[1].parent == assocs[0]
+    assert assocs[2].parent == assocs[1]
+    assert assocs[3].parent == assocs[2]
+
+    # And, because we can...
+    outer.body = ResolveAssociatesTransformer().visit(outer.body)
+    call = FindNodes(CallStatement).visit(outer.body)[0]
+    assert call.name == 'rick_is'
+    assert call.arguments == ('never%going_to%give_you%up',)
+    # Q. E. D.
+
+
 @pytest.mark.parametrize('frontend', available_frontends(
     (OFP, 'Prefix/elemental support not implemented'))
 )