From 22bae99029e0852fc5a1e07154d2d203139c0018 Mon Sep 17 00:00:00 2001
From: rupertford <rupert.ford@stfc.ac.uk>
Date: Thu, 7 Dec 2023 13:29:45 +0000
Subject: [PATCH] pr #2396. Added support for adding the target attribute to
 the original type so that the pointers can point to it. More tests are
 needed.

---
 src/psyclone/psyir/frontend/fparser2.py       | 39 +++++++++++++++++++
 .../frontend/fparser2_select_type_test.py     |  8 ++--
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/psyclone/psyir/frontend/fparser2.py b/src/psyclone/psyir/frontend/fparser2.py
index 9daf6b56f3..77ae71968b 100644
--- a/src/psyclone/psyir/frontend/fparser2.py
+++ b/src/psyclone/psyir/frontend/fparser2.py
@@ -3444,6 +3444,45 @@ def _type_construct_handler(self, node, parent):
             currentparent.addchild(elsebody)
             self.process_nodes(parent=elsebody, nodes=stmts[default_idx])
 
+        # Ensure that the type selector variable declaration has the
+        # pointer or a target attribute. The difficulty is that it
+        # will be in an UnknownFortranType
+        type_selector_name = selector
+        symbol_table = stmt.scope.symbol_table
+        symbol = symbol_table.lookup(type_selector_name)
+        datatype = symbol.datatype
+        if not isinstance(datatype, UnknownFortranType):
+            raise InternalError(
+                f"Expected the datatype of the select type selector to be an "
+                f"UnknownFortranType, but found '{type(datatype).__name__}'.")
+        # Re-create the datatype as an fparser ast
+        dummy_code = (
+            f"subroutine dummy()\n"
+            f"  {datatype.declaration}\n"
+            f"end subroutine\n")
+        parser = ParserFactory().create(std="f2008")
+        reader = FortranStringReader(dummy_code)
+        fp2_ast = parser(reader)
+        type_decl_stmt = fp2_ast.children[0].children[1].children[0]
+        attr_spec_list = type_decl_stmt.children[1]
+        found = False
+        if attr_spec_list:
+            for attr_spec in attr_spec_list.children:
+                attr_spec_str = attr_spec.string
+                if attr_spec_str.upper() in ["TARGET", "POINTER"]:
+                    found = True
+                    break
+        if not found:
+            # TARGET needs to be added as an attribute
+            if attr_spec_list:
+                attr_spec_list.append(Fortran2003.Attr_Spec("TARGET"))
+            else:
+                attr_spec_list = Fortran2003.Attr_Spec_List("TARGET")
+            type_decl_stmt.items = (
+                type_decl_stmt.items[0], attr_spec_list,
+                type_decl_stmt.items[2])
+            attr_spec_list.parent = type_decl_stmt
+            datatype._declaration = str(type_decl_stmt)
         return stmt
 
     def _case_construct_handler(self, node, parent):
diff --git a/src/psyclone/tests/psyir/frontend/fparser2_select_type_test.py b/src/psyclone/tests/psyir/frontend/fparser2_select_type_test.py
index 7b23075394..58333379d9 100644
--- a/src/psyclone/tests/psyir/frontend/fparser2_select_type_test.py
+++ b/src/psyclone/tests/psyir/frontend/fparser2_select_type_test.py
@@ -242,7 +242,8 @@ def test_class(fortran_reader, fortran_writer, tmpdir):
         "  END SELECT\n"
         "end subroutine\n"
         "end module\n")
-    expected = (
+    expected1 = "CLASS(*), TARGET :: type"
+    expected2 = (
         "    character(256) :: type_string\n\n"
         "    type(type2), pointer :: ptr_type2\n\n"
         "    INTEGER, pointer :: ptr_INTEGER\n\n"
@@ -284,9 +285,8 @@ def test_class(fortran_reader, fortran_writer, tmpdir):
         "    end if\n")
     psyir = fortran_reader.psyir_from_source(code)
     result = fortran_writer(psyir)
-    print(result)
-    exit(1)
-    assert expected in result
+    assert expected1 in result
+    assert expected2 in result
     if_blocks = psyir.walk(IfBlock)
     assert "was_class_type" in if_blocks[0].annotations
     assert "was_class_type" in if_blocks[2].annotations