diff --git a/loki/frontend/omni.py b/loki/frontend/omni.py index 3c64f4d86..ce5508212 100644 --- a/loki/frontend/omni.py +++ b/loki/frontend/omni.py @@ -579,6 +579,11 @@ def visit_varDecl(self, o, **kwargs): # why we restore the return_type _type = _type.dtype.return_type + # If the return type has a shape, we need to apply this as a dimension to the + # variable, otherwise it will be missing from the declaration + if _type.shape: + variable = variable.clone(dimensions=_type.shape) + else: raise ValueError diff --git a/loki/tests/test_subroutine.py b/loki/tests/test_subroutine.py index 9a6d543b4..8376f209a 100644 --- a/loki/tests/test_subroutine.py +++ b/loki/tests/test_subroutine.py @@ -2305,3 +2305,43 @@ def test_resolve_typebound_var_missing_definition(frontend, tmp_path): assert tt_invalid_val == 'tt%invalid%val' assert tt_invalid_val.type.dtype == BasicType.DEFERRED assert tt_invalid_val.parent.type.dtype == BasicType.DEFERRED + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('dim_decl', [':: add_to_a(n)', ', DIMENSION(n) :: add_to_a']) +def test_function_array_return_type(frontend, dim_decl): + """ + Verify array return types are correctly represented with all frontends + """ + fcode = f""" +subroutine member_functions + implicit none + integer :: i + real(kind=8) :: a(3) + contains + function add_to_a(b, n) + integer, intent(in) :: n + real(kind=8), intent(in) :: b(n) + real(kind=8) {dim_decl} + + do i = 1, n + add_to_a(i) = a(i) + b(i) + end do + end function +end subroutine member_functions + """.strip() + routine = Subroutine.from_source(fcode, frontend=frontend) + add_to_a = routine['add_to_a'] + return_type = add_to_a.procedure_type.return_type + assert return_type.dtype == BasicType.REAL + assert return_type.shape == ('n',) + ret_var = add_to_a.variable_map['add_to_a'] + assert ret_var.type.dtype == BasicType.REAL + assert ret_var.type.shape == ('n',) + assert ret_var.dimensions == ('n',) + + if frontend == OMNI: + # OMNI frontend puts the shape declaration always on the variable + assert ':: add_to_a(n)' in routine.to_fortran() + else: + assert dim_decl in routine.to_fortran()