[Flang] Fix ALLOCATE with MOLD where MOLD is a scalar

[PeteSteinfeld]

Summary

We were failing tests where an ALLOCATE statement that allocated an
array had a non-character scalar MOLD argument.

I fixed this be merging the code for ALLOCATE statements with MOLD
and SOURCE arguments.

[jeanPerier]

flang/lib/Lower/Allocatable.cpp

223

I do not think you should reverse this. Even if unused by this test, this is still broken code. You could probably still expose the issue with a mold allocation to a character scalar allocatable.

377

What you did is correct for simple numerical objects, but the character case will still be buggy I think (I did not check, but I suspect the bounds will not be set properly for allocate(some_char_array(10,20), mold=scalar_char)).
I think it is also a bit risky to completely ignore the mold here, there are some other cases where the MOLD may be relevant: at least polymorphism, parametrized derived types.

If what is missing are the bounds-spec application, it would make more sense to me to set the bounds properly in genMoldAllocation.

Pushing this a bit further, the code for SOURCE and MOLD allocation should be exactly similar except the last runtime call (genRuntimeAllocate for MOLD vs genRuntimeAllocateSource for SOURCE). Source allocation is a MOLD allocation followed by an assignment (RuntimeAllocateSource simply calls RuntimeAllocate and some assign version).

Could you try sharing genSourceAllocation above for both source and MOLD allocation except for last runtime call?

[clementval]

flang/lib/Lower/Allocatable.cpp

374–375

This will not work with polymorphic entities and unlimited polymorphic entities where we need to get the dynamic type information from the MOLD.

class(*), allocatable :: a

allocate(a, mold=3)

[PeteSteinfeld]

Changes after excellent feedback from Jean and Valentin.

[clementval]

LGTM

flang/lib/Lower/Allocatable.cpp

374

You can also remove the brace here to be consistent

376

brace

flang/lib/Lower/Allocatable.cpp

static void genRuntimeAllocateApplyMold(fir::FirOpBuilder &builder,
                                         fir::ExtendedValue mold, int rank) {
   mlir::func::FuncOp callee =
       box.isPointer()
           ? fir::runtime::getRuntimeFunc<mkRTKey(PointerApplyMold)>(loc,
                                                                     builder)
           : fir::runtime::getRuntimeFunc<mkRTKey(AllocatableApplyMold)>(
                 loc, builder);
   llvm::SmallVector<mlir::Value> args{
       fir::factory::getMutableIRBox(builder, loc, box), fir::getBase(mold),
       builder.createIntegerConstant(
           loc, callee.getFunctionType().getInputs()[2], rank)};
   llvm::SmallVector<mlir::Value> operands;
   for (auto [fst, snd] : llvm::zip(args, callee.getFunctionType().getInputs()))
     operands.emplace_back(builder.createConvert(loc, snd, fst));
   builder.create<fir::CallOp>(loc, callee, operands);
 }
 
for (const auto &allocOption :
           },
           allocOption.u);
   }
 
   void lowerAllocation(const Allocation &alloc) {
     fir::MutableBoxValue boxAddr =
         genMutableBoxValue(converter, loc, alloc.getAllocObj());
 
     if (sourceExpr) {
-      genSourceAllocation(alloc, boxAddr);
-    } else if (moldExpr) {
-      genMoldAllocation(alloc, boxAddr);
-    } else {
+      genSourceMoldAllocation(alloc, boxAddr, /*isSource=*/true);
+    } else if (moldExpr)
+      genSourceMoldAllocation(alloc, boxAddr, /*isSource=*/false);
+    else
       genSimpleAllocation(alloc, boxAddr);
-    }
-  }
 
   static bool lowerBoundsAreOnes(const Allocation &alloc) {
     for (const Fortran::parser::AllocateShapeSpec &shapeSpec :
          alloc.getShapeSpecs())
       if (std::get<0>(shapeSpec.t))
         return false;
     return true;
   }
if (sourceExpr && sourceExpr->Rank() > 0 &&
         mlir::Value ub = builder.create<mlir::arith::SubIOp>(
             loc, builder.create<mlir::arith::AddIOp>(loc, extent, lb), one);
         mlir::Value dimIndex = builder.createIntegerConstant(loc, i32Ty, i);
         genRuntimeSetBounds(builder, loc, box, dimIndex, lb, ub);
       }
     }
   }
 
-  void genSourceAllocation(const Allocation &alloc,
-                           const fir::MutableBoxValue &box) {
+  void genSourceMoldAllocation(const Allocation &alloc,
+                               const fir::MutableBoxValue &box, bool isSource) {
+    fir::ExtendedValue exv = isSource ? sourceExv : moldExv;
+    ;
     // Generate a sequence of runtime calls.
     errorManager.genStatCheck(builder, loc);
     genAllocateObjectInit(box);
     if (alloc.hasCoarraySpec())
       TODO(loc, "coarray allocation");
     // Set length of the allocate object if it has. Otherwise, get the length
     // from source for the deferred length parameter.
     if (lenParams.empty() && box.isCharacter() &&
         !box.hasNonDeferredLenParams())
-      lenParams.push_back(fir::factory::readCharLen(builder, loc, sourceExv));
-    if (alloc.type.IsPolymorphic())
-      genRuntimeAllocateApplyMold(builder, loc, box, sourceExv,
+      lenParams.push_back(fir::factory::readCharLen(builder, loc, exv));
+    if (!isSource || alloc.type.IsPolymorphic())
+      genRuntimeAllocateApplyMold(builder, loc, box, exv,
                                   alloc.getSymbol().Rank());
     genSetDeferredLengthParameters(alloc, box);
     genAllocateObjectBounds(alloc, box);
-    mlir::Value stat =
-        genRuntimeAllocateSource(builder, loc, box, sourceExv, errorManager);
-    fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
-    errorManager.assignStat(builder, loc, stat);
-  }
-
-  void genMoldAllocation(const Allocation &alloc,
-                         const fir::MutableBoxValue &box) {
-    genRuntimeAllocateApplyMold(builder, loc, box, moldExv,
-                                alloc.getSymbol().Rank());
-    errorManager.genStatCheck(builder, loc);
-    mlir::Value stat = genRuntimeAllocate(builder, loc, box, errorManager);
+    mlir::Value stat;
+    if (isSource)
+      stat = genRuntimeAllocateSource(builder, loc, box, exv, errorManager);
+    else
+      stat = genRuntimeAllocate(builder, loc, box, errorManager);
     fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
     errorManager.assignStat(builder, loc, stat);
   }
 
   /// Generate call to PointerNullifyDerived or AllocatableInitDerived
   /// to set the dynamic type information.
   void genInitDerived(const fir::MutableBoxValue &box, mlir::Value typeDescAddr,
                       int rank, int corank = 0) {

flang/test/Lower/allocate-mold.f90

 ! CHECK: %[[HEAP_A:.*]] = fir.alloca !fir.heap<i32> {uniq_name = "_QFscalar_mold_allocationEa.addr"}
 ! CHECK: %[[ADDR_A:.*]] = fir.load %[[HEAP_A]] : !fir.ref<!fir.heap<i32>>
 ! CHECK: %[[BOX_ADDR_A:.*]] = fir.embox %[[ADDR_A]] : (!fir.heap<i32>) -> !fir.box<!fir.heap<i32>>
 ! CHECK: fir.store %[[BOX_ADDR_A]] to %[[A]] : !fir.ref<!fir.box<!fir.heap<i32>>>
 ! CHECK: %[[A_REF_BOX_NONE1:.*]] = fir.convert %[[A]] : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %{{.*}} = fir.call @_FortranAAllocatableApplyMold(%[[A_REF_BOX_NONE1]], %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32) -> none
 ! CHECK: %[[A_REF_BOX_NONE2:.*]] = fir.convert %[[A]] : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %{{.*}} = fir.call @_FortranAAllocatableAllocate(%[[A_REF_BOX_NONE2]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
+
+subroutine array_scalar_mold_allocation()
+  real, allocatable :: a(:)
+
+  allocate (a(10), mold=3.0)
+end subroutine array_scalar_mold_allocation
+
+! CHECK-LABEL: func.func @_QParray_scalar_mold_allocation() {
+! CHECK: %[[A:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", uniq_name = "_QFarray_scalar_mold_allocationEa"}
+! CHECK: %[[HEAP_A:.*]] = fir.alloca !fir.heap<!fir.array<?xf32>> {uniq_name = "_QFarray_scalar_mold_allocationEa.addr"}
+! CHECK: %[[EXT0:.*]] = fir.alloca index {uniq_name = "_QFarray_scalar_mold_allocationEa.ext0"}
+! CHECK: %[[ZERO:.*]] = fir.zero_bits !fir.heap<!fir.array<?xf32>>
+! CHECK: fir.store %[[ZERO]] to %[[HEAP_A]] : !fir.ref<!fir.heap<!fir.array<?xf32>>>
+! CHECK: %[[LOADED_A:.*]] = fir.load %[[HEAP_A]] : !fir.ref<!fir.heap<!fir.array<?xf32>>>
+! CHECK: %[[SHAPESHIFT:.*]] = fir.shape_shift {{.*}}, {{.*}} : (index, index) -> !fir.shapeshift<1>
+! CHECK: %[[BOX_SHAPESHIFT:.*]] = fir.embox %[[LOADED_A]](%[[SHAPESHIFT]]) : (!fir.heap<!fir.array<?xf32>>, !fir.shapeshift<1>) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+! CHECK: fir.store %[[BOX_SHAPESHIFT]] to %[[A]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: %[[REF_BOX_A0:.*]] = fir.convert %1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableApplyMold(%[[REF_BOX_A0]], {{.*}}, {{.*}}) fastmath<contract> : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32) -> none
+! CHECK: %[[C10:.*]] = arith.constant 10 : i32
+! CHECK: %[[REF_BOX_A1:.*]] = fir.convert %1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableSetBounds(%[[REF_BOX_A1]], {{.*}},{{.*}}, {{.*}}) fastmath<contract> : (!fir.ref<!fir.box<none>>, i32, i64, i64) -> none
+! CHECK: %[[REF_BOX_A2:.*]] = fir.convert %[[A]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableAllocate(%[[REF_BOX_A2]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32