#3473 Add Loop.variable_reference property

Author: sergisiso

src/psyclone/psyir/nodes/loop.py

@@ -39,7 +39,7 @@
 
 ''' This module contains the Loop node implementation.'''
 
-from typing import Union
+from typing import Union, Optional
 
 from psyclone.core import VariablesAccessMap
 from psyclone.psyir.nodes.datanode import DataNode
@@ -274,22 +274,28 @@ def _check_completeness(self):
                 f" '{', '.join([str(child) for child in self.children])}'.")
 
     @property
-    def variable(self):
+    def variable_reference(self) -> Reference:
+        '''
+        :returns: the control variable reference for this loop.
+        '''
+        self._check_completeness()
+        return self._children[0]
+
+    @property
+    def variable(self) -> Optional[DataSymbol]:
         '''
         :returns: the control variable for this loop.
-        :rtype: :py:class:`psyclone.psyir.symbols.DataSymbol`
         '''
         if len(self._children) < 1:
             return None
         return self._children[0].symbol
 
     @variable.setter
-    def variable(self, var):
+    def variable(self, var: DataSymbol):
         '''
         Setter for the variable associated with this loop.
 
         :param var: the control variable reference.
-        :type var: :py:class:`psyclone.psyir.symbols.DataSymbol`
 
         '''
         self._check_variable(var)

src/psyclone/tests/psyir/nodes/loop_test.py

@@ -126,6 +126,9 @@ def test_loop_navigation_properties():
     loop.addchild(Literal("1", ScalarType.integer_single_type()))
 
     # If it's not fully complete, it still returns an error
+    with pytest.raises(InternalError) as err:
+        _ = loop.variable_reference
+    assert error_str in str(err.value)
     with pytest.raises(InternalError) as err:
         _ = loop.start_expr
     assert error_str in str(err.value)
@@ -143,6 +146,7 @@ def test_loop_navigation_properties():
     loop.addchild(Schedule(parent=loop))
     loop.loop_body.addchild(Return(parent=loop.loop_body))
 
+    assert loop.variable_reference.name == "i"
     assert loop.start_expr.value == "0"
     assert loop.stop_expr.value == "2"
     assert loop.step_expr.value == "1"
@@ -274,7 +278,7 @@ def test_loop_create(fortran_writer):
         Reference(DataSymbol("i", ScalarType.real_single_type())))
     loop = Loop.create(DataSymbol("i", ScalarType.integer_single_type()),
                        start, stop, step, [child_node])
-    ref = loop.children[0]
+    ref = loop.variable_reference
     schedule = loop.loop_body
     assert isinstance(schedule, Schedule)
     check_links(loop, [ref, start, stop, step, schedule])

src/psyclone/tests/psyir/nodes/node_test.py

@@ -1362,7 +1362,7 @@ def test_replace_with_error1():
     # The first child of a loop is the loop start value which should
     # be a DataNode.
     with pytest.raises(GenerationError) as info:
-        loop.children[0].replace_with(new_node)
+        loop.variable_reference.replace_with(new_node)
     assert ("Item 'Assignment' can't be child 0 of 'Loop'. The valid "
             "format is: 'Reference, DataNode, DataNode, DataNode, Schedule'"
             in str(info.value))

src/psyclone/tests/psyir/nodes/reference_test.py

@@ -232,7 +232,8 @@ def test_reference_next_accesses(fortran_reader):
     assert b.next_accesses()[0] == b
 
     # Check the next_access of the loop variable
-    assert loop.children[0].next_accesses() == [loop.loop_body.children[0].rhs]
+    assert (loop.variable_reference.next_accesses() ==
+            [loop.loop_body.children[0].rhs])
 
     # Check that a loop accessing a variable before
     # the reference doesn't result in a false positive.