From fb7366cab1a4bb237b752f5cfacbdbef3464df9e Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin@nvidia.com>
Date: Thu, 27 Jul 2023 12:36:04 -0700
Subject: [PATCH] [flang] Fixed shape computation for elementals with optional
 dummys.

It looks like a regression after D151737: shape of the elemental
call became rank-0.

Reviewed By: klausler

Differential Revision: https://reviews.llvm.org/D156386
---
 flang/lib/Evaluate/shape.cpp                  | 17 ++++++++++++++---
 .../shape-of-elemental-with-optional-arg.f90  | 19 +++++++++++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)
 create mode 100644 flang/test/Lower/shape-of-elemental-with-optional-arg.f90

diff --git a/flang/lib/Evaluate/shape.cpp b/flang/lib/Evaluate/shape.cpp
index 9ddd7e585dae..652b59d901d6 100644
--- a/flang/lib/Evaluate/shape.cpp
+++ b/flang/lib/Evaluate/shape.cpp
@@ -811,13 +811,24 @@ auto GetShapeHelper::operator()(const ProcedureRef &call) const -> Result {
       if (auto chars{characteristics::Procedure::FromActuals(
               call.proc(), call.arguments(), *context_)}) {
         std::size_t j{0};
+        std::size_t anyArrayArgRank{0};
         for (const auto &arg : call.arguments()) {
-          if (arg && arg->Rank() > 0 && j < chars->dummyArguments.size() &&
-              !chars->dummyArguments[j].IsOptional()) {
-            return (*this)(*arg);
+          if (arg && arg->Rank() > 0 && j < chars->dummyArguments.size()) {
+            anyArrayArgRank = arg->Rank();
+            if (!chars->dummyArguments[j].IsOptional()) {
+              return (*this)(*arg);
+            }
           }
           ++j;
         }
+        if (anyArrayArgRank) {
+          // All dummy array arguments of the procedure are OPTIONAL.
+          // We cannot take the shape from just any array argument,
+          // because all of them might be OPTIONAL dummy arguments
+          // of the caller. Return unknown shape ranked according
+          // to the last actual array argument.
+          return Shape(anyArrayArgRank, MaybeExtentExpr{});
+        }
       }
     }
     return ScalarShape();
diff --git a/flang/test/Lower/shape-of-elemental-with-optional-arg.f90 b/flang/test/Lower/shape-of-elemental-with-optional-arg.f90
new file mode 100644
index 000000000000..311a7a70131c
--- /dev/null
+++ b/flang/test/Lower/shape-of-elemental-with-optional-arg.f90
@@ -0,0 +1,19 @@
+! Test that the shape of the elemental call is properly
+! computed as being rank 1, even though the only dummy
+! argument is optional.
+
+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+! RUN: bbc -emit-hlfir %s -o - | FileCheck %s
+
+subroutine test
+  interface
+     elemental function callee(arg1)
+       integer, intent(in), optional :: arg1
+       integer                       :: fun
+     end function callee
+  end interface
+  integer :: arr(2)
+  print *, callee(arr)
+end subroutine test
+! The PRINT statement must be lowered into a ranked print:
+! CHECK: fir.call @_FortranAioOutputDescriptor