[flang-commits] [flang] [flang] Fix folding of RANK(assumed-type assumed-rank) (PR #101027)

[Peter Klausler via flang-commits]

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

>From ce27da9d7642661108941b71947b0beaafc04ff1 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler at nvidia.com>
Date: Mon, 29 Jul 2024 08:13:38 -0700
Subject: [PATCH] [flang] Fix folding of RANK(assumed-type assumed-rank)

The code that deals with the special case of RANK(assumed-rank)
in intrinsic function folding wasn't handling the even more special
case of assumed-type assumed-rank dummy arguments.
---
 flang/lib/Evaluate/fold-integer.cpp           | 27 ++++++++++---------
 flang/lib/Evaluate/variable.cpp               |  3 ++-
 .../test/Evaluate/fold-assumed-type-rank.f90  |  6 +++++
 3 files changed, 23 insertions(+), 13 deletions(-)
 create mode 100644 flang/test/Evaluate/fold-assumed-type-rank.f90

diff --git a/flang/lib/Evaluate/fold-integer.cpp b/flang/lib/Evaluate/fold-integer.cpp
index 981cdff7f350b..e4195120ad9d6 100644
--- a/flang/lib/Evaluate/fold-integer.cpp
+++ b/flang/lib/Evaluate/fold-integer.cpp
@@ -1208,20 +1208,23 @@ Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
           cx->u)};
     }
   } else if (name == "rank") {
-    if (const auto *array{UnwrapExpr<Expr<SomeType>>(args[0])}) {
-      if (auto named{ExtractNamedEntity(*array)}) {
-        const Symbol &symbol{named->GetLastSymbol()};
-        if (IsAssumedRank(symbol)) {
-          // DescriptorInquiry can only be placed in expression of kind
-          // DescriptorInquiry::Result::kind.
-          return ConvertToType<T>(Expr<
-              Type<TypeCategory::Integer, DescriptorInquiry::Result::kind>>{
-              DescriptorInquiry{*named, DescriptorInquiry::Field::Rank}});
-        }
+    if (args[0]) {
+      const Symbol *symbol{nullptr};
+      if (auto dataRef{ExtractDataRef(args[0])}) {
+        symbol = &dataRef->GetLastSymbol();
+      } else {
+        symbol = args[0]->GetAssumedTypeDummy();
+      }
+      if (symbol && IsAssumedRank(*symbol)) {
+        // DescriptorInquiry can only be placed in expression of kind
+        // DescriptorInquiry::Result::kind.
+        return ConvertToType<T>(
+            Expr<Type<TypeCategory::Integer, DescriptorInquiry::Result::kind>>{
+                DescriptorInquiry{
+                    NamedEntity{*symbol}, DescriptorInquiry::Field::Rank}});
       }
-      return Expr<T>{args[0].value().Rank()};
+      return Expr<T>{args[0]->Rank()};
     }
-    return Expr<T>{args[0].value().Rank()};
   } else if (name == "selected_char_kind") {
     if (const auto *chCon{UnwrapExpr<Constant<TypeOf<std::string>>>(args[0])}) {
       if (std::optional<std::string> value{chCon->GetScalarValue()}) {
diff --git a/flang/lib/Evaluate/variable.cpp b/flang/lib/Evaluate/variable.cpp
index 247386a365de9..b074ae6d811ac 100644
--- a/flang/lib/Evaluate/variable.cpp
+++ b/flang/lib/Evaluate/variable.cpp
@@ -250,7 +250,8 @@ DescriptorInquiry::DescriptorInquiry(NamedEntity &&base, Field field, int dim)
   const Symbol &last{base_.GetLastSymbol()};
   CHECK(IsDescriptor(last));
   CHECK((field == Field::Len && dim == 0) ||
-      (field != Field::Len && dim >= 0 && dim < last.Rank()));
+      (field != Field::Len && dim >= 0 &&
+          (dim < last.Rank() || IsAssumedRank(last))));
 }
 
 // LEN()
diff --git a/flang/test/Evaluate/fold-assumed-type-rank.f90 b/flang/test/Evaluate/fold-assumed-type-rank.f90
new file mode 100644
index 0000000000000..ce296c8e27abf
--- /dev/null
+++ b/flang/test/Evaluate/fold-assumed-type-rank.f90
@@ -0,0 +1,6 @@
+! RUN: %flang_fc1 -fdebug-unparse %s 2>&1 | FileCheck %s
+subroutine sub3(ar_at)
+  type(*) :: ar_at(..)
+!CHECK:  PRINT *, int(int(rank(ar_at),kind=8),kind=4)
+  print *, rank(ar_at)
+end