From b7691d7be0362499eaadbec13cec5f2e248836f6 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval@gmail.com>
Date: Thu, 23 Feb 2023 17:28:07 +0100
Subject: [PATCH] [flang] Accept polymorphic component element in storage_size

When an element is extracted from a polymorphic array, it is
represented as a PolymorphicValue. The PolymorphicValue is
not a boxed value but holds the original polyrmophic array
and the element itself. This was raising an error in storage_size
lowering since we expect a boxed value to take advantage of
fir.box_elesize.
This patch handles PolymorphicValue correctly.

Reviewed By: PeteSteinfeld

Differential Revision: https://reviews.llvm.org/D144643
---
 flang/lib/Optimizer/Builder/IntrinsicCall.cpp |  5 ++++-
 flang/test/Lower/Intrinsics/storage_size.f90  | 29 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
index 007d258b3e14..926aea53e20f 100644
--- a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
@@ -4925,8 +4925,11 @@ IntrinsicLibrary::genStorageSize(mlir::Type resultType,
         builder.getKindMap().getIntegerBitsize(fir::toInt(constOp)));
   }
 
-  if (box.getType().isa<fir::ReferenceType>())
+  if (args[0].getBoxOf<fir::PolymorphicValue>()) {
+    box = builder.createBox(loc, args[0], /*isPolymorphic=*/true);
+  } else if (box.getType().isa<fir::ReferenceType>()) {
     box = builder.create<fir::LoadOp>(loc, box);
+  }
   mlir::Value eleSize = builder.create<fir::BoxEleSizeOp>(loc, kindTy, box);
   mlir::Value c8 = builder.createIntegerConstant(loc, kindTy, 8);
   return builder.create<mlir::arith::MulIOp>(loc, eleSize, c8);
diff --git a/flang/test/Lower/Intrinsics/storage_size.f90 b/flang/test/Lower/Intrinsics/storage_size.f90
index 2c975a194186..6cc883ebf4da 100644
--- a/flang/test/Lower/Intrinsics/storage_size.f90
+++ b/flang/test/Lower/Intrinsics/storage_size.f90
@@ -9,6 +9,10 @@ module storage_size_test
     integer :: b
   end type
 
+  type :: p3
+    class(p1), pointer :: p(:)
+  end type
+
 contains
 
   integer function unlimited_polymorphic_pointer(p) result(size)
@@ -113,4 +117,29 @@ contains
 ! CHECK: %[[RES:.*]] = fir.load %[[SIZE]] : !fir.ref<i64>
 ! CHECK: return %[[RES]] : i64
 
+  integer function polymorphic_value(t) result(size)
+    type(p3) :: t
+    size = storage_size(t%p(1))
+  end function
+
+! CHECK-LABEL: func.func @_QMstorage_size_testPpolymorphic_value(
+! CHECK-SAME: %[[T:.*]]: !fir.ref<!fir.type<_QMstorage_size_testTp3{p:!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>}>> {fir.bindc_name = "t"}) -> i32 {
+! CHECK: %[[ALLOCA:.*]] = fir.alloca i32 {bindc_name = "size", uniq_name = "_QMstorage_size_testFpolymorphic_valueEsize"}
+! CHECK: %[[FIELD_P:.*]] = fir.field_index p, !fir.type<_QMstorage_size_testTp3{p:!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>}>
+! CHECK: %[[COORD_P:.*]] = fir.coordinate_of %[[T]], %[[FIELD_P]] : (!fir.ref<!fir.type<_QMstorage_size_testTp3{p:!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>}>>, !fir.field) -> !fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>>
+! CHECK: %[[LOAD_COORD_P:.*]] = fir.load %[[COORD_P]] : !fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>>
+! CHECK: %[[C0:.*]] = arith.constant 0 : index
+! CHECK: %[[BOX_DIMS:.*]]:3 = fir.box_dims %[[LOAD_COORD_P]], %[[C0]] : (!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>, index) -> (index, index, index)
+! CHECK: %[[C1:.*]] = arith.constant 1 : i64
+! CHECK: %[[DIMI64:.*]] = fir.convert %[[BOX_DIMS]]#0 : (index) -> i64
+! CHECK: %[[IDX:.*]] = arith.subi %[[C1]], %[[DIMI64]] : i64
+! CHECK: %[[COORD_OF:.*]] = fir.coordinate_of %[[LOAD_COORD_P]], %[[IDX]] : (!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>, i64) -> !fir.ref<!fir.type<_QMstorage_size_testTp1{a:i32}>>
+! CHECK: %[[BOXED:.*]] = fir.embox %[[COORD_OF]] source_box %[[LOAD_COORD_P]] : (!fir.ref<!fir.type<_QMstorage_size_testTp1{a:i32}>>, !fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMstorage_size_testTp1{a:i32}>>>>) -> !fir.class<!fir.type<_QMstorage_size_testTp1{a:i32}>>
+! CHECK: %[[ELE_SIZE:.*]] = fir.box_elesize %[[BOXED]] : (!fir.class<!fir.type<_QMstorage_size_testTp1{a:i32}>>) -> i32
+! CHECK: %[[C8:.*]] = arith.constant 8 : i32
+! CHECK: %[[SIZE:.*]] = arith.muli %[[ELE_SIZE]], %[[C8]] : i32
+! CHECK: fir.store %[[SIZE]] to %[[ALLOCA]] : !fir.ref<i32>
+! CHECK: %[[RET:.*]] = fir.load %[[ALLOCA]] : !fir.ref<i32>
+! CHECK: return %[[RET]] : i32
+
 end module
-- 
2.11.4.GIT