[flang] Support allocate array from scalar source in runtime

[peixin]

Summary

As Fortran 2018 9.7.1.2(7), the value of each element of allocate object
becomes the value of source when the allocate object is array and the
source is scalar.

[peixin]

The running tests are in https://github.com/llvm/llvm-project/issues/60090#issuecomment-1397092468.

[PeteSteinfeld]

Thanks for implementing this and, especially, for adding the comments.

Can you please add a test to .../llvm-project/flang/unittests/Runtime?

[peixin]

Sorry for the late reply due to the Chinese Lunar New Year.

Add the unittests as @PeteSteinfeld mentioned.

[PeteSteinfeld]

Thanks, @peixin! All builds and tests correctly and looks good.

[peixin]

Thanks @PeteSteinfeld .

Will wait for one more day before landing in case other reviewers have comments.

flang/lib/Lower/Allocatable.cpp

private:
 
   void genSourceAllocation(const Allocation &alloc,
                            const fir::MutableBoxValue &box) {
     // Generate a sequence of runtime calls.
     errorManager.genStatCheck(builder, loc);
     genAllocateObjectInit(box);
     if (alloc.hasCoarraySpec())
       TODO(loc, "coarray allocation");
-    if (alloc.getShapeSpecs().size() > 0 && sourceExv.rank() == 0)
-      TODO(loc, "allocate array object with scalar SOURCE specifier");
     // 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);
     genSetDeferredLengthParameters(alloc, box);

flang/runtime/allocatable.cpp

int RTNAME(AllocatableAllocateSource)(Descriptor &alloc,
     const char *sourceFile, int sourceLine) {
   if (alloc.Elements() == 0) {
     return StatOk;
   }
   int stat{RTNAME(AllocatableAllocate)(
       alloc, hasStat, errMsg, sourceFile, sourceLine)};
   if (stat == StatOk) {
     Terminator terminator{sourceFile, sourceLine};
-    // 9.7.1.2(7)
-    Assign(alloc, source, terminator, /*skipRealloc=*/true);
+    DoFromSourceAssign(alloc, source, terminator);
   }
   return stat;
 }
 
 int RTNAME(AllocatableDeallocate)(Descriptor &descriptor, bool hasStat,
     const Descriptor *errMsg, const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};
   if (!descriptor.IsAllocatable()) {

flang/runtime/assign.h

 //===-- runtime/assign.h-----------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-// Internal APIs for data assignment (both intrinsic assignment and TBP defined
-// generic ASSIGNMENT(=)).
-
 #ifndef FORTRAN_RUNTIME_ASSIGN_INTERNAL_H_
 #define FORTRAN_RUNTIME_ASSIGN_INTERNAL_H_
 
 namespace Fortran::runtime {
 class Descriptor;
 class Terminator;
 
-// Assigns one object to another via intrinsic assignment (F'2018 10.2.1.3) or
-// type-bound (only!) defined assignment (10.2.1.4), as appropriate.  Performs
-// finalization, scalar expansion, & allocatable (re)allocation as needed.
-// Does not perform intrinsic assignment implicit type conversion.  Both
-// descriptors must be initialized.  Recurses as needed to handle components.
-// Do not perform allocatable reallocation if \p skipRealloc is true, which is
-// used for allocate statement with source specifier.
-void Assign(
-    Descriptor &, const Descriptor &, Terminator &, bool skipRealloc = false);
+// Assign one object to another via allocate statement from source specifier.
+// Note that if allocate object and source expression have the same rank, the
+// value of the allocate object becomes the value provided; otherwise the value
+// of each element of allocate object becomes the value provided (9.7.1.2(7)).
+void DoFromSourceAssign(Descriptor &, const Descriptor &, Terminator &);
 
 } // namespace Fortran::runtime
 #endif // FORTRAN_RUNTIME_ASSIGN_INTERNAL_H_

flang/runtime/assign.cpp

static void DoElementalDefinedAssignment(const Descriptor &to,
   for (std::size_t j{0}; j < toElements;
        ++j, to.IncrementSubscripts(toAt), from.IncrementSubscripts(fromAt)) {
     toElementDesc.set_base_addr(to.Element<char>(toAt));
     fromElementDesc.set_base_addr(from.Element<char>(fromAt));
     DoScalarDefinedAssignment(toElementDesc, fromElementDesc, special);
   }
 }
 
-void Assign(Descriptor &to, const Descriptor &from, Terminator &terminator,
-    bool skipRealloc) {
+// Assigns one object to another via intrinsic assignment (F'2018 10.2.1.3) or
+// type-bound (only!) defined assignment (10.2.1.4), as appropriate.  Performs
+// finalization, scalar expansion, & allocatable (re)allocation as needed.
+// Does not perform intrinsic assignment implicit type conversion.  Both
+// descriptors must be initialized.  Recurses as needed to handle components.
+// Do not perform allocatable reallocation if \p skipRealloc is true, which is
+// used for allocate statement with source specifier.
+static void Assign(Descriptor &to, const Descriptor &from,
+    Terminator &terminator, bool skipRealloc = false) {
   DescriptorAddendum *toAddendum{to.Addendum()};
   const typeInfo::DerivedType *toDerived{
       toAddendum ? toAddendum->derivedType() : nullptr};
   const DescriptorAddendum *fromAddendum{from.Addendum()};
   const typeInfo::DerivedType *fromDerived{
       fromAddendum ? fromAddendum->derivedType() : nullptr};
   bool wasJustAllocated{false};
   if (to.IsAllocatable()) {
if (to.rank() == from.rank() && to.IsContiguous() && from.IsContiguous()) {
            to.IncrementSubscripts(toAt), from.IncrementSubscripts(fromAt)) {
         std::memmove(to.Element<char>(toAt), from.Element<const char>(fromAt),
             elementBytes);
       }
     }
   }
 }
 
+void DoFromSourceAssign(
+    Descriptor &alloc, const Descriptor &source, Terminator &terminator) {
+  if (alloc.rank() > 0 && source.rank() == 0) {
+    // The value of each element of allocate object becomes the value of source.
+    DescriptorAddendum *allocAddendum{alloc.Addendum()};
+    SubscriptValue allocAt[maxRank];
+    alloc.GetLowerBounds(allocAt);
+    if (allocAddendum) {
+      const typeInfo::DerivedType *allocDerived{allocAddendum->derivedType()};
+      for (std::size_t n{alloc.Elements()}; n-- > 0;
+           alloc.IncrementSubscripts(allocAt)) {
+        Descriptor allocElement{*Descriptor::Create(*allocDerived,
+            reinterpret_cast<void *>(alloc.Element<char>(allocAt)), 0)};
+        Assign(allocElement, source, terminator, /*skipRealloc=*/true);
+      }
+    } else { // intrinsic type
+      for (std::size_t n{alloc.Elements()}; n-- > 0;
+           alloc.IncrementSubscripts(allocAt)) {
+        std::memmove(alloc.Element<char>(allocAt), source.raw().base_addr,
+            alloc.ElementBytes());
+      }
+    }
+  } else {
+    Assign(alloc, source, terminator, /*skipRealloc=*/true);
+  }
+}
+
 extern "C" {
 void RTNAME(Assign)(Descriptor &to, const Descriptor &from,
     const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};
   Assign(to, from, terminator);
 }
 
 } // extern "C"
 } // namespace Fortran::runtime

flang/runtime/pointer.cpp

int RTNAME(PointerAllocateSource)(Descriptor &pointer, const Descriptor &source,
     int sourceLine) {
   if (pointer.Elements() == 0) {
     return StatOk;
   }
   int stat{RTNAME(PointerAllocate)(
       pointer, hasStat, errMsg, sourceFile, sourceLine)};
   if (stat == StatOk) {
     Terminator terminator{sourceFile, sourceLine};
-    // 9.7.1.2(7)
-    Assign(pointer, source, terminator, /*skipRealloc=*/true);
+    DoFromSourceAssign(pointer, source, terminator);
   }
   return stat;
 }
 
 int RTNAME(PointerDeallocate)(Descriptor &pointer, bool hasStat,
     const Descriptor *errMsg, const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};
   if (!pointer.IsPointer()) {