[flang] Carry dynamic type when emboxing polymorphic pointer

In order to be passed as passed-object in the dynamic dispatch, the
polymorphic pointer entity are emboxed. In this process, the dynamic
type must be preserve and pass to fir.embox as the tdesc operand. This
patch introduce a new ExtendedValue that allow to carry over the
dynamic type when the value is unboxed.

Depends on D136820

Reviewed By: PeteSteinfeld

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

7 files changed, 104 insertions, 1 deletions

diff --git a/flang/include/flang/Optimizer/Builder/BoxValue.h b/flang/include/flang/Optimizer/Builder/BoxValue.h
index 0764e62..7fc7fee 100644
--- a/flang/include/flang/Optimizer/Builder/BoxValue.h
+++ b/flang/include/flang/Optimizer/Builder/BoxValue.h
@@ -32,6 +32,7 @@ class BoxValue;
 class CharBoxValue;
 class CharArrayBoxValue;
 class MutableBoxValue;
+class PolymorphicValue;
 class ProcBoxValue;
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const CharBoxValue &);
@@ -40,6 +41,7 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &, const CharArrayBoxValue &);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const ProcBoxValue &);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const MutableBoxValue &);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const BoxValue &);
+llvm::raw_ostream &operator<<(llvm::raw_ostream &, const PolymorphicValue &);
 
 //===----------------------------------------------------------------------===//
 //
@@ -96,6 +98,24 @@ protected:
   mlir::Value len;
 };
 
+/// Polymorphic value associated with a dynamic type descriptor.
+class PolymorphicValue : public AbstractBox {
+public:
+  PolymorphicValue(mlir::Value addr, mlir::Value tdesc)
+      : AbstractBox{addr}, tdesc{tdesc} {}
+
+  PolymorphicValue clone(mlir::Value newBase) const { return {newBase, tdesc}; }
+
+  mlir::Value getTdesc() const { return tdesc; }
+
+  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
+                                       const PolymorphicValue &);
+  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }
+
+protected:
+  mlir::Value tdesc;
+};
+
 /// Abstract base class.
 /// Expressions of type array have at minimum a shape. These expressions may
 /// have lbound attributes (dynamic values) that affect the interpretation of
@@ -456,7 +476,7 @@ class ExtendedValue : public details::matcher<ExtendedValue> {
 public:
   using VT =
       std::variant<UnboxedValue, CharBoxValue, ArrayBoxValue, CharArrayBoxValue,
-                   ProcBoxValue, BoxValue, MutableBoxValue>;
+                   ProcBoxValue, BoxValue, MutableBoxValue, PolymorphicValue>;
 
   ExtendedValue() : box{UnboxedValue{}} {}
   template <typename A, typename = std::enable_if_t<
@@ -492,6 +512,7 @@ public:
     return match([](const fir::UnboxedValue &box) -> unsigned { return 0; },
                  [](const fir::CharBoxValue &box) -> unsigned { return 0; },
                  [](const fir::ProcBoxValue &box) -> unsigned { return 0; },
+                 [](const fir::PolymorphicValue &box) -> unsigned { return 0; },
                  [](const auto &box) -> unsigned { return box.rank(); });
   }
 
diff --git a/flang/lib/Frontend/CompilerInvocation.cpp b/flang/lib/Frontend/CompilerInvocation.cpp
index 761300b..e79ca8d 100644
--- a/flang/lib/Frontend/CompilerInvocation.cpp
+++ b/flang/lib/Frontend/CompilerInvocation.cpp
@@ -869,4 +869,5 @@ void CompilerInvocation::setLoweringOptions() {
 
   // Lower TRANSPOSE as a runtime call under -O0.
   loweringOpts.setOptimizeTranspose(codegenOpts.OptimizationLevel > 0);
+  loweringOpts.setPolymorphicTypeImpl(true);
 }
diff --git a/flang/lib/Lower/ConvertExpr.cpp b/flang/lib/Lower/ConvertExpr.cpp
index 82e298f..acc8004 100644
--- a/flang/lib/Lower/ConvertExpr.cpp
+++ b/flang/lib/Lower/ConvertExpr.cpp
@@ -1106,6 +1106,9 @@ public:
             fir::emitFatalError(loc, "derived type components must not be "
                                      "represented by fir::BoxValue");
           },
+          [&](const fir::PolymorphicValue &) {
+            TODO(loc, "polymorphic component in derived type assignment");
+          },
           [&](const fir::MutableBoxValue &toBox) {
             if (toBox.isPointer()) {
               Fortran::lower::associateMutableBox(
diff --git a/flang/lib/Optimizer/Builder/BoxValue.cpp b/flang/lib/Optimizer/Builder/BoxValue.cpp
index 83d8ba9..9f58182 100644
--- a/flang/lib/Optimizer/Builder/BoxValue.cpp
+++ b/flang/lib/Optimizer/Builder/BoxValue.cpp
@@ -85,6 +85,12 @@ llvm::raw_ostream &fir::operator<<(llvm::raw_ostream &os,
 }
 
 llvm::raw_ostream &fir::operator<<(llvm::raw_ostream &os,
+                                   const fir::PolymorphicValue &p) {
+  return os << "polymorphicvalue: { addr: " << p.getAddr()
+            << ", tdesc: " << p.getTdesc() << " }";
+}
+
+llvm::raw_ostream &fir::operator<<(llvm::raw_ostream &os,
                                    const fir::ArrayBoxValue &box) {
   os << "boxarray { addr: " << box.getAddr();
   if (box.getLBounds().size()) {
diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index d17f88b..6ad37e2 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -527,6 +527,12 @@ mlir::Value fir::FirOpBuilder::createBox(mlir::Location loc,
         return create<fir::LoadOp>(
             loc, fir::factory::getMutableIRBox(*this, loc, x));
       },
+      [&](const fir::PolymorphicValue &p) -> mlir::Value {
+        mlir::Value empty;
+        mlir::ValueRange emptyRange;
+        return create<fir::EmboxOp>(loc, boxTy, itemAddr, empty, empty,
+                                    emptyRange, p.getTdesc());
+      },
       [&](const auto &) -> mlir::Value {
         mlir::Value empty;
         mlir::ValueRange emptyRange;
diff --git a/flang/lib/Optimizer/Builder/MutableBox.cpp b/flang/lib/Optimizer/Builder/MutableBox.cpp
index 00692d3a..3f36310 100644
--- a/flang/lib/Optimizer/Builder/MutableBox.cpp
+++ b/flang/lib/Optimizer/Builder/MutableBox.cpp
@@ -419,6 +419,13 @@ fir::factory::genMutableBoxRead(fir::FirOpBuilder &builder, mlir::Location loc,
   }
   if (rank)
     return fir::ArrayBoxValue{addr, extents, lbounds};
+  if (box.isPolymorphic()) {
+    auto loadedBox = builder.create<fir::LoadOp>(loc, box.getAddr());
+    mlir::Type tdescType =
+        fir::TypeDescType::get(mlir::NoneType::get(builder.getContext()));
+    auto tdesc = builder.create<fir::BoxTypeDescOp>(loc, tdescType, loadedBox);
+    return fir::PolymorphicValue(addr, tdesc);
+  }
   return addr;
 }
 
@@ -467,6 +474,10 @@ void fir::factory::associateMutableBox(fir::FirOpBuilder &builder,
                                        mlir::ValueRange lbounds) {
   MutablePropertyWriter writer(builder, loc, box);
   source.match(
+      [&](const fir::PolymorphicValue &p) {
+        writer.updateMutableBox(p.getAddr(), /*lbounds=*/llvm::None,
+                                /*extents=*/llvm::None, /*lengths=*/llvm::None);
+      },
       [&](const fir::UnboxedValue &addr) {
         writer.updateMutableBox(addr, /*lbounds=*/llvm::None,
                                 /*extents=*/llvm::None, /*lengths=*/llvm::None);
@@ -566,6 +577,10 @@ void fir::factory::associateMutableBoxWithRemap(
   };
   MutablePropertyWriter writer(builder, loc, box);
   source.match(
+      [&](const fir::PolymorphicValue &p) {
+        writer.updateMutableBox(cast(p.getAddr()), lbounds, extents,
+                                /*lengths=*/llvm::None);
+      },
       [&](const fir::UnboxedValue &addr) {
         writer.updateMutableBox(cast(addr), lbounds, extents,
                                 /*lengths=*/llvm::None);
diff --git a/flang/test/Lower/allocatable-polymorphic.f90 b/flang/test/Lower/allocatable-polymorphic.f90
index 87a78aa..0338993 100644
--- a/flang/test/Lower/allocatable-polymorphic.f90
+++ b/flang/test/Lower/allocatable-polymorphic.f90
@@ -40,17 +40,33 @@ contains
     class(p1), pointer :: p
     class(p1), pointer :: c1, c2
     class(p1), pointer, dimension(:) :: c3, c4
+    integer :: i
 
     print*, 'test allocation of polymorphic pointers'
 
     allocate(p)
+    call p%proc1()
 
     allocate(p1::c1)
     allocate(p2::c2)
 
+    call c1%proc1()
+    call c2%proc1()
+
+    call c1%proc2()
+    call c2%proc2()
+
     allocate(p1::c3(10))
     allocate(p2::c4(20))
 
+    do i = 1, 10
+      call c3(i)%proc2()
+    end do
+
+    do i = 1, 20
+      call c4(i)%proc2()
+    end do
+
   end subroutine
 
 ! CHECK-LABEL: func.func @_QMpolyPtest_pointer()
@@ -99,6 +115,28 @@ contains
 ! CHECK: %[[BOX_ADDR:.*]] = fir.box_addr %[[C2_LOAD]] : (!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>) -> !fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
 ! CHECK: fir.store %[[BOX_ADDR]] to %[[C2_ADDR]] : !fir.ref<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>
 
+! call c1%proc1()
+! CHECK: %[[C1_DESC_CAST:.*]] = fir.convert %[[C1_DESC]] : (!fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>) -> !fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>
+! CHECK: fir.dispatch "proc1"(%[[C1_DESC_CAST]] : !fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>)
+
+! call c2%proc1()
+! CHECK: %[[C2_DESC_CAST:.*]] = fir.convert %[[C2_DESC]] : (!fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>) -> !fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>
+! CHECK: fir.dispatch "proc1"(%[[C2_DESC_CAST]] : !fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>)
+
+! call c1%proc2()
+! CHECK: %[[C1_LOAD:.*]] = fir.load %[[C1_ADDR]] : !fir.ref<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>
+! CHECK: %[[C1_DESC_LOAD:.*]] = fir.load %[[C1_DESC]] : !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>
+! CHECK: %[[C1_TDESC:.*]] = fir.box_tdesc %[[C1_DESC_LOAD]] : (!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>) -> !fir.tdesc<none>
+! CHECK: %[[C1_BOXED:.*]] = fir.embox %[[C1_LOAD]] tdesc %[[C1_TDESC]] : (!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>, !fir.tdesc<none>) -> !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: fir.dispatch "proc2"(%[[C1_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) (%[[C1_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) {pass_arg_pos = 0 : i32}
+
+! call c2%proc2()
+! CHECK: %[[C2_LOAD:.*]] = fir.load %[[C2_ADDR]] : !fir.ref<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>
+! CHECK: %[[C2_DESC_LOAD:.*]] = fir.load %[[C2_DESC]] : !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>
+! CHECK: %[[C2_TDESC:.*]] = fir.box_tdesc %[[C2_DESC_LOAD]] : (!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>) -> !fir.tdesc<none>
+! CHECK: %[[C2_BOXED:.*]] = fir.embox %[[C2_LOAD]] tdesc %[[C2_TDESC]] : (!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>, !fir.tdesc<none>) -> !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: fir.dispatch "proc2"(%[[C2_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) (%[[C2_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) {pass_arg_pos = 0 : i32}
+
 ! CHECK: %[[TYPE_DESC_P1:.*]] = fir.address_of(@_QMpolyE.dt.p1) : !fir.ref<!fir.type<{{.*}}>>
 ! CHECK: %[[C3_CAST:.*]] = fir.convert %[[C3_DESC]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %[[TYPE_DESC_P1_CAST:.*]] = fir.convert %[[TYPE_DESC_P1]] : (!fir.ref<!fir.type<_QM__fortran_type_infoTderivedtype{binding:!fir.box<!fir.ptr<!fir.array<?x!fir.type<{{.*}}>>) -> !fir.ref<none>
@@ -121,6 +159,19 @@ contains
 ! CHECK: %[[C4_CAST:.*]] = fir.convert %[[C4_DESC]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %{{.*}} = fir.call @_FortranAPointerAllocate(%[[C4_CAST]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
 
+! CHECK-LABEL: fir.do_loop
+! CHECK: %[[C3_LOAD:.*]] = fir.load %[[C3_DESC]] : !fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>>
+! CHECK: %[[C3_COORD:.*]] = fir.coordinate_of %[[C3_LOAD]], %{{.*}} : (!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>, i64) -> !fir.ref<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: %[[C3_TDESC:.*]] = fir.box_tdesc %[[C3_LOAD]] : (!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>) -> !fir.tdesc<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: %[[C3_BOXED:.*]] = fir.embox %[[C3_COORD]] tdesc %[[C3_TDESC]] : (!fir.ref<!fir.type<_QMpolyTp1{a:i32,b:i32}>>, !fir.tdesc<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) -> !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: fir.dispatch "proc2"(%[[C3_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) (%[[C3_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) {pass_arg_pos = 0 : i32}
+
+! CHECK-LABEL: fir.do_loop
+! CHECK: %[[C4_LOAD:.*]] = fir.load %[[C4_DESC]] : !fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>>
+! CHECK: %[[C4_COORD:.*]] = fir.coordinate_of %[[C4_LOAD]], %{{.*}} : (!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>, i64) -> !fir.ref<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: %[[C4_TDESC:.*]] = fir.box_tdesc %[[C4_LOAD]] : (!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>) -> !fir.tdesc<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: %[[C4_BOXED:.*]] = fir.embox %[[C4_COORD]] tdesc %[[C4_TDESC]] : (!fir.ref<!fir.type<_QMpolyTp1{a:i32,b:i32}>>, !fir.tdesc<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) -> !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
+! CHECK: fir.dispatch "proc2"(%[[C4_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) (%[[C4_BOXED]] : !fir.class<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) {pass_arg_pos = 0 : i32}
 
 end module