[flang] Support allocate array from scalar source in runtime

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.

Fix #60090.

Reviewed By: PeteSteinfeld

Differential Revision: https://reviews.llvm.org/D142112

diff --git a/flang/lib/Lower/Allocatable.cpp b/flang/lib/Lower/Allocatable.cpp
index f724d964d0f78e656bb4c77a5cbe7a2f176c01c6..9b0906191c7a900d9e3335ab5a707956dca003fe 100644 (file)
--- a/flang/lib/Lower/Allocatable.cpp
+++ b/flang/lib/Lower/Allocatable.cpp
@@ -559,8 +559,6 @@ private:
     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() &&

diff --git a/flang/runtime/allocatable.cpp b/flang/runtime/allocatable.cpp
index 3ec9bdaf63beb8cf305ab6099b9aa3876ddc8739..60e1073056c7dfea3aa96dbb6b5c53b115937108 100644 (file)
--- a/flang/runtime/allocatable.cpp
+++ b/flang/runtime/allocatable.cpp
@@ -127,8 +127,7 @@ int RTNAME(AllocatableAllocateSource)(Descriptor &alloc,
       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;
 }

diff --git a/flang/runtime/assign.cpp b/flang/runtime/assign.cpp
index 63ec732cd1cb7f2c3b09e1ee28b711353c006444..741049aa9207bd3af10d0e75278206391475e465 100644 (file)
--- a/flang/runtime/assign.cpp
+++ b/flang/runtime/assign.cpp
@@ -60,8 +60,15 @@ static void DoElementalDefinedAssignment(const Descriptor &to,
   }
 }
 
-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};
@@ -276,6 +283,34 @@ void Assign(Descriptor &to, const Descriptor &from, Terminator &terminator,
   }
 }
 
+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) {

diff --git a/flang/runtime/assign.h b/flang/runtime/assign.h
index 57dd9f38926ab062522cd2088edf082d80b1cedc..6b7f442bbfcedf84d6fe9b965d23d3cd273f2ae9 100644 (file)
--- a/flang/runtime/assign.h
+++ b/flang/runtime/assign.h
@@ -6,9 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-// 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_
 
@@ -16,15 +13,11 @@ 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_

diff --git a/flang/runtime/pointer.cpp b/flang/runtime/pointer.cpp
index 763348d0c365bf9a074075a2717cb217accd8901..115e49bdc68065d0dc5ccea47beaa799fa3ae252 100644 (file)
--- a/flang/runtime/pointer.cpp
+++ b/flang/runtime/pointer.cpp
@@ -142,8 +142,7 @@ int RTNAME(PointerAllocateSource)(Descriptor &pointer, const Descriptor &source,
       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;
 }

diff --git a/flang/unittests/Runtime/Allocatable.cpp b/flang/unittests/Runtime/Allocatable.cpp
index 11cb2f2db33e84e651ecfd0030c1cdb04f5155b2..b6bc759a0ecf5b95013af3b5d865c7e0fd9a0d09 100644 (file)
--- a/flang/unittests/Runtime/Allocatable.cpp
+++ b/flang/unittests/Runtime/Allocatable.cpp
@@ -71,3 +71,23 @@ TEST(AllocatableTest, MoveAlloc) {
     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();
+}

diff --git a/flang/unittests/Runtime/Pointer.cpp b/flang/unittests/Runtime/Pointer.cpp
index e00fb9bd5711fa3ca15b6d57418e13f3d42b7c0a..09ae3c4b4d966c59249ef09e25326bcb361520bf 100644 (file)
--- a/flang/unittests/Runtime/Pointer.cpp
+++ b/flang/unittests/Runtime/Pointer.cpp
@@ -63,3 +63,23 @@ TEST(Pointer, DeallocatePolymorphic) {
   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();
+}