[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()};
+    const typeInfo::DerivedType *allocDerived{
+        allocAddendum ? allocAddendum->derivedType() : nullptr};
+    SubscriptValue allocAt[maxRank];
+    alloc.GetLowerBounds(allocAt);
+    if (allocDerived) {
+      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()) {

flang/unittests/Runtime/Allocatable.cpp

TEST(AllocatableTest, MoveAlloc) {
   stat = RTNAME(MoveAlloc)(*a, *a, true, errMsg.get(), __FILE__, __LINE__);
   EXPECT_EQ(stat, 109);
   std::string_view errStr{errMsg->OffsetElement(), errMsg->ElementBytes()};
   auto trim_pos = errStr.find_last_not_of(' ');
   if (trim_pos != errStr.npos)
     errStr.remove_suffix(errStr.size() - trim_pos - 1);
   EXPECT_EQ(errStr, "MOVE_ALLOC passed the same address as to and from");
 }
+
+TEST(AllocatableTest, AllocateFromScalarSource) {
+  using Fortran::common::TypeCategory;
+  // REAL(4), ALLOCATABLE :: a(:)
+  auto a{createAllocatable(TypeCategory::Real, 4)};
+  // ALLOCATE(a(2:11), SOURCE=3.4)
+  float sourecStorage{3.4F};
+  auto s{Descriptor::Create(TypeCategory::Real, 4,
+      reinterpret_cast<void *>(&sourecStorage), 0, nullptr,
+      CFI_attribute_pointer)};
+  RTNAME(AllocatableSetBounds)(*a, 0, 2, 11);
+  RTNAME(AllocatableAllocateSource)
+  (*a, *s, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
+  EXPECT_TRUE(a->IsAllocated());
+  EXPECT_EQ(a->Elements(), 10u);
+  EXPECT_EQ(a->GetDimension(0).LowerBound(), 2);
+  EXPECT_EQ(a->GetDimension(0).UpperBound(), 11);
+  EXPECT_EQ(*a->OffsetElement<float>(), 3.4F);
+  a->Destroy();
+}

flang/unittests/Runtime/Pointer.cpp

TEST(Pointer, DeallocatePolymorphic) {
   auto p{Descriptor::Create(TypeCode{Fortran::common::TypeCategory::Integer, 4},
       4, nullptr, 0, nullptr, CFI_attribute_pointer)};
   RTNAME(PointerAllocate)
   (*p, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
   // DEALLOCATE(p)
   RTNAME(PointerDeallocatePolymorphic)
   (*p, nullptr, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
 }
+
+TEST(Pointer, AllocateFromScalarSource) {
+  // REAL(4), POINTER :: p(:)
+  auto p{Descriptor::Create(TypeCode{Fortran::common::TypeCategory::Real, 4}, 4,
+      nullptr, 1, nullptr, CFI_attribute_pointer)};
+  // ALLOCATE(p(2:11), SOURCE=3.4)
+  float sourecStorage{3.4F};
+  auto s{Descriptor::Create(Fortran::common::TypeCategory::Real, 4,
+      reinterpret_cast<void *>(&sourecStorage), 0, nullptr,
+      CFI_attribute_pointer)};
+  RTNAME(PointerSetBounds)(*p, 0, 2, 11);
+  RTNAME(PointerAllocateSource)
+  (*p, *s, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
+  EXPECT_TRUE(RTNAME(PointerIsAssociated)(*p));
+  EXPECT_EQ(p->Elements(), 10u);
+  EXPECT_EQ(p->GetDimension(0).LowerBound(), 2);
+  EXPECT_EQ(p->GetDimension(0).UpperBound(), 11);
+  EXPECT_EQ(*p->OffsetElement<float>(), 3.4F);
+  p->Destroy();
+}