[flang][hlfir] Fix multiple return declaration type

When the ENTRY statement is used, the same source can return different
types depending on the entry point. These different return values are
storage associated (share the same storage). Previously, this led to the
declaration of the results to all have the largest type. This patch adds
a convert between the stack allocation and the declaration so that the
hlfir.decl gets the right type.

I haven't managed to generate code where this convert converted a
reference to an allocation for a smaller type into an allocation for a
larger one, but I have added an assert just in case.

This is a different solution to https://reviews.llvm.org/D152725, see
discussion there.

Differential Revision: https://reviews.llvm.org/D152931

2 files changed, 117 insertions, 1 deletions

diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index ee21d63..f3efbfa 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -3993,8 +3993,38 @@ private:
         Fortran::lower::mapSymbolAttributes(*this, altResult, localSymbols,
                                             stmtCtx);
       } else {
+        // catch cases where the allocation for the function result storage type
+        // doesn't match the type of this symbol
+        mlir::Value preAlloc = primaryFuncResultStorage;
+        mlir::Type resTy = primaryFuncResultStorage.getType();
+        mlir::Type symTy = genType(altResult);
+        mlir::Type wrappedSymTy = fir::ReferenceType::get(symTy);
+        if (resTy != wrappedSymTy) {
+          // check size of the pointed to type so we can't overflow by writing
+          // double precision to a single precision allocation, etc
+          LLVM_ATTRIBUTE_UNUSED auto getBitWidth = [this](mlir::Type ty) {
+            // 15.6.2.6.3: differering result types should be integer, real,
+            // complex or logical
+            if (auto cmplx = mlir::dyn_cast_or_null<fir::ComplexType>(ty)) {
+              fir::KindTy kind = cmplx.getFKind();
+              return 2 * builder->getKindMap().getRealBitsize(kind);
+            }
+            if (auto logical = mlir::dyn_cast_or_null<fir::LogicalType>(ty)) {
+              fir::KindTy kind = logical.getFKind();
+              return builder->getKindMap().getLogicalBitsize(kind);
+            }
+            return ty.getIntOrFloatBitWidth();
+          };
+          assert(getBitWidth(fir::unwrapRefType(resTy)) >= getBitWidth(symTy));
+
+          // convert the storage to the symbol type so that the hlfir.declare
+          // gets the correct type for this symbol
+          preAlloc = builder->create<fir::ConvertOp>(getCurrentLocation(),
+                                                     wrappedSymTy, preAlloc);
+        }
+
         Fortran::lower::mapSymbolAttributes(*this, altResult, localSymbols,
-                                            stmtCtx, primaryFuncResultStorage);
+                                            stmtCtx, preAlloc);
       }
     }
 
diff --git a/flang/test/Lower/HLFIR/entry_return.f90 b/flang/test/Lower/HLFIR/entry_return.f90
new file mode 100644
index 0000000..d2fb80c
--- /dev/null
+++ b/flang/test/Lower/HLFIR/entry_return.f90
@@ -0,0 +1,86 @@
+! RUN: bbc -emit-hlfir -o - %s | FileCheck %s
+! test multiple return values with different types coming from ENTRY statements
+
+complex function f1()
+  logical e1
+  entry e1()
+  e1 = .false.
+end function
+! CHECK-LABEL:   func.func @_QPf1() -> !fir.complex<4> {
+! CHECK:           %[[VAL_0:.*]] = fir.alloca !fir.complex<4> {bindc_name = "f1", uniq_name = "_QFf1Ef1"}
+! CHECK:           %[[VAL_1:.*]]:2 = hlfir.declare %[[VAL_0]] {uniq_name = "_QFf1Ef1"} : (!fir.ref<!fir.complex<4>>) -> (!fir.ref<!fir.complex<4>>, !fir.ref<!fir.complex<4>>)
+! CHECK:           %[[VAL_2:.*]] = fir.convert %[[VAL_1]]#1 : (!fir.ref<!fir.complex<4>>) -> !fir.ref<!fir.logical<4>>
+! CHECK:           %[[VAL_3:.*]]:2 = hlfir.declare %[[VAL_2]] {uniq_name = "_QFf1Ee1"} : (!fir.ref<!fir.logical<4>>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>)
+! CHECK:           cf.br ^bb1
+! CHECK:         ^bb1:
+! CHECK:           %[[VAL_4:.*]] = arith.constant false
+! CHECK:           %[[VAL_5:.*]] = fir.convert %[[VAL_4]] : (i1) -> !fir.logical<4>
+! CHECK:           hlfir.assign %[[VAL_5]] to %[[VAL_3]]#0 : !fir.logical<4>, !fir.ref<!fir.logical<4>>
+! CHECK:           %[[VAL_6:.*]] = fir.load %[[VAL_1]]#1 : !fir.ref<!fir.complex<4>>
+! CHECK:           return %[[VAL_6]] : !fir.complex<4>
+! CHECK:         }
+
+! // CHECK-LABEL:   func.func @_QPe1() -> !fir.logical<4> {
+! CHECK:           %[[VAL_0:.*]] = fir.alloca !fir.complex<4> {bindc_name = "f1", uniq_name = "_QFf1Ef1"}
+! CHECK:           %[[VAL_1:.*]]:2 = hlfir.declare %[[VAL_0]] {uniq_name = "_QFf1Ef1"} : (!fir.ref<!fir.complex<4>>) -> (!fir.ref<!fir.complex<4>>, !fir.ref<!fir.complex<4>>)
+! CHECK:           %[[VAL_2:.*]] = fir.convert %[[VAL_1]]#1 : (!fir.ref<!fir.complex<4>>) -> !fir.ref<!fir.logical<4>>
+! CHECK:           %[[VAL_3:.*]]:2 = hlfir.declare %[[VAL_2]] {uniq_name = "_QFf1Ee1"} : (!fir.ref<!fir.logical<4>>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>)
+! CHECK:           cf.br ^bb1
+! CHECK:         ^bb1:
+! CHECK:           %[[VAL_4:.*]] = arith.constant false
+! CHECK:           %[[VAL_5:.*]] = fir.convert %[[VAL_4]] : (i1) -> !fir.logical<4>
+! CHECK:           hlfir.assign %[[VAL_5]] to %[[VAL_3]]#0 : !fir.logical<4>, !fir.ref<!fir.logical<4>>
+! CHECK:           %[[VAL_6:.*]] = fir.load %[[VAL_3]]#1 : !fir.ref<!fir.logical<4>>
+! CHECK:           return %[[VAL_6]] : !fir.logical<4>
+! CHECK:         }
+
+logical function f2()
+  complex e2
+  entry e2()
+  e2 = complex(1.0, 2.0)
+end function
+! CHECK-LABEL:   func.func @_QPf2() -> !fir.logical<4> {
+! CHECK:           %[[VAL_0:.*]] = fir.alloca !fir.complex<4> {bindc_name = "e2", uniq_name = "_QFf2Ee2"}
+! CHECK:           %[[VAL_1:.*]]:2 = hlfir.declare %[[VAL_0]] {uniq_name = "_QFf2Ee2"} : (!fir.ref<!fir.complex<4>>) -> (!fir.ref<!fir.complex<4>>, !fir.ref<!fir.complex<4>>)
+! CHECK:           %[[VAL_2:.*]] = fir.convert %[[VAL_1]]#1 : (!fir.ref<!fir.complex<4>>) -> !fir.ref<!fir.logical<4>>
+! CHECK:           %[[VAL_3:.*]]:2 = hlfir.declare %[[VAL_2]] {uniq_name = "_QFf2Ef2"} : (!fir.ref<!fir.logical<4>>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>)
+! CHECK:           cf.br ^bb1
+! CHECK:         ^bb1:
+! CHECK:           %[[VAL_4:.*]] = arith.constant 1.000000e+00 : f32
+! CHECK:           %[[VAL_5:.*]] = arith.constant 2.000000e+00 : f32
+! CHECK:           %[[VAL_6:.*]]:3 = hlfir.associate %[[VAL_4]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
+! CHECK:           %[[VAL_7:.*]]:3 = hlfir.associate %[[VAL_5]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
+! CHECK:           %[[VAL_8:.*]] = fir.call @_QPcomplex(%[[VAL_6]]#1, %[[VAL_7]]#1) fastmath<contract> : (!fir.ref<f32>, !fir.ref<f32>) -> f32
+! CHECK:           hlfir.end_associate %[[VAL_6]]#1, %[[VAL_6]]#2 : !fir.ref<f32>, i1
+! CHECK:           hlfir.end_associate %[[VAL_7]]#1, %[[VAL_7]]#2 : !fir.ref<f32>, i1
+! CHECK:           %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f32
+! CHECK:           %[[VAL_10:.*]] = fir.undefined !fir.complex<4>
+! CHECK:           %[[VAL_11:.*]] = fir.insert_value %[[VAL_10]], %[[VAL_8]], [0 : index] : (!fir.complex<4>, f32) -> !fir.complex<4>
+! CHECK:           %[[VAL_12:.*]] = fir.insert_value %[[VAL_11]], %[[VAL_9]], [1 : index] : (!fir.complex<4>, f32) -> !fir.complex<4>
+! CHECK:           hlfir.assign %[[VAL_12]] to %[[VAL_1]]#0 : !fir.complex<4>, !fir.ref<!fir.complex<4>>
+! CHECK:           %[[VAL_13:.*]] = fir.load %[[VAL_3]]#1 : !fir.ref<!fir.logical<4>>
+! CHECK:           return %[[VAL_13]] : !fir.logical<4>
+! CHECK:         }
+
+! CHECK-LABEL:   func.func @_QPe2() -> !fir.complex<4> {
+! CHECK:           %[[VAL_0:.*]] = fir.alloca !fir.complex<4> {bindc_name = "e2", uniq_name = "_QFf2Ee2"}
+! CHECK:           %[[VAL_1:.*]]:2 = hlfir.declare %[[VAL_0]] {uniq_name = "_QFf2Ee2"} : (!fir.ref<!fir.complex<4>>) -> (!fir.ref<!fir.complex<4>>, !fir.ref<!fir.complex<4>>)
+! CHECK:           %[[VAL_2:.*]] = fir.convert %[[VAL_1]]#1 : (!fir.ref<!fir.complex<4>>) -> !fir.ref<!fir.logical<4>>
+! CHECK:           %[[VAL_3:.*]]:2 = hlfir.declare %[[VAL_2]] {uniq_name = "_QFf2Ef2"} : (!fir.ref<!fir.logical<4>>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>)
+! CHECK:           cf.br ^bb1
+! CHECK:         ^bb1:
+! CHECK:           %[[VAL_4:.*]] = arith.constant 1.000000e+00 : f32
+! CHECK:           %[[VAL_5:.*]] = arith.constant 2.000000e+00 : f32
+! CHECK:           %[[VAL_6:.*]]:3 = hlfir.associate %[[VAL_4]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
+! CHECK:           %[[VAL_7:.*]]:3 = hlfir.associate %[[VAL_5]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
+! CHECK:           %[[VAL_8:.*]] = fir.call @_QPcomplex(%[[VAL_6]]#1, %[[VAL_7]]#1) fastmath<contract> : (!fir.ref<f32>, !fir.ref<f32>) -> f32
+! CHECK:           hlfir.end_associate %[[VAL_6]]#1, %[[VAL_6]]#2 : !fir.ref<f32>, i1
+! CHECK:           hlfir.end_associate %[[VAL_7]]#1, %[[VAL_7]]#2 : !fir.ref<f32>, i1
+! CHECK:           %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f32
+! CHECK:           %[[VAL_10:.*]] = fir.undefined !fir.complex<4>
+! CHECK:           %[[VAL_11:.*]] = fir.insert_value %[[VAL_10]], %[[VAL_8]], [0 : index] : (!fir.complex<4>, f32) -> !fir.complex<4>
+! CHECK:           %[[VAL_12:.*]] = fir.insert_value %[[VAL_11]], %[[VAL_9]], [1 : index] : (!fir.complex<4>, f32) -> !fir.complex<4>
+! CHECK:           hlfir.assign %[[VAL_12]] to %[[VAL_1]]#0 : !fir.complex<4>, !fir.ref<!fir.complex<4>>
+! CHECK:           %[[VAL_13:.*]] = fir.load %[[VAL_1]]#1 : !fir.ref<!fir.complex<4>>
+! CHECK:           return %[[VAL_13]] : !fir.complex<4>
+! CHECK:         }