[llvm] [flang][runtime] Check SOURCE= conformability on ALLOCATE (PR #144113)

[Peter Klausler via llvm-commits]

https://github.com/klausler updated https://github.com/llvm/llvm-project/pull/144113

>From 4dc981f4f7b272d75bc3a557c12bfca6ef7d89a4 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler at nvidia.com>
Date: Thu, 12 Jun 2025 16:55:28 -0700
Subject: [PATCH] [flang][runtime] Check SOURCE= conformability on ALLOCATE

The SOURCE= expression of an ALLOCATE statement, when present and
not scalar, must conform to the shape of the allocated objects.
Check this at runtime, and return a recoverable error, or crash,
when appropriate.

Fixes https://github.com/llvm/llvm-project/issues/143900.
---
 flang-rt/lib/runtime/allocatable.cpp | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/flang-rt/lib/runtime/allocatable.cpp b/flang-rt/lib/runtime/allocatable.cpp
index ef18da6ea0786..f724f0a20884b 100644
--- a/flang-rt/lib/runtime/allocatable.cpp
+++ b/flang-rt/lib/runtime/allocatable.cpp
@@ -165,6 +165,26 @@ int RTDEF(AllocatableAllocateSource)(Descriptor &alloc,
       alloc, /*asyncObject=*/nullptr, hasStat, errMsg, sourceFile, sourceLine)};
   if (stat == StatOk) {
     Terminator terminator{sourceFile, sourceLine};
+    if (alloc.rank() != source.rank() && source.rank() != 0) {
+      terminator.Crash("ALLOCATE object has rank %d while SOURCE= has rank %d",
+          alloc.rank(), source.rank());
+    }
+    if (int rank{source.rank()}; rank > 0) {
+      SubscriptValue allocExtent[maxRank], sourceExtent[maxRank];
+      alloc.GetShape(allocExtent);
+      source.GetShape(sourceExtent);
+      for (int j{0}; j < rank; ++j) {
+        if (allocExtent[j] != sourceExtent[j]) {
+          if (!hasStat) {
+            terminator.Crash("ALLOCATE object has extent %jd on dimension %d, "
+                             "but SOURCE= has extent %jd",
+                static_cast<std::intmax_t>(allocExtent[j]), j + 1,
+                static_cast<std::intmax_t>(sourceExtent[j]));
+          }
+          return StatInvalidExtent;
+        }
+      }
+    }
     DoFromSourceAssign(alloc, source, terminator);
   }
   return stat;