[flang] Accept BaseBoxType in couple of fir.array_* operations

[clementval]

Summary

Couple of operation are expecting BoxType but can totally handle
ClassType as well. This patch updates couple of locations to support
BaseBoxType instead of BoxType only.

[jeanPerier]

LGTM

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Lower/ConvertExpr.cpp

llvm::SmallVector<mlir::Value> computeSliceShape(mlir::Value slice) {
     return slicedShape;
   }
 
   /// Get the shape from an ArrayOperand. The shape of the array is adjusted if
   /// the array was sliced.
   llvm::SmallVector<mlir::Value> getShape(ArrayOperand array) {
     if (array.slice)
       return computeSliceShape(array.slice);
-    if (array.memref.getType().isa<fir::BoxType>())
+    if (array.memref.getType().isa<fir::BaseBoxType>())
       return fir::factory::readExtents(builder, getLoc(),
                                        fir::BoxValue{array.memref});
     return fir::factory::getExtents(array.shape);
   }
 
   /// Get the shape from an ArrayLoad.
   llvm::SmallVector<mlir::Value> getShape(fir::ArrayLoadOp arrayLoad) {
     return getShape(ArrayOperand{arrayLoad.getMemref(), arrayLoad.getShape(),

flang/lib/Optimizer/CodeGen/CGOps.cpp

unsigned fir::cg::XReboxOp::getRank() {
   if (auto seqTy = fir::dyn_cast_ptrOrBoxEleTy(getBox().getType())
                        .dyn_cast<fir::SequenceType>())
     return seqTy.getDimension();
   return 0;
 }
 
 unsigned fir::cg::XArrayCoorOp::getRank() {
   auto memrefTy = getMemref().getType();
-  if (memrefTy.isa<fir::BoxType>())
+  if (memrefTy.isa<fir::BaseBoxType>())
     if (auto seqty =
             fir::dyn_cast_ptrOrBoxEleTy(memrefTy).dyn_cast<fir::SequenceType>())
       return seqty.getDimension();
   return getShape().size();
 }

flang/lib/Optimizer/CodeGen/CodeGen.cpp

doRewrite(fir::cg::XArrayCoorOp coor, mlir::Type ty, OpAdaptor adaptor,
     unsigned shiftOffset = coor.shiftOffset();
     unsigned sliceOffset = coor.sliceOffset();
     auto sliceOps = coor.getSlice().begin();
     mlir::Value one = genConstantIndex(loc, idxTy, rewriter, 1);
     mlir::Value prevExt = one;
     mlir::Value offset = genConstantIndex(loc, idxTy, rewriter, 0);
     const bool isShifted = !coor.getShift().empty();
     const bool isSliced = !coor.getSlice().empty();
-    const bool baseIsBoxed = coor.getMemref().getType().isa<fir::BoxType>();
+    const bool baseIsBoxed = coor.getMemref().getType().isa<fir::BaseBoxType>();
 
     // For each dimension of the array, generate the offset calculation.
     for (unsigned i = 0; i < rank; ++i, ++indexOffset, ++shapeOffset,
                   ++shiftOffset, sliceOffset += 3, sliceOps += 3) {
       mlir::Value index =
           integerCast(loc, rewriter, idxTy, operands[indexOffset]);
       mlir::Value lb =
           isShifted ? integerCast(loc, rewriter, idxTy, operands[shiftOffset])

flang/lib/Optimizer/Dialect/FIROps.cpp

if (auto shapeOp = getShape()) {
     unsigned shapeTyRank = 0;
     if (auto s = shapeTy.dyn_cast<fir::ShapeType>()) {
       shapeTyRank = s.getRank();
     } else if (auto ss = shapeTy.dyn_cast<fir::ShapeShiftType>()) {
       shapeTyRank = ss.getRank();
     } else {
       auto s = shapeTy.cast<fir::ShiftType>();
       shapeTyRank = s.getRank();
-      if (!getMemref().getType().isa<fir::BoxType>())
+      if (!getMemref().getType().isa<fir::BaseBoxType>())
         return emitOpError("shift can only be provided with fir.box memref");
     }
     if (arrDim && arrDim != shapeTyRank)
       return emitOpError("rank of dimension mismatched");
     if (shapeTyRank != getIndices().size())
       return emitOpError("number of indices do not match dim rank");
   }
 
if (auto shapeOp = getShape()) {
     unsigned shapeTyRank = 0u;
     if (auto s = shapeTy.dyn_cast<fir::ShapeType>()) {
       shapeTyRank = s.getRank();
     } else if (auto ss = shapeTy.dyn_cast<fir::ShapeShiftType>()) {
       shapeTyRank = ss.getRank();
     } else {
       auto s = shapeTy.cast<fir::ShiftType>();
       shapeTyRank = s.getRank();
-      if (!getMemref().getType().isa<fir::BoxType>())
+      if (!getMemref().getType().isa<fir::BaseBoxType>())
         return emitOpError("shift can only be provided with fir.box memref");
     }
     if (arrDim && arrDim != shapeTyRank)
       return emitOpError("rank of dimension mismatched");
   }
 
   if (auto sliceOp = getSlice()) {
     if (auto sl = mlir::dyn_cast_or_null<fir::SliceOp>(sliceOp.getDefiningOp()))

flang/test/Lower/polymorphic.f90

 ! 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}>>>>
 
+  subroutine implicit_loop_with_polymorphic()
+    class(p1), allocatable :: p(:)
+    allocate(p(3))
+    p%a = [ 1, 2, 3 ]
+  end subroutine
+
+! CHECK-LABEL: func.func @_QMpolymorphic_testPimplicit_loop_with_polymorphic() {
+! CHECK: %{{.*}} = fir.array_load %{{.*}}(%{{.*}}) [%{{.*}}] : (!fir.class<!fir.heap<!fir.array<?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>>, !fir.shift<1>, !fir.slice<1>) -> !fir.array<?xi32>
+! CHECK: %{{.*}} = fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} unordered iter_args(%{{.*}} = %{{.*}}) -> (!fir.array<?xi32>) {
+! CHECK: %{{.*}} = fir.array_fetch %{{.*}}, %{{.*}} : (!fir.array<3xi32>, index) -> i32
+! CHECK: %{{.*}} = fir.array_update %{{.*}}, %{{.*}}, %{{.*}} : (!fir.array<?xi32>, i32, index) -> !fir.array<?xi32>
+! CHECK: }
+! CHECK: fir.array_merge_store %{{.*}}, %{{.*}} to %{{.*}}[%{{.*}}] : !fir.array<?xi32>, !fir.array<?xi32>, !fir.class<!fir.heap<!fir.array<?x!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>>>, !fir.slice<1>
+
 end module