Repository: ecmwf-ifs/loki
Pull request #550: FParser: Better `Subroutine` and `Module` construction
Author: login: mlange05
Created: 2025-05-09T08:09:47Z | Merged: 2025-05-13T08:42:05Z

Another small housekeeping PR that aims to tidy up the construction of `Subroutine` and `Module` objects in the FParser frontend. The key changes are a slight generalisation of utilities to move comments between spec and body, and a common utility to pre-instantiate sub-procedure scopes.

_As a bonus, this gets rid of the last use of a naked dict-based `Transformer` in the AST traversal._ 

To make this safer, I also extended one of the `Subroutine` tests to capture the corner case that forces us to pre-instantiate internal and module procedures. We do this, so that all subroutines are created up front, so that `InlineCall`s between two internal routines are handled cleanly and the cross-calling type info is generated safely.

I also piggy-backed a small change that makes sure we use the `get_source()` utility throughout, even for node types with distinct `end_xxx` AST nodes. This all is in preparation to separate Function and Subroutine constructors, which have a lot of natural overlap. Hope this makes sense.


Diff:
diff --git a/loki/frontend/fparser.py b/loki/frontend/fparser.py
index dcda5dc05..f15f2549c 100644
--- a/loki/frontend/fparser.py
+++ b/loki/frontend/fparser.py
@@ -7,6 +7,7 @@
 
 # pylint: disable=too-many-lines
 import re
+from itertools import takewhile
 
 from codetiming import Timer
 
@@ -27,9 +28,8 @@
 
 from loki import ir
 from loki.ir import (
-    GenericVisitor, Transformer, FindNodes, attach_pragmas,
-    process_dimension_pragmas, detach_pragmas, pragmas_attached,
-    AttachScopes
+    GenericVisitor, FindNodes, AttachScopes, attach_pragmas,
+    detach_pragmas, pragmas_attached, process_dimension_pragmas
 )
 import loki.expression.symbols as sym
 from loki.expression.operations import (
@@ -229,6 +229,41 @@ def rget_child(node, node_type):
     return None
 
 
+def _get_comments_from_section(sec, include_pragmas=False, reverse=False):
+    """
+    Extract leading or trailing :any:`Comment` or `:any:`CommentBlock`
+    nodes from a :any:`Section`.
+
+    Parameters
+    ----------
+    sec : :any:`Section`
+        Code section from which to extract comment nodes
+    include_pragmas : bool
+        Flag to enable matching :any:`Pragma` nodes
+    reverse : bool
+        Flag to enable matching trailing comment nodes
+
+    Returns
+    -------
+    tuple of :any:`Node`
+        Leading or trailing comment or pragma nodes
+    """
+
+    _matches = (ir.Comment, ir.CommentBlock)
+    if include_pragmas:
+        _matches += (ir.Pragma,)
+
+    def is_comment(n):
+        return isinstance(n, _matches)
+
+    # Pick out comments from the beginning of the section and update in-place
+    nodes = reversed(sec.body) if reverse else sec.body
+    comments = tuple(takewhile(is_comment, nodes))
+    sec._update(body=tuple(filter(lambda n: n not in comments, sec.body)))
+
+    return reversed(comments) if reverse else comments
+
+
 class FParser2IR(GenericVisitor):
     # pylint: disable=unused-argument  # Stop warnings about unused arguments
 
@@ -300,6 +335,35 @@ def visit_Intrinsic_Stmt(self, o, **kwargs):
     # Base blocks
     #
 
+    def create_contained_procedures(self, o, **kwargs):
+        """
+        Helper utility that creates :any:`Subroutine` objects before
+        the full parse to ensure the scope hierarchy is in place.
+
+        Notes
+        -----
+        We first make sure the procedure objects for all internal
+        procedures are instantiated before parsing the actual spec and
+        body of the parent routine.
+
+        This way, all procedure types should exist in the scope and
+        any use of their symbol (e.g. in a :any:`CallStatement` or
+        :any:`InlineCall`) can be matched against a type.
+        """
+        if not o:
+            return
+
+        member_asts = tuple(
+            c for c in o.children
+            if isinstance(c, (Fortran2003.Subroutine_Subprogram, Fortran2003.Function_Subprogram))
+        )
+
+        # Instantiate the procedure objects from their initial "stmt
+        # line" to fill the type cache of the scope. This is needed to
+        # get `ProcedureType` objecst and identify `InlineCall` objects.
+        for c in member_asts:
+            self.visit(get_child(c, (Fortran2003.Subroutine_Stmt, Fortran2003.Function_Stmt)), **kwargs)
+
     def visit_Specification_Part(self, o, **kwargs):
         """
         The specification part of a program-unit
@@ -1496,9 +1560,7 @@ def visit_Associate_Construct(self, o, **kwargs):
         pre = as_tuple(self.visit(c, **kwargs) for c in o.children[:assoc_stmt_index])
 
         # Extract source object for construct
-        lines = (assoc_stmt.item.span[0], end_assoc_stmt.item.span[1])
-        string = ''.join(self.raw_source[lines[0]-1:lines[1]]).strip('\n')
-        source = Source(lines=lines, string=string)
+        source = self.get_source(assoc_stmt, end_node=end_assoc_stmt)
 
         # Handle the associates
         associations = self.visit(assoc_stmt, **kwargs)
@@ -1581,9 +1643,7 @@ def visit_Interface_Block(self, o, **kwargs):
         pre = as_tuple(self.visit(c, **kwargs) for c in o.children[:interface_stmt_index])
 
         # Extract source object for construct
-        lines = (interface_stmt.item.span[0], end_interface_stmt.item.span[1])
-        string = ''.join(self.raw_source[lines[0]-1:lines[1]]).strip('\n')
-        source = Source(lines=lines, string=string)
+        source = self.get_source(interface_stmt, end_node=end_interface_stmt)
 
         # The interface spec
         abstract = False
@@ -1733,29 +1793,10 @@ def visit_Subroutine_Subprogram(self, o, **kwargs):
         kwargs['scope'] = routine
 
         # Extract source object for construct
-        lines = (subroutine_stmt.item.span[0], end_subroutine_stmt.item.span[1])
-        string = ''.join(self.raw_source[lines[0]-1:lines[1]]).strip('\n')
-        source = Source(lines=lines, string=string)
-
-        # We make sure the subroutine objects for all member routines are
-        # instantiated before parsing the actual spec and body of the parent routine.
-        # This way, all procedure types should exist and any use of their symbol
-        # (e.g. in a CallStatement) should have type.dtype.procedure initialized correctly
-        contains_ast = get_child(o, Fortran2003.Internal_Subprogram_Part)
-        if contains_ast is not None:
-            member_asts = [
-                c for c in contains_ast.children
-                if isinstance(c, (Fortran2003.Subroutine_Subprogram, Fortran2003.Function_Subprogram))
-            ]
-            # Note that we overwrite this variable subsequently with the fully parsed subroutines
-            # where the visit-method for the subroutine/function statement will pick out the existing
-            # subroutine objects using the weakref pointers stored in the symbol table.
-            # I know, it's not pretty but alternatively we could hand down this array as part of
-            # kwargs but that feels like carrying around a lot of bulk, too.
-            contains = [
-                self.visit(get_child(c, (Fortran2003.Subroutine_Stmt, Fortran2003.Function_Stmt)), **kwargs)[0]
-                for c in member_asts
-            ]
+        source = self.get_source(subroutine_stmt, end_node=end_subroutine_stmt)
+
+        # Pre-populate internal procedure scopes in the type hierarchy
+        self.create_contained_procedures(get_child(o, Fortran2003.Internal_Subprogram_Part), **kwargs)
 
         # Hack: Collect all spec and body parts and use all but the
         # last body as spec. Reason is that Fparser misinterprets statement
@@ -1804,10 +1845,7 @@ def visit_Subroutine_Subprogram(self, o, **kwargs):
                 spec.insert(spec.body.index(decls[0]), return_var_decl)
 
         # Now all declarations are well-defined and we can parse the member routines
-        if contains_ast is not None:
-            contains = self.visit(contains_ast, **kwargs)
-        else:
-            contains = None
+        contains = self.visit(get_child(o, Fortran2003.Internal_Subprogram_Part), **kwargs)
 
         # Finally, take care of the body
         if body_ast is None:
@@ -1825,24 +1863,11 @@ def visit_Subroutine_Subprogram(self, o, **kwargs):
             spec._update(body=spec.body + body.body[:idx])
             body._update(body=body.body[idx:])
 
-        # Another big hack: fparser allocates all comments before and after the
-        # spec to the spec. We remove them from the beginning to get the docstring.
-        comment_map = {}
-        docs = []
-        for node in spec.body:
-            if not isinstance(node, (ir.Comment, ir.CommentBlock)):
-                break
-            docs.append(node)
-            comment_map[node] = None
+        # Extract the leading comments of the specification as "docstring" section
+        docs = _get_comments_from_section(spec) if spec else ()
 
-        # Now we move comments from the end to the beginning of the body as those
-        # can potentially be pragmas.
-        for node in reversed(spec.body):
-            if not isinstance(node, (ir.Pragma, ir.Comment, ir.CommentBlock)):
-                break
-            body.prepend(node)
-            comment_map[node] = None
-        spec = Transformer(comment_map, invalidate_source=False).visit(spec)
+        # Move trailing comments from spec to the body as those can be pragmas.
+        body.prepend(_get_comments_from_section(spec, include_pragmas=True, reverse=True))
 
         # Finally, call the subroutine constructor on the object again to register all
         # bits and pieces in place and rescope all symbols
@@ -2026,68 +2051,33 @@ def visit_Module(self, o, **kwargs):
         assert end_module_stmt_index + 1 == len(o.children)
 
         # Extract source object for construct
-        lines = (module_stmt.item.span[0], end_module_stmt.item.span[1])
-        string = ''.join(self.raw_source[lines[0]-1:lines[1]]).strip('\n')
-        source = Source(lines=lines, string=string)
+        source = self.get_source(module_stmt, end_node=end_module_stmt)
 
         # Instantiate the object
         module = self.visit(module_stmt, **kwargs)
         kwargs['scope'] = module
 
-        # We make sure the subroutine objects for all member routines are
-        # instantiated before parsing the actual spec of the module.
-        # This way, all procedure types should exist and any use of their symbol
-        # (e.g. in a derived type procedure binding etc) should be initialized correctly
-        contains_ast = get_child(o, Fortran2003.Module_Subprogram_Part)
-        if contains_ast is not None:
-            member_asts = [
-                c for c in contains_ast.children
-                if isinstance(c, (Fortran2003.Subroutine_Subprogram, Fortran2003.Function_Subprogram))
-            ]
-            # Note that we overwrite this variable subsequently with the fully parsed subroutines
-            # where the visit-method for the subroutine/function statement will pick out the existing
-            # subroutine objects using the weakref pointers stored in the symbol table.
-            # I know, it's not pretty but alternatively we could hand down this array as part of
-            # kwargs but that feels like carrying around a lot of bulk, too.
-            contains = [
-                self.visit(get_child(c, (Fortran2003.Subroutine_Stmt, Fortran2003.Function_Stmt)), **kwargs)[0]
-                for c in member_asts
-            ]
+        # Pre-populate internal procedure scopes in the type hierarchy
+        self.create_contained_procedures(get_child(o, Fortran2003.Module_Subprogram_Part), **kwargs)
 
         # Build the spec
-        spec_ast = get_child(o, Fortran2003.Specification_Part)
-        if spec_ast:
-            spec = self.visit(spec_ast, **kwargs)
-            spec = sanitize_ir(spec, FP, pp_registry=sanitize_registry[FP], pp_info=self.pp_info)
-
-            # Infer any additional shape information from `!$loki dimension` pragmas
-            spec = attach_pragmas(spec, ir.VariableDeclaration)
-            spec = process_dimension_pragmas(spec)
-            spec = detach_pragmas(spec, ir.VariableDeclaration)
-
-            # Another big hack: fparser allocates all comments before and after the
-            # spec to the spec. We remove them from the beginning to get the docstring.
-            comment_map = {}
-            docs = []
-            for node in spec.body:
-                if not isinstance(node, (ir.Comment, ir.CommentBlock)):
-                    break
-                docs.append(node)
-                comment_map[node] = None
-            spec = Transformer(comment_map, invalidate_source=False).visit(spec)
-        else:
-            docs = []
-            spec = None
+        spec = self.visit(get_child(o, Fortran2003.Specification_Part), **kwargs)
+        spec = sanitize_ir(spec, FP, pp_registry=sanitize_registry[FP], pp_info=self.pp_info)
+
+        # Infer any additional shape information from `!$loki dimension` pragmas
+        spec = attach_pragmas(spec, ir.VariableDeclaration)
+        spec = process_dimension_pragmas(spec)
+        spec = detach_pragmas(spec, ir.VariableDeclaration)
+
+        # Extract the leading comments of the specification as "docstring" section
+        docs = _get_comments_from_section(spec) if spec else ()
 
         # As variables may be defined out of sequence, we need to re-generate
         # symbols in the spec part to make them coherent with the symbol table
         spec = AttachScopes().visit(spec, scope=module, recurse_to_declaration_attributes=True)
 
         # Now that all declarations are well-defined we can parse the member routines
-        if contains_ast is not None:
-            contains = self.visit(contains_ast, **kwargs)
-        else:
-            contains = None
+        contains = self.visit(get_child(o, Fortran2003.Module_Subprogram_Part), **kwargs)
 
         # Finally, call the module constructor on the object again to register all
         # bits and pieces in place and rescope all symbols
@@ -2258,9 +2248,7 @@ def visit_Case_Construct(self, o, **kwargs):
         pre = as_tuple(self.visit(c, **kwargs) for c in o.children[:select_case_stmt_index])
 
         # Extract source object for construct
-        lines = (select_case_stmt.item.span[0], end_select_stmt.item.span[1])
-        string = ''.join(self.raw_source[lines[0]-1:lines[1]]).strip('\n')
-        source = Source(lines=lines, string=string)
+        source = self.get_source(select_case_stmt, end_node=end_select_stmt)
 
         # Handle the SELECT CASE statement
         expr = self.visit(select_case_stmt, **kwargs)
@@ -2378,9 +2366,7 @@ def visit_Select_Type_Construct(self, o, **kwargs):
         pre = as_tuple(self.visit(c, **kwargs) for c in o.children[:select_type_stmt_index])
 
         # Extract source object for construct
-        lines = (select_type_stmt.item.span[0], end_select_stmt.item.span[1])
-        string = ''.join(self.raw_source[lines[0]-1:lines[1]]).strip('\n')
-        source = Source(lines=lines, string=string)
+        source = self.get_source(select_type_stmt, end_node=end_select_stmt)
 
         # Handle the SELECT TYPE statement
         expr = self.visit(select_type_stmt, **kwargs)
@@ -2903,9 +2889,7 @@ def visit_Where_Construct(self, o, **kwargs):
         pre = as_tuple(self.visit(c, **kwargs) for c in o.children[:where_stmt_index])
 
         # Extract source object for construct
-        lines = (where_stmt.item.span[0], end_where_stmt.item.span[1])
-        string = ''.join(self.raw_source[lines[0]-1:lines[1]]).strip('\n')
-        source = Source(lines=lines, string=string)
+        source = self.get_source(where_stmt, end_node=end_where_stmt)
 
         # Find all ELSEWHERE statements
         where_stmts, where_stmts_index = zip(*(
diff --git a/loki/tests/test_sourcefile.py b/loki/tests/test_sourcefile.py
index 69af37b7f..8d890ee53 100644
--- a/loki/tests/test_sourcefile.py
+++ b/loki/tests/test_sourcefile.py
@@ -346,7 +346,10 @@ def test_sourcefile_lazy_comments(frontend):
 
     assert isinstance(source.ir.body[0], Comment)
     assert isinstance(source.ir.body[2], Comment)
-    assert isinstance(myroutine.body.body[0], Comment)
+    if frontend == OMNI:
+        assert isinstance(myroutine.body.body[0], Comment)
+    else:
+        assert isinstance(myroutine.docstring[0], Comment)
 
     code = source.to_fortran()
     assert '! Comment outside' in code
diff --git a/loki/tests/test_subroutine.py b/loki/tests/test_subroutine.py
index d5d171f02..4f2065e66 100644
--- a/loki/tests/test_subroutine.py
+++ b/loki/tests/test_subroutine.py
@@ -15,7 +15,7 @@
 from loki.frontend import available_frontends, OMNI, REGEX
 from loki.ir import (
     nodes as ir, FindNodes, FindVariables, FindTypedSymbols,
-    Transformer
+    FindInlineCalls, Transformer
 )
 from loki.types import (
     BasicType, DerivedType, ProcedureType, SymbolAttributes
@@ -885,7 +885,7 @@ def test_member_procedures(tmp_path, frontend):
     integer, intent(in) :: in1
     integer, intent(out) :: out1
 
-    out1 = 5 * in1 + localvar
+    out1 = 5 * in1 + localvar + member_function(1)
   end subroutine member_procedure
 
   ! Below is disabled because f90wrap (wrongly) exhibits that
@@ -915,6 +915,12 @@ def test_member_procedures(tmp_path, frontend):
     assert routine.symbol_attrs.lookup('in1') is not None
     assert routine.members[0].get_symbol_scope('in1') is routine.members[0]
 
+    # Check that inline function is correctly identified
+    inline_calls = list(FindInlineCalls().visit(routine.members[0].body))
+    assert len(inline_calls) == 1
+    assert inline_calls[0].function.name == 'member_function'
+    assert inline_calls[0].function.type.dtype.procedure == routine.members[1]
+
     assert routine.members[1].name == 'member_function'
     assert routine.members[1].symbol_attrs.lookup('in2') is not None
     assert routine.members[1].get_symbol_scope('in2') is routine.members[1]
@@ -927,8 +933,8 @@ def test_member_procedures(tmp_path, frontend):
 
     # Test results of the generated and compiled code
     out1, out2 = function(1, 2)
-    assert out1 == 7
-    assert out2 == 23
+    assert out1 == 12
+    assert out2 == 38
     clean_test(filepath)
 
 


Reviews:
user: login: reuterbal
body: Very nice, no objections! Small stylistic comments/questions, otherwise all good from my side
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:06:43Z

Review comments:
diff_hunk: @@ -300,6 +335,35 @@ def visit_Intrinsic_Stmt(self, o, **kwargs):
     # Base blocks
     #
 
+    def create_contained_procedures(self, o, **kwargs):
+        """
+        Helper utility that creates :any:`Subroutine` objects before
+        the full parse to ensure the scope hierarchy is in place.
+
+        Notes
+        -----
+        We first make sure the procedure objects for all internal
+        procedures are instantiated before parsing the actual spec and
+        body of the parent routine.
+
+        This way, all procedure types should exist in the scope and
+        any use of their symbol (e.g. in a :any:`CallStatement` or
+        :any:`InlineCall`) can be matched against a type.
path: loki/frontend/fparser.py
user: login: reuterbal
body: What can `o` be in this case? I presume Module, Subroutine or Function?
created_at: 2025-05-12T19:55:43Z
author_association: COLLABORATOR
diff_hunk: @@ -1825,24 +1863,11 @@ def visit_Subroutine_Subprogram(self, o, **kwargs):
             spec._update(body=spec.body + body.body[:idx])
             body._update(body=body.body[idx:])
 
-        # Another big hack: fparser allocates all comments before and after the
-        # spec to the spec. We remove them from the beginning to get the docstring.
-        comment_map = {}
-        docs = []
-        for node in spec.body:
-            if not isinstance(node, (ir.Comment, ir.CommentBlock)):
-                break
-            docs.append(node)
-            comment_map[node] = None
+        # Extract the leading comments of the specification as "docstring" section
+        docs = _get_comments_from_section(spec) if spec else ()
 
-        # Now we move comments from the end to the beginning of the body as those
-        # can potentially be pragmas.
-        for node in reversed(spec.body):
-            if not isinstance(node, (ir.Pragma, ir.Comment, ir.CommentBlock)):
-                break
-            body.prepend(node)
-            comment_map[node] = None
-        spec = Transformer(comment_map, invalidate_source=False).visit(spec)
+        # Move trailing comments from spec to the body as those can be pragmas.
path: loki/frontend/fparser.py
user: login: reuterbal
body: [no action only and observation] The comment's not 100% accurate anymore. We're moving comments and directives, but the latter are now identified as different node types (which wasn't the case at the time this was first built, iirc). The conceptual step is still required, though, and we can't bail out early either, because you might have interleaved comments and preprocesser directives and you wouldn't want to miss any of the latter
created_at: 2025-05-12T20:03:10Z
author_association: COLLABORATOR
diff_hunk: @@ -229,6 +229,41 @@ def rget_child(node, node_type):
     return None
 
 
+def _get_comments_from_section(sec, include_pragmas=False, reverse=False):
path: loki/frontend/fparser.py
user: login: reuterbal
body: You could consider making this a common frontend utility. We're doing the same hack for OMNI: https://github.com/ecmwf-ifs/loki/blob/main/loki/frontend/omni.py#L495
created_at: 2025-05-12T20:05:50Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-09T08:14:25Z
body: Documentation for this branch can be viewed at https://sites.ecmwf.int/docs/loki/550/index.html
author_association: NONE
user: login: codecov[bot]
created_at: 2025-05-10T04:48:12Z
body: ## [Codecov](https://app.codecov.io/gh/ecmwf-ifs/loki/pull/550?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 96.22%. Comparing base [(`9f3c1b4`)](https://app.codecov.io/gh/ecmwf-ifs/loki/commit/9f3c1b4d5bcf5305f1753c3a739ebfce86b4126a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs) to head [(`444c5f9`)](https://app.codecov.io/gh/ecmwf-ifs/loki/commit/444c5f9f9b772be9ae91aec76d6474f6d037a4d9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs).
> Report is 15 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #550      +/-   ##
==========================================
- Coverage   96.22%   96.22%   -0.01%     
==========================================
  Files         235      235              
  Lines       42141    42118      -23     
==========================================
- Hits        40552    40528      -24     
- Misses       1589     1590       +1     
```

| [Flag](https://app.codecov.io/gh/ecmwf-ifs/loki/pull/550/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs) | Coverage Δ | |
|---|---|---|
| [lint_rules](https://app.codecov.io/gh/ecmwf-ifs/loki/pull/550/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs) | `96.40% <ø> (ø)` | |
| [loki](https://app.codecov.io/gh/ecmwf-ifs/loki/pull/550/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs) | `96.22% <100.00%> (-0.01%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ecmwf-ifs/loki/pull/550?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf-ifs).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE