From d054959786eeec6311218efd857c0cca658b6f77 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler@nvidia.com>
Date: Fri, 1 Apr 2022 16:01:04 -0700
Subject: [PATCH] [flang] Fix shape analysis of RESHAPE result

Shape analysis of RESHAPE(..., SHAPE=s) should of course return
the SHAPE= actual argument when it is constant; but when it is
not, its length is still known, and thus so is the rank of the
result of RESHAPE(), and shape analysis should at least return
a shape vector of the right length rather than a result that
makes the result appear to be a scalar, which can lead to some
bogus error messages.

Also, while here: rename a private GetShapeHelper::AsShape()
routine so that it can't be confused with the ones in the API
of shape.h.

Differential Revision: https://reviews.llvm.org/D123712
---
 flang/include/flang/Evaluate/shape.h |  2 +-
 flang/lib/Evaluate/intrinsics.cpp    |  2 +-
 flang/lib/Evaluate/shape.cpp         | 40 ++++++++++++++++++++----------------
 flang/test/Semantics/call03.f90      |  1 +
 4 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/flang/include/flang/Evaluate/shape.h b/flang/include/flang/Evaluate/shape.h
index 246f346..5cee3a3 100644
--- a/flang/include/flang/Evaluate/shape.h
+++ b/flang/include/flang/Evaluate/shape.h
@@ -161,7 +161,7 @@ public:
 private:
   static Result ScalarShape() { return Shape{}; }
   static Shape ConstantShape(const Constant<ExtentType> &);
-  Result AsShape(ExtentExpr &&) const;
+  Result AsShapeResult(ExtentExpr &&) const;
   static Shape CreateShape(int rank, NamedEntity &);
 
   template <typename T>
diff --git a/flang/lib/Evaluate/intrinsics.cpp b/flang/lib/Evaluate/intrinsics.cpp
index 41731ac..d751c74 100644
--- a/flang/lib/Evaluate/intrinsics.cpp
+++ b/flang/lib/Evaluate/intrinsics.cpp
@@ -1513,7 +1513,7 @@ std::optional<SpecificCall> IntrinsicInterface::Match(
           if (auto shape{GetShape(context, *arg)}) {
             if (auto constShape{AsConstantShape(context, *shape)}) {
               shapeArgSize = constShape->At(ConstantSubscripts{1}).ToInt64();
-              CHECK(shapeArgSize >= 0);
+              CHECK(*shapeArgSize >= 0);
               argOk = true;
             }
           }
diff --git a/flang/lib/Evaluate/shape.cpp b/flang/lib/Evaluate/shape.cpp
index 0eefd93..15e44b9 100644
--- a/flang/lib/Evaluate/shape.cpp
+++ b/flang/lib/Evaluate/shape.cpp
@@ -53,7 +53,7 @@ Shape GetShapeHelper::ConstantShape(const Constant<ExtentType> &arrayConstant) {
   return result;
 }
 
-auto GetShapeHelper::AsShape(ExtentExpr &&arrayExpr) const -> Result {
+auto GetShapeHelper::AsShapeResult(ExtentExpr &&arrayExpr) const -> Result {
   if (context_) {
     arrayExpr = Fold(*context_, std::move(arrayExpr));
   }
@@ -63,17 +63,17 @@ auto GetShapeHelper::AsShape(ExtentExpr &&arrayExpr) const -> Result {
   if (auto *constructor{UnwrapExpr<ArrayConstructor<ExtentType>>(arrayExpr)}) {
     Shape result;
     for (auto &value : *constructor) {
-      if (auto *expr{std::get_if<ExtentExpr>(&value.u)}) {
-        if (expr->Rank() == 0) {
-          result.emplace_back(std::move(*expr));
-          continue;
-        }
+      auto *expr{std::get_if<ExtentExpr>(&value.u)};
+      if (expr && expr->Rank() == 0) {
+        result.emplace_back(std::move(*expr));
+      } else {
+        return std::nullopt;
       }
-      return std::nullopt;
     }
     return result;
+  } else {
+    return std::nullopt;
   }
-  return std::nullopt;
 }
 
 Shape GetShapeHelper::CreateShape(int rank, NamedEntity &base) {
@@ -847,15 +847,6 @@ auto GetShapeHelper::operator()(const ProcedureRef &call) const -> Result {
           }
         }
       }
-    } else if (intrinsic->name == "reshape") {
-      if (call.arguments().size() >= 2 && call.arguments().at(1)) {
-        // SHAPE(RESHAPE(array,shape)) -> shape
-        if (const auto *shapeExpr{
-                call.arguments().at(1).value().UnwrapExpr()}) {
-          auto shape{std::get<Expr<SomeInteger>>(shapeExpr->u)};
-          return AsShape(ConvertToType<ExtentType>(std::move(shape)));
-        }
-      }
     } else if (intrinsic->name == "pack") {
       if (call.arguments().size() >= 3 && call.arguments().at(2)) {
         // SHAPE(PACK(,,VECTOR=v)) -> SHAPE(v)
@@ -891,6 +882,18 @@ auto GetShapeHelper::operator()(const ProcedureRef &call) const -> Result {
           }
         }
       }
+    } else if (intrinsic->name == "reshape") {
+      if (call.arguments().size() >= 2 && call.arguments().at(1)) {
+        // SHAPE(RESHAPE(array,shape)) -> shape
+        if (const auto *shapeExpr{
+                call.arguments().at(1).value().UnwrapExpr()}) {
+          auto shapeArg{std::get<Expr<SomeInteger>>(shapeExpr->u)};
+          if (auto result{AsShapeResult(
+                  ConvertToType<ExtentType>(std::move(shapeArg)))}) {
+            return result;
+          }
+        }
+      }
     } else if (intrinsic->name == "spread") {
       // SHAPE(SPREAD(ARRAY,DIM,NCOPIES)) = SHAPE(ARRAY) with NCOPIES inserted
       // at position DIM.
@@ -966,7 +969,8 @@ auto GetShapeHelper::operator()(const ProcedureRef &call) const -> Result {
       // TODO: shapes of other non-elemental intrinsic results
     }
   }
-  return std::nullopt;
+  // The rank is always known even if the extents are not.
+  return Shape(static_cast<std::size_t>(call.Rank()), MaybeExtentExpr{});
 }
 
 // Check conformance of the passed shapes.
diff --git a/flang/test/Semantics/call03.f90 b/flang/test/Semantics/call03.f90
index 24e7e40..db76033 100644
--- a/flang/test/Semantics/call03.f90
+++ b/flang/test/Semantics/call03.f90
@@ -228,6 +228,7 @@ module m01
     real :: a(*)
     !ERROR: Scalar actual argument may not be associated with assumed-shape dummy argument 'x='
     call assumedshape(scalar)
+    call assumedshape(reshape(matrix,shape=[size(matrix)])) ! ok
     !ERROR: Rank of dummy argument is 1, but actual argument has rank 2
     call assumedshape(matrix)
     !ERROR: Assumed-size array may not be associated with assumed-shape dummy argument 'x='
-- 
2.7.4