From 032eb4c77f736344b981f7e292cefc63d4cd71f7 Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Tue, 17 Sep 2024 07:58:17 +0000
Subject: [PATCH 01/11] Extract: Use `Transformer` in `outline_pragma_regions`

---
 loki/transformations/extract/outline.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index b6082b82a..576d9dcb3 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -9,7 +9,7 @@
 from loki.expression import Variable
 from loki.ir import (
     CallStatement, Import, PragmaRegion, Section, FindNodes,
-    FindVariables, MaskedTransformer, Transformer, is_loki_pragma,
+    FindVariables, Transformer, is_loki_pragma,
     get_pragma_parameters, pragma_regions_attached
 )
 from loki.logging import info
@@ -121,16 +121,12 @@ def outline_pragma_regions(routine):
                 # insert into list of new routines
                 routines.append(region_routine)
 
-                # Register start and end nodes in transformer mask for original routine
-                starts += [region.pragma_post]
-                stops += [region.pragma]
-
                 # Replace end pragma by call in original routine
                 call_arguments = region_in_args + region_inout_args + region_out_args
                 call = CallStatement(name=Variable(name=name), arguments=call_arguments)
-                mask_map[region.pragma_post] = call
+                mask_map[region] = call
 
-    routine.body = MaskedTransformer(active=True, start=starts, stop=stops, mapper=mask_map).visit(routine.body)
+            routine.body = Transformer(mapper=mask_map).visit(routine.body)
     info('%s: converted %d region(s) to calls', routine.name, counter)
 
     return routines

From 53de1f8ca8f3f6b886737d1d8e05493e3f3bce6d Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Tue, 17 Sep 2024 08:47:51 +0000
Subject: [PATCH 02/11] Extract: Add `outline_region` for single region
 extraction

This now allows us to call outlining extraction on a single pragma
region, while the overarching utility uses this in its own pragma
loop.
---
 loki/transformations/extract/outline.py | 160 ++++++++++++++----------
 1 file changed, 96 insertions(+), 64 deletions(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 576d9dcb3..76bf8373d 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -17,7 +17,95 @@
 from loki.tools import as_tuple, CaseInsensitiveDict
 
 
-__all__ = ['outline_pragma_regions']
+
+__all__ = ['outline_region', 'outline_pragma_regions']
+
+
+def outline_region(region, name, imports, intent_map=None):
+    """
+    Creates a new :any:`Subroutine` object from a given :any:`PragmaRegion`.
+
+    Parameters
+    ----------
+    region : :any:`PragmaRegion`
+        The region that holds the body for which to create a subroutine.
+    name : str
+        Name of the new subroutine
+    imports : tuple of :any:`Import`, optional
+        List of imports to replicate in the new subroutine
+    intent_map : dict, optional
+        Mapping of instent strings to list of variables to override intents
+
+    Returns
+    -------
+    tuple of :any:`CallStatement` and :any:`Subroutine`
+        The newly created call and respectice subroutine.
+    """
+    intent_map = intent_map or {}
+    imports = as_tuple(imports)
+    imported_symbols = {var for imp in imports for var in imp.symbols}
+
+    # Create the external subroutine containing the routine's imports and the region's body
+    spec = Section(body=imports)
+    body = Section(body=Transformer().visit(region.body))
+    region_routine = Subroutine(name, spec=spec, body=body)
+
+    # Use dataflow analysis to find in, out and inout variables to that region
+    # (ignoring any symbols that are external imports)
+    region_in_args = region.uses_symbols - region.defines_symbols - imported_symbols
+    region_inout_args = region.uses_symbols & region.defines_symbols - imported_symbols
+    region_out_args = region.defines_symbols - region.uses_symbols - imported_symbols
+
+    # Remove any parameters from in args
+    region_in_args = {arg for arg in region_in_args if not arg.type.parameter}
+
+    # Extract arguments given in pragma annotations
+    region_var_map = CaseInsensitiveDict(
+        (v.name, v.clone(dimensions=None))
+        for v in FindVariables().visit(region.body)
+        if v.clone(dimensions=None) not in imported_symbols
+    )
+    pragma_in_args = {region_var_map[v.lower()] for v in intent_map.get('in', '').split(',') if v}
+    pragma_inout_args = {region_var_map[v.lower()] for v in intent_map.get('inout', '').split(',') if v}
+    pragma_out_args = {region_var_map[v.lower()] for v in intent_map.get('out', '').split(',') if v}
+
+    # Override arguments according to pragma annotations
+    region_in_args = (region_in_args - (pragma_inout_args | pragma_out_args)) | pragma_in_args
+    region_inout_args = (region_inout_args - (pragma_in_args | pragma_out_args)) | pragma_inout_args
+    region_out_args = (region_out_args - (pragma_in_args | pragma_inout_args)) | pragma_out_args
+
+    # Now fix the order
+    region_inout_args = as_tuple(region_inout_args)
+    region_in_args = as_tuple(region_in_args)
+    region_out_args = as_tuple(region_out_args)
+
+    # Set the list of variables used in region routine (to create declarations)
+    # and put all in the new scope
+    region_routine_variables = {v.clone(dimensions=v.type.shape or None)
+                                for v in FindVariables().visit(region_routine.body)
+                                if v.name in region_var_map}
+    region_routine.variables = as_tuple(region_routine_variables)
+    region_routine.rescope_symbols()
+
+    # Build the call signature
+    region_routine_var_map = region_routine.variable_map
+    region_routine_arguments = []
+    for intent, args in zip(('in', 'inout', 'out'), (region_in_args, region_inout_args, region_out_args)):
+        for arg in args:
+            local_var = region_routine_var_map[arg.name]
+            local_var = local_var.clone(type=local_var.type.clone(intent=intent))
+            region_routine_var_map[arg.name] = local_var
+            region_routine_arguments += [local_var]
+
+    # We need to update the list of variables again to avoid duplicate declarations
+    region_routine.variables = as_tuple(region_routine_var_map.values())
+    region_routine.arguments = as_tuple(region_routine_arguments)
+
+    # Create the call according to the wrapped code region
+    call_arguments = region_in_args + region_inout_args + region_out_args
+    call = CallStatement(name=Variable(name=name), arguments=call_arguments)
+
+    return call, region_routine
 
 
 def outline_pragma_regions(routine):
@@ -46,12 +134,11 @@ def outline_pragma_regions(routine):
     -------
     list of :any:`Subroutine`
         the list of newly created subroutines.
-
     """
     counter = 0
-    routines, starts, stops = [], [], []
-    imports = {var for imprt in FindNodes(Import).visit(routine.spec) for var in imprt.symbols}
-    mask_map = {}
+    routines = []
+    imports = FindNodes(Import).visit(routine.spec)
+    mapper = {}
     with pragma_regions_attached(routine):
         with dataflow_analysis_attached(routine):
             for region in FindNodes(PragmaRegion).visit(routine.body):
@@ -63,70 +150,15 @@ def outline_pragma_regions(routine):
                 name = parameters.get('name', f'{routine.name}_outlined_{counter}')
                 counter += 1
 
-                # Create the external subroutine containing the routine's imports and the region's body
-                spec = Section(body=Transformer().visit(FindNodes(Import).visit(routine.spec)))
-                body = Section(body=Transformer().visit(region.body))
-                region_routine = Subroutine(name, spec=spec, body=body)
-
-                # Use dataflow analysis to find in, out and inout variables to that region
-                # (ignoring any symbols that are external imports)
-                region_in_args = region.uses_symbols - region.defines_symbols - imports
-                region_inout_args = region.uses_symbols & region.defines_symbols - imports
-                region_out_args = region.defines_symbols - region.uses_symbols - imports
-
-                # Remove any parameters from in args
-                region_in_args = {arg for arg in region_in_args if not arg.type.parameter}
-
-                # Extract arguments given in pragma annotations
-                region_var_map = CaseInsensitiveDict(
-                    (v.name, v.clone(dimensions=None))
-                    for v in FindVariables().visit(region.body) if v.clone(dimensions=None) not in imports
-                )
-                pragma_in_args = {region_var_map[v.lower()] for v in parameters.get('in', '').split(',') if v}
-                pragma_inout_args = {region_var_map[v.lower()] for v in parameters.get('inout', '').split(',') if v}
-                pragma_out_args = {region_var_map[v.lower()] for v in parameters.get('out', '').split(',') if v}
-
-                # Override arguments according to pragma annotations
-                region_in_args = (region_in_args - (pragma_inout_args | pragma_out_args)) | pragma_in_args
-                region_inout_args = (region_inout_args - (pragma_in_args | pragma_out_args)) | pragma_inout_args
-                region_out_args = (region_out_args - (pragma_in_args | pragma_inout_args)) | pragma_out_args
-
-                # Now fix the order
-                region_inout_args = as_tuple(region_inout_args)
-                region_in_args = as_tuple(region_in_args)
-                region_out_args = as_tuple(region_out_args)
-
-                # Set the list of variables used in region routine (to create declarations)
-                # and put all in the new scope
-                region_routine_variables = {v.clone(dimensions=v.type.shape or None)
-                                            for v in FindVariables().visit(region_routine.body)
-                                            if v.name in region_var_map}
-                region_routine.variables = as_tuple(region_routine_variables)
-                region_routine.rescope_symbols()
-
-                # Build the call signature
-                region_routine_var_map = region_routine.variable_map
-                region_routine_arguments = []
-                for intent, args in zip(('in', 'inout', 'out'), (region_in_args, region_inout_args, region_out_args)):
-                    for arg in args:
-                        local_var = region_routine_var_map[arg.name]
-                        local_var = local_var.clone(type=local_var.type.clone(intent=intent))
-                        region_routine_var_map[arg.name] = local_var
-                        region_routine_arguments += [local_var]
-
-                # We need to update the list of variables again to avoid duplicate declarations
-                region_routine.variables = as_tuple(region_routine_var_map.values())
-                region_routine.arguments = as_tuple(region_routine_arguments)
+                call, region_routine = outline_region(region, name, imports, intent_map=parameters)
 
                 # insert into list of new routines
                 routines.append(region_routine)
 
-                # Replace end pragma by call in original routine
-                call_arguments = region_in_args + region_inout_args + region_out_args
-                call = CallStatement(name=Variable(name=name), arguments=call_arguments)
-                mask_map[region] = call
+                # Replace region by call in original routine
+                mapper[region] = call
 
-            routine.body = Transformer(mapper=mask_map).visit(routine.body)
+            routine.body = Transformer(mapper=mapper).visit(routine.body)
     info('%s: converted %d region(s) to calls', routine.name, counter)
 
     return routines

From 99a17ea0dfeb2988b90ddef05d300e1384944f9b Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Wed, 25 Sep 2024 15:06:23 +0000
Subject: [PATCH 03/11] Extract: Bug fix for `outline_pragma_regions`

Need to explicit give an empty tuple for kwargument when creating
the CallStatement, as dataflow analysis will iterate it.
---
 loki/transformations/extract/outline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 76bf8373d..51e4c0221 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -103,7 +103,7 @@ def outline_region(region, name, imports, intent_map=None):
 
     # Create the call according to the wrapped code region
     call_arguments = region_in_args + region_inout_args + region_out_args
-    call = CallStatement(name=Variable(name=name), arguments=call_arguments)
+    call = CallStatement(name=Variable(name=name), arguments=call_arguments, kwarguments=())
 
     return call, region_routine
 

From 58a8e7ed0ea52b1d131b43b5e67a1cf93d57736d Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Thu, 26 Sep 2024 08:04:26 +0000
Subject: [PATCH 04/11] Extract: Add C-import special-case to routine
 extraction

---
 loki/transformations/extract/outline.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 51e4c0221..d62fd8031 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -44,6 +44,10 @@ def outline_region(region, name, imports, intent_map=None):
     intent_map = intent_map or {}
     imports = as_tuple(imports)
     imported_symbols = {var for imp in imports for var in imp.symbols}
+    # Special-case for IFS-style C-imports
+    imported_symbols |= {
+        str(imp.module).split('.', maxsplit=1)[0] for imp in imports if imp.c_import
+    }
 
     # Create the external subroutine containing the routine's imports and the region's body
     spec = Section(body=imports)

From 9082e63632a920ddf6f328d434bd9eca03dcbbad Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Sat, 19 Oct 2024 09:12:57 +0000
Subject: [PATCH 05/11] Extract: Allow derived-type accesses in routine
 extraction

---
 loki/transformations/extract/outline.py       | 27 +++++--
 .../extract/tests/test_outline.py             | 79 +++++++++++++++++++
 2 files changed, 99 insertions(+), 7 deletions(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index d62fd8031..33ad9e5c6 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -54,11 +54,15 @@ def outline_region(region, name, imports, intent_map=None):
     body = Section(body=Transformer().visit(region.body))
     region_routine = Subroutine(name, spec=spec, body=body)
 
+    # Filter derived-type component accesses and only use the root parent
+    region_uses_symbols = {s.parents[0] if s.parent else s for s in region.uses_symbols}
+    region_defines_symbols = {s.parents[0] if s.parent else s for s in region.defines_symbols}
+
     # Use dataflow analysis to find in, out and inout variables to that region
     # (ignoring any symbols that are external imports)
-    region_in_args = region.uses_symbols - region.defines_symbols - imported_symbols
-    region_inout_args = region.uses_symbols & region.defines_symbols - imported_symbols
-    region_out_args = region.defines_symbols - region.uses_symbols - imported_symbols
+    region_in_args = region_uses_symbols - region_defines_symbols - imported_symbols
+    region_inout_args = region_uses_symbols & region_defines_symbols - imported_symbols
+    region_out_args = region_defines_symbols - region_uses_symbols - imported_symbols
 
     # Remove any parameters from in args
     region_in_args = {arg for arg in region_in_args if not arg.type.parameter}
@@ -85,10 +89,19 @@ def outline_region(region, name, imports, intent_map=None):
 
     # Set the list of variables used in region routine (to create declarations)
     # and put all in the new scope
-    region_routine_variables = {v.clone(dimensions=v.type.shape or None)
-                                for v in FindVariables().visit(region_routine.body)
-                                if v.name in region_var_map}
-    region_routine.variables = as_tuple(region_routine_variables)
+    region_routine_variables = {
+        v.clone(dimensions=v.type.shape or None)
+        for v in FindVariables().visit(region_routine.body)
+        if v.name in region_var_map
+    }
+    # Filter out derived-type component variables from declarations
+    region_routine_variables = {
+        v.parents[0] if v.parent else v for v in region_routine_variables
+    }
+    # Order the local devlaration list to put arguments first
+    region_routine_args = tuple(v for v in region_routine_variables if v.type.intent)
+    region_routine_locals = tuple(v for v in region_routine_variables if not v.type.intent)
+    region_routine.variables = region_routine_args + region_routine_locals
     region_routine.rescope_symbols()
 
     # Build the call signature
diff --git a/loki/transformations/extract/tests/test_outline.py b/loki/transformations/extract/tests/test_outline.py
index a17ae7361..a9ad85419 100644
--- a/loki/transformations/extract/tests/test_outline.py
+++ b/loki/transformations/extract/tests/test_outline.py
@@ -361,3 +361,82 @@ def test_outline_pragma_regions_imports(tmp_path, builder, frontend):
     mod_function(a, b)
     assert np.all(a == [1] * 10)
     assert np.all(b == range(1,11))
+
+
+@pytest.mark.parametrize('frontend', available_frontends())
+def test_outline_pragma_regions_derived_args(tmp_path, builder, frontend):
+    """
+    Test subroutine extraction with derived-type arguments.
+    """
+
+    fcode = """
+module test_outline_dertype_mod
+  implicit none
+
+  type rick
+    integer :: a(10), b(10)
+  end type rick
+contains
+
+  subroutine test_outline_imps(a, b)
+    integer, intent(out) :: a(10), b(10)
+    type(rick) :: dave
+    integer :: j
+
+    dave%a(:) = a(:)
+    dave%b(:) = b(:)
+
+!$loki outline
+    do j=1,10
+      dave%a(j) = j + 1
+    end do
+
+    dave%b(:) = dave%b(:) + 42
+!$loki end outline
+
+    a(:) = dave%a(:)
+    b(:) = dave%b(:)
+  end subroutine test_outline_imps
+end module test_outline_dertype_mod
+"""
+    module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path])
+    refname = f'ref_{module.name}_{frontend}'
+    reference = jit_compile_lib([module], path=tmp_path, name=refname, builder=builder)
+    function = getattr(getattr(reference, module.name), module.subroutines[0].name)
+
+    # Test the reference solution
+    a = np.zeros(shape=(10,), dtype=np.int32)
+    b = np.zeros(shape=(10,), dtype=np.int32)
+    function(a, b)
+    assert np.all(a == range(2,12))
+    assert np.all(b == 42)
+    (tmp_path/f'{module.name}.f90').unlink()
+
+    assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 6
+    assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 0
+
+    # Apply transformation
+    routines = outline_pragma_regions(module.subroutines[0])
+
+    assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 4
+    assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 1
+
+    # Check for a single derived-type argument
+    assert len(routines) == 1
+    assert len(routines[0].arguments) == 1
+    assert routines[0].arguments[0] == 'dave'
+    assert routines[0].arguments[0].type.dtype.name == 'rick'
+    assert routines[0].arguments[0].type.intent == 'inout'
+
+    # Insert created routines into module
+    module.contains.append(routines)
+
+    obj = jit_compile_lib([module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder)
+    mod_function = getattr(getattr(obj, module.name), module.subroutines[0].name)
+
+    # Test the transformed module solution
+    a = np.zeros(shape=(10,), dtype=np.int32)
+    b = np.zeros(shape=(10,), dtype=np.int32)
+    mod_function(a, b)
+    assert np.all(a == range(2,12))
+    assert np.all(b == 42)

From 5d81e253ef85805af94ba843a3ccfc2dec3dcfba Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Fri, 18 Oct 2024 18:16:11 +0000
Subject: [PATCH 06/11] Extract: Resolve intent overrides from caller when
 extracting

---
 loki/transformations/extract/outline.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 33ad9e5c6..3165c35ed 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -73,9 +73,9 @@ def outline_region(region, name, imports, intent_map=None):
         for v in FindVariables().visit(region.body)
         if v.clone(dimensions=None) not in imported_symbols
     )
-    pragma_in_args = {region_var_map[v.lower()] for v in intent_map.get('in', '').split(',') if v}
-    pragma_inout_args = {region_var_map[v.lower()] for v in intent_map.get('inout', '').split(',') if v}
-    pragma_out_args = {region_var_map[v.lower()] for v in intent_map.get('out', '').split(',') if v}
+    pragma_in_args = {v.clone(scope=region_routine) for v in intent_map['in']}
+    pragma_inout_args = {v.clone(scope=region_routine) for v in intent_map['inout']}
+    pragma_out_args = {v.clone(scope=region_routine) for v in intent_map['out']}
 
     # Override arguments according to pragma annotations
     region_in_args = (region_in_args - (pragma_inout_args | pragma_out_args)) | pragma_in_args
@@ -109,8 +109,12 @@ def outline_region(region, name, imports, intent_map=None):
     region_routine_arguments = []
     for intent, args in zip(('in', 'inout', 'out'), (region_in_args, region_inout_args, region_out_args)):
         for arg in args:
-            local_var = region_routine_var_map[arg.name]
-            local_var = local_var.clone(type=local_var.type.clone(intent=intent))
+            local_var = region_routine_var_map.get(arg.name, arg)
+            # Sanitise argument types
+            local_var = local_var.clone(
+                type=local_var.type.clone(intent=intent, allocatable=None, target=None)
+            )
+
             region_routine_var_map[arg.name] = local_var
             region_routine_arguments += [local_var]
 
@@ -155,6 +159,7 @@ def outline_pragma_regions(routine):
     counter = 0
     routines = []
     imports = FindNodes(Import).visit(routine.spec)
+    parent_vmap = routine.variable_map
     mapper = {}
     with pragma_regions_attached(routine):
         with dataflow_analysis_attached(routine):
@@ -167,7 +172,13 @@ def outline_pragma_regions(routine):
                 name = parameters.get('name', f'{routine.name}_outlined_{counter}')
                 counter += 1
 
-                call, region_routine = outline_region(region, name, imports, intent_map=parameters)
+                # Extract explicitly requested symbols from context
+                intent_map = {}
+                intent_map['in'] = tuple(parent_vmap[v.lower()] for v in parameters.get('in', '').split(',') if v)
+                intent_map['inout'] = tuple(parent_vmap[v.lower()] for v in parameters.get('inout', '').split(',') if v)
+                intent_map['out'] = tuple(parent_vmap[v.lower()] for v in parameters.get('out', '').split(',') if v)
+
+                call, region_routine = outline_region(region, name, imports, intent_map=intent_map)
 
                 # insert into list of new routines
                 routines.append(region_routine)

From e2aa27f35fe462c2e56e64360c71addd93fd90bc Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Tue, 1 Oct 2024 08:17:16 +0000
Subject: [PATCH 07/11] Extract: Add test for outlininig extraction with
 associates

---
 .../extract/tests/test_outline.py             | 91 +++++++++++++++++++
 1 file changed, 91 insertions(+)

diff --git a/loki/transformations/extract/tests/test_outline.py b/loki/transformations/extract/tests/test_outline.py
index a9ad85419..2a35c32b1 100644
--- a/loki/transformations/extract/tests/test_outline.py
+++ b/loki/transformations/extract/tests/test_outline.py
@@ -13,6 +13,7 @@
 from loki.frontend import available_frontends
 from loki.ir import FindNodes, Section, Assignment, CallStatement, Intrinsic
 from loki.tools import as_tuple
+from loki.types import BasicType
 
 from loki.transformations.extract.outline import outline_pragma_regions
 
@@ -440,3 +441,93 @@ def test_outline_pragma_regions_derived_args(tmp_path, builder, frontend):
     mod_function(a, b)
     assert np.all(a == range(2,12))
     assert np.all(b == 42)
+
+
+@pytest.mark.parametrize('frontend', available_frontends())
+def test_outline_pragma_regions_associates(tmp_path, builder, frontend):
+    """
+    Test subroutine extraction with derived-type arguments.
+    """
+
+    fcode = """
+module test_outline_assoc_mod
+  implicit none
+
+  type rick
+    integer :: a(10), b(10)
+  end type rick
+contains
+
+  subroutine test_outline_imps(a, b)
+    integer, intent(out) :: a(10), b(10)
+    type(rick) :: dave
+    integer :: j
+
+    associate(c=>dave%a, d=>dave%b)
+
+    c(:) = a(:)
+    d(:) = b(:)
+
+!$loki outline
+    do j=1,10
+      c(j) = j + 1
+    end do
+
+    d(:) = d(:) + 42
+!$loki end outline
+
+    a(:) = c(:)
+    b(:) = d(:)
+    end associate
+  end subroutine test_outline_imps
+end module test_outline_assoc_mod
+"""
+    module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path])
+    routine = module.subroutines[0]
+    refname = f'ref_{module.name}_{frontend}'
+    reference = jit_compile_lib([module], path=tmp_path, name=refname, builder=builder)
+    function = getattr(getattr(reference, module.name), routine.name)
+
+    # Test the reference solution
+    a = np.zeros(shape=(10,), dtype=np.int32)
+    b = np.zeros(shape=(10,), dtype=np.int32)
+    function(a, b)
+    assert np.all(a == range(2,12))
+    assert np.all(b == 42)
+    (tmp_path/f'{module.name}.f90').unlink()
+
+    assert len(FindNodes(Assignment).visit(routine.body)) == 6
+    assert len(FindNodes(CallStatement).visit(routine.body)) == 0
+
+    # Apply transformation
+    outlined = outline_pragma_regions(routine)
+
+    assert len(FindNodes(Assignment).visit(routine.body)) == 4
+    calls = FindNodes(CallStatement).visit(routine.body)
+    assert len(calls) == 1
+    assert calls[0].arguments == ('d', 'c')
+
+    # Check for a single derived-type argument
+    assert len(outlined) == 1
+    assert len(outlined[0].arguments) == 2
+    assert outlined[0].arguments[0].name == 'd'
+    assert outlined[0].arguments[0].type.shape == (10,)
+    assert outlined[0].arguments[0].type.dtype == BasicType.INTEGER
+    assert outlined[0].arguments[0].type.intent == 'inout'
+    assert outlined[0].arguments[1].name == 'c'
+    assert outlined[0].arguments[1].type.shape == (10,)
+    assert outlined[0].arguments[1].type.dtype == BasicType.INTEGER
+    assert outlined[0].arguments[1].type.intent == 'out'
+
+    # Insert created routines into module
+    module.contains.append(outlined)
+
+    obj = jit_compile_lib(
+        [module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder
+    )
+    mod_function = getattr(getattr(obj, module.name), routine.name)
+    a = np.zeros(shape=(10,), dtype=np.int32)
+    b = np.zeros(shape=(10,), dtype=np.int32)
+    mod_function(a, b)
+    assert np.all(a == range(2,12))
+    assert np.all(b == 42)

From fd0242fe8b599ad95f911482edb8ee8884b9cd86 Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Tue, 8 Oct 2024 03:43:43 +0000
Subject: [PATCH 08/11] Extract: Order declarations when extracting routines

---
 loki/transformations/extract/outline.py       | 72 ++++++++++++-------
 .../extract/tests/test_outline.py             | 10 +--
 2 files changed, 53 insertions(+), 29 deletions(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 3165c35ed..78dbcfc15 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -6,7 +6,7 @@
 # nor does it submit to any jurisdiction.
 
 from loki.analyse import dataflow_analysis_attached
-from loki.expression import Variable
+from loki.expression import symbols as sym, Variable
 from loki.ir import (
     CallStatement, Import, PragmaRegion, Section, FindNodes,
     FindVariables, Transformer, is_loki_pragma,
@@ -14,13 +14,40 @@
 )
 from loki.logging import info
 from loki.subroutine import Subroutine
-from loki.tools import as_tuple, CaseInsensitiveDict
+from loki.tools import as_tuple
+from loki.types import BasicType, DerivedType
 
 
 
 __all__ = ['outline_region', 'outline_pragma_regions']
 
 
+def order_variables_by_type(variables, imports=None):
+    """
+    Apply a default ordering to variables based on their type, so that
+    their use in declaration lists is unified.
+    """
+    variables = sorted(variables, key=str)  # Lexicographical base order
+
+    derived = tuple(
+        v for v in variables
+        if isinstance(v.type.dtype, DerivedType) or v.type.dtype == BasicType.DEFERRED
+    )
+
+    if imports:
+        # Order derived types by the order of their type in imports
+        imported_symbols = tuple(s for i in imports for s in i.symbols if not i.c_import)
+        derived = tuple(sorted(derived, key=lambda x: imported_symbols.index(x.type.dtype.name)))
+
+    # Order declarations by type and put arrays before scalars
+    non_derived = tuple(v for v in variables if v not in derived)
+    arrays = tuple(v for v in non_derived if isinstance(v, sym.Array))
+    scalars = tuple(v for v in non_derived if isinstance(v, sym.Scalar))
+    assert len(derived) + len(arrays) + len(scalars) == len(variables)
+
+    return derived + arrays + scalars
+
+
 def outline_region(region, name, imports, intent_map=None):
     """
     Creates a new :any:`Subroutine` object from a given :any:`PragmaRegion`.
@@ -68,11 +95,6 @@ def outline_region(region, name, imports, intent_map=None):
     region_in_args = {arg for arg in region_in_args if not arg.type.parameter}
 
     # Extract arguments given in pragma annotations
-    region_var_map = CaseInsensitiveDict(
-        (v.name, v.clone(dimensions=None))
-        for v in FindVariables().visit(region.body)
-        if v.clone(dimensions=None) not in imported_symbols
-    )
     pragma_in_args = {v.clone(scope=region_routine) for v in intent_map['in']}
     pragma_inout_args = {v.clone(scope=region_routine) for v in intent_map['inout']}
     pragma_out_args = {v.clone(scope=region_routine) for v in intent_map['out']}
@@ -89,23 +111,18 @@ def outline_region(region, name, imports, intent_map=None):
 
     # Set the list of variables used in region routine (to create declarations)
     # and put all in the new scope
-    region_routine_variables = {
-        v.clone(dimensions=v.type.shape or None)
-        for v in FindVariables().visit(region_routine.body)
-        if v.name in region_var_map
-    }
+    region_routine_variables = tuple(
+        v.clone(dimensions=v.type.shape or None, scope=region_routine)
+        for v in FindVariables().visit(region.body)
+        if v.clone(dimensions=None) not in imported_symbols
+    )
     # Filter out derived-type component variables from declarations
-    region_routine_variables = {
+    region_routine_variables = tuple(
         v.parents[0] if v.parent else v for v in region_routine_variables
-    }
-    # Order the local devlaration list to put arguments first
-    region_routine_args = tuple(v for v in region_routine_variables if v.type.intent)
-    region_routine_locals = tuple(v for v in region_routine_variables if not v.type.intent)
-    region_routine.variables = region_routine_args + region_routine_locals
-    region_routine.rescope_symbols()
+    )
 
     # Build the call signature
-    region_routine_var_map = region_routine.variable_map
+    region_routine_var_map = {v.name: v for v in region_routine_variables}
     region_routine_arguments = []
     for intent, args in zip(('in', 'inout', 'out'), (region_in_args, region_inout_args, region_out_args)):
         for arg in args:
@@ -118,12 +135,19 @@ def outline_region(region, name, imports, intent_map=None):
             region_routine_var_map[arg.name] = local_var
             region_routine_arguments += [local_var]
 
-    # We need to update the list of variables again to avoid duplicate declarations
-    region_routine.variables = as_tuple(region_routine_var_map.values())
-    region_routine.arguments = as_tuple(region_routine_arguments)
+    # Order the arguments and local declaration lists and put arguments first
+    region_routine_locals = tuple(
+        v for v in region_routine_variables if not v in region_routine_arguments
+    )
+    region_routine_arguments = order_variables_by_type(region_routine_arguments, imports=imports)
+    region_routine_locals = order_variables_by_type(region_routine_locals, imports=imports)
+
+    region_routine.variables = region_routine_arguments + region_routine_locals
+    region_routine.arguments = region_routine_arguments
 
     # Create the call according to the wrapped code region
-    call_arguments = region_in_args + region_inout_args + region_out_args
+    call_arg_map = {v.name: v for v in region_in_args + region_inout_args + region_out_args}
+    call_arguments = tuple(call_arg_map[a.name] for a in region_routine_arguments)
     call = CallStatement(name=Variable(name=name), arguments=call_arguments, kwarguments=())
 
     return call, region_routine
diff --git a/loki/transformations/extract/tests/test_outline.py b/loki/transformations/extract/tests/test_outline.py
index 2a35c32b1..c6ba5ffa0 100644
--- a/loki/transformations/extract/tests/test_outline.py
+++ b/loki/transformations/extract/tests/test_outline.py
@@ -505,19 +505,19 @@ def test_outline_pragma_regions_associates(tmp_path, builder, frontend):
     assert len(FindNodes(Assignment).visit(routine.body)) == 4
     calls = FindNodes(CallStatement).visit(routine.body)
     assert len(calls) == 1
-    assert calls[0].arguments == ('d', 'c')
+    assert calls[0].arguments == ('c', 'd')
 
     # Check for a single derived-type argument
     assert len(outlined) == 1
     assert len(outlined[0].arguments) == 2
-    assert outlined[0].arguments[0].name == 'd'
+    assert outlined[0].arguments[0].name == 'c'
     assert outlined[0].arguments[0].type.shape == (10,)
     assert outlined[0].arguments[0].type.dtype == BasicType.INTEGER
-    assert outlined[0].arguments[0].type.intent == 'inout'
-    assert outlined[0].arguments[1].name == 'c'
+    assert outlined[0].arguments[0].type.intent == 'out'
+    assert outlined[0].arguments[1].name == 'd'
     assert outlined[0].arguments[1].type.shape == (10,)
     assert outlined[0].arguments[1].type.dtype == BasicType.INTEGER
-    assert outlined[0].arguments[1].type.intent == 'out'
+    assert outlined[0].arguments[1].type.intent == 'inout'
 
     # Insert created routines into module
     module.contains.append(outlined)

From 76dc83d386516e717f1236ff00b80abc387e6a82 Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Wed, 9 Oct 2024 04:47:42 +0000
Subject: [PATCH 09/11] Extract: Ensure correct re-scoping after cloning
 arguments

---
 loki/transformations/extract/outline.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 78dbcfc15..4fa0cda51 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -129,7 +129,8 @@ def outline_region(region, name, imports, intent_map=None):
             local_var = region_routine_var_map.get(arg.name, arg)
             # Sanitise argument types
             local_var = local_var.clone(
-                type=local_var.type.clone(intent=intent, allocatable=None, target=None)
+                type=local_var.type.clone(intent=intent, allocatable=None, target=None),
+                scope=region_routine
             )
 
             region_routine_var_map[arg.name] = local_var
@@ -145,6 +146,9 @@ def outline_region(region, name, imports, intent_map=None):
     region_routine.variables = region_routine_arguments + region_routine_locals
     region_routine.arguments = region_routine_arguments
 
+    # Ensure everything has been rescoped
+    region_routine.rescope_symbols()
+
     # Create the call according to the wrapped code region
     call_arg_map = {v.name: v for v in region_in_args + region_inout_args + region_out_args}
     call_arguments = tuple(call_arg_map[a.name] for a in region_routine_arguments)

From 736948ea8bae0002f7600b77153612d59f18f47e Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Thu, 7 Nov 2024 12:44:40 +0000
Subject: [PATCH 10/11] Extract: Better use of internal Loki API in
 outline_pragma_regions

---
 loki/transformations/extract/outline.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 4fa0cda51..8d22deac3 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -186,7 +186,7 @@ def outline_pragma_regions(routine):
     """
     counter = 0
     routines = []
-    imports = FindNodes(Import).visit(routine.spec)
+    imports = routine.imports
     parent_vmap = routine.variable_map
     mapper = {}
     with pragma_regions_attached(routine):
@@ -202,9 +202,9 @@ def outline_pragma_regions(routine):
 
                 # Extract explicitly requested symbols from context
                 intent_map = {}
-                intent_map['in'] = tuple(parent_vmap[v.lower()] for v in parameters.get('in', '').split(',') if v)
-                intent_map['inout'] = tuple(parent_vmap[v.lower()] for v in parameters.get('inout', '').split(',') if v)
-                intent_map['out'] = tuple(parent_vmap[v.lower()] for v in parameters.get('out', '').split(',') if v)
+                intent_map['in'] = tuple(parent_vmap[v] for v in parameters.get('in', '').split(',') if v)
+                intent_map['inout'] = tuple(parent_vmap[v] for v in parameters.get('inout', '').split(',') if v)
+                intent_map['out'] = tuple(parent_vmap[v] for v in parameters.get('out', '').split(',') if v)
 
                 call, region_routine = outline_region(region, name, imports, intent_map=intent_map)
 

From 66995cf9b5a3442ceb0be59ec2244bb65f247543 Mon Sep 17 00:00:00 2001
From: Michael Lange <michael.lange@ecmwf.int>
Date: Thu, 7 Nov 2024 13:58:32 +0100
Subject: [PATCH 11/11] Extract: Better way to filter derived types in arg
 reordering for outlining

Co-authored-by: Balthasar Reuter <6384870+reuterbal@users.noreply.github.com>
---
 loki/transformations/extract/outline.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py
index 8d22deac3..bd97d2155 100644
--- a/loki/transformations/extract/outline.py
+++ b/loki/transformations/extract/outline.py
@@ -8,14 +8,14 @@
 from loki.analyse import dataflow_analysis_attached
 from loki.expression import symbols as sym, Variable
 from loki.ir import (
-    CallStatement, Import, PragmaRegion, Section, FindNodes,
+    CallStatement, PragmaRegion, Section, FindNodes,
     FindVariables, Transformer, is_loki_pragma,
     get_pragma_parameters, pragma_regions_attached
 )
 from loki.logging import info
 from loki.subroutine import Subroutine
 from loki.tools import as_tuple
-from loki.types import BasicType, DerivedType
+from loki.types import BasicType
 
 
 
@@ -31,7 +31,7 @@ def order_variables_by_type(variables, imports=None):
 
     derived = tuple(
         v for v in variables
-        if isinstance(v.type.dtype, DerivedType) or v.type.dtype == BasicType.DEFERRED
+        if not isinstance(v, (sym.Scalar, sym.Array)) or not isinstance(v.type.dtype, BasicType)
     )
 
     if imports: