From 70c33dd913376cd87b46887da51da2b5d939f10f Mon Sep 17 00:00:00 2001
From: Marcin Copik <mcopik@gmail.com>
Date: Fri, 8 Sep 2023 23:04:41 +0200
Subject: [PATCH] Add support for Fortran modules in scope parent assignment
 pass

---
 dace/frontend/fortran/ast_internal_classes.py |  3 +-
 dace/frontend/fortran/ast_transforms.py       |  3 +-
 tests/fortran/parent_test.py                  | 37 +++++++++++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/dace/frontend/fortran/ast_internal_classes.py b/dace/frontend/fortran/ast_internal_classes.py
index 171b941858..70a43e21b8 100644
--- a/dace/frontend/fortran/ast_internal_classes.py
+++ b/dace/frontend/fortran/ast_internal_classes.py
@@ -15,7 +15,8 @@ def __init__(self, *args, **kwargs):  # real signature unknown
             Union[
                 Subroutine_Subprogram_Node,
                 Function_Subprogram_Node,
-                Main_Program_Node
+                Main_Program_Node,
+                Module_Node
             ]
         ] = None
         for k, v in kwargs.items():
diff --git a/dace/frontend/fortran/ast_transforms.py b/dace/frontend/fortran/ast_transforms.py
index 24ac6edeca..e2a7246aed 100644
--- a/dace/frontend/fortran/ast_transforms.py
+++ b/dace/frontend/fortran/ast_transforms.py
@@ -326,7 +326,8 @@ def visit(self, node: ast_internal_classes.FNode, parent_node: Optional[ast_inte
         parent_node_types = [
             ast_internal_classes.Subroutine_Subprogram_Node,
             ast_internal_classes.Function_Subprogram_Node,
-            ast_internal_classes.Main_Program_Node
+            ast_internal_classes.Main_Program_Node,
+            ast_internal_classes.Module_Node
         ]
 
         if parent_node is not None and type(parent_node) in parent_node_types:
diff --git a/tests/fortran/parent_test.py b/tests/fortran/parent_test.py
index e68f03db8c..b1d08eaf37 100644
--- a/tests/fortran/parent_test.py
+++ b/tests/fortran/parent_test.py
@@ -48,7 +48,44 @@ def test_fortran_frontend_parent():
         for execution in subroutine.execution_part.execution:
             assert execution.parent == subroutine
 
+def test_fortran_frontend_module():
+    """
+    Tests that the Fortran frontend can parse array accesses and that the accessed indices are correct.
+    """
+    test_string = """
+                    module test_module
+                    implicit none
+                        ! good enough approximation
+                        integer, parameter :: pi = 4
+                    end module test_module
+
+                    PROGRAM access_test
+                    implicit none
+                    double precision d(4)
+                    d(1)=0
+                    CALL array_access_test_function(d)
+                    end
+
+                    SUBROUTINE array_access_test_function(d)
+                    double precision d(4)
+
+                    d(2)=5.5
+
+                    END SUBROUTINE array_access_test_function
+                    """
+    ast, functions = fortran_parser.create_ast_from_string(test_string, "array_access_test")
+    ast_transforms.ParentScopeAssigner().visit(ast)
+
+    assert ast.parent is None
+    assert ast.main_program.parent == None
+
+    module = ast.modules[0]
+    assert module.parent == None
+    specification = module.specification_part.specifications[0]
+    assert specification.parent == module
+
 
 if __name__ == "__main__":
 
     test_fortran_frontend_parent()
+    test_fortran_frontend_module()