[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.

379

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

376

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–382

You can also remove the brace here to be consistent

384

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),
+      box.getAddr(), 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) {
+    } else if ((moldExpr && moldExpr->Rank() != 0) &&
+               alloc.type.category() !=
+                   Fortran::semantics::DeclTypeSpec::Category::Character)
+      // When the MOLD is a scalar and the allocation is not CHARACTER, just
+      // use simple allocation because the MOLD argument doesn't have any
+      // relevant information.
       genMoldAllocation(alloc, boxAddr);
-    } else {
+    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;
   }

flang/test/Lower/allocate-mold.f90

 ! RUN: bbc -emit-fir %s -o - | FileCheck %s
 
 ! Test lowering of ALLOCATE statement with a MOLD argument for scalars
 
 subroutine scalar_mold_allocation()
   integer, allocatable :: a
   allocate(a, mold=9)
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPscalar_mold_allocation() {
 ! CHECK: %[[A:.*]] = fir.alloca !fir.box<!fir.heap<i32>> {bindc_name = "a", uniq_name = "_QFscalar_mold_allocationEa"}
 ! 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
+! CHECK: %[[ZERO:.*]] = fir.zero_bits !fir.heap<i32>
+! CHECK: fir.store %[[ZERO]] to %[[HEAP_A]] : !fir.ref<!fir.heap<i32>>
+! CHECK: %[[ALLOCATED:.*]] = fir.allocmem i32 {fir.must_be_heap = true, uniq_name = "_QFscalar_mold_allocationEa.alloc"}
+! CHECK: fir.store %[[ALLOCATED]] to %[[HEAP_A]] : !fir.ref<!fir.heap<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: %[[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: %[[C10:.*]] = arith.constant 10 : i32
+! CHECK: %[[C10I:.*]] = fir.convert %[[C10]] : (i32) -> index
+! CHECK: %[[C0:.*]] = arith.constant 0 : index
+! CHECK: %[[COMPARE:.*]] = arith.cmpi sgt, %[[C10I]], %[[C0]] : index
+! CHECK: %[[CHOICE:.*]] = arith.select %[[COMPARE]], %[[C10I]], %[[C0]] : index
+! CHECK: %[[ALLOCATED:.*]] = fir.allocmem !fir.array<?xf32>, %[[CHOICE]] {fir.must_be_heap = true, uniq_name = "_QFarray_scalar_mold_allocationEa.alloc"}
+! CHECK: fir.store %[[ALLOCATED]] to %[[HEAP_A]] : !fir.ref<!fir.heap<!fir.array<?xf32>>>
+! CHECK: fir.store %[[CHOICE]] to %[[EXT0]] : !fir.ref<index>