From fe2069284d10f78c6472f5934983c5740dea8039 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval@gmail.com>
Date: Mon, 7 Nov 2022 18:23:22 +0100
Subject: [PATCH] [flang] Allow fir.class as fir.save_result operand #0

Reviewed By: jeanPerier, PeteSteinfeld

Differential Revision: https://reviews.llvm.org/D137546
---
 flang/include/flang/Optimizer/Dialect/FIRTypes.td |  2 +-
 flang/test/Lower/polymorphic.f90                  | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/flang/include/flang/Optimizer/Dialect/FIRTypes.td b/flang/include/flang/Optimizer/Dialect/FIRTypes.td
index 0d06e1d..eaf43a6 100644
--- a/flang/include/flang/Optimizer/Dialect/FIRTypes.td
+++ b/flang/include/flang/Optimizer/Dialect/FIRTypes.td
@@ -635,7 +635,7 @@ def AnyAddressableLike : TypeConstraint<Or<[fir_ReferenceType.predicate,
     FunctionType.predicate]>, "any addressable">;
 
 def ArrayOrBoxOrRecord : TypeConstraint<Or<[fir_SequenceType.predicate,
-    fir_BoxType.predicate, fir_RecordType.predicate]>,
+    IsBaseBoxTypePred, fir_RecordType.predicate]>,
     "fir.box, fir.array or fir.type">;
 
 
diff --git a/flang/test/Lower/polymorphic.f90 b/flang/test/Lower/polymorphic.f90
index 232dfad..d828c0c 100644
--- a/flang/test/Lower/polymorphic.f90
+++ b/flang/test/Lower/polymorphic.f90
@@ -58,4 +58,22 @@ module polymorphic_test
 ! CHECK-LABEL: test_allocate_unlimited_polymorphic_non_derived
 ! CHECK-NOT: _FortranAPointerNullifyDerived
 ! CHECK: fir.call @_FortranAPointerAllocate
+
+  function test_fct_ret_class()
+    class(p1), pointer :: test_fct_ret_class
+  end function
+
+  subroutine call_fct()
+    class(p1), pointer :: p
+    p => test_fct_ret_class()
+  end subroutine
+
+! CHECK-LABEL: func.func @_QMpolymorphic_testPtest_fct_ret_class() -> !fir.class<!fir.ptr<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>
+! CHECK: return %{{.*}} : !fir.class<!fir.ptr<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>
+
+! CHECK-lABEL: func.func @_QMpolymorphic_testPcall_fct()
+! CHECK: %[[RESULT:.*]] = fir.alloca !fir.class<!fir.ptr<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>> {bindc_name = ".result"}
+! CHECK: %[[CALL_RES:.*]] = fir.call @_QMpolymorphic_testPtest_fct_ret_class() : () -> !fir.class<!fir.ptr<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>
+! CHECK: fir.save_result %[[CALL_RES]] to %[[RESULT]] : !fir.class<!fir.ptr<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>, !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>>
+
 end module
-- 
2.7.4