[flang-commits] [flang] [flang] Fix spurious error due to bad expression shape calculation (PR #124323)

[Peter Klausler via flang-commits]

https://github.com/klausler created https://github.com/llvm/llvm-project/pull/124323

GetShape() needed to be called with a FoldingContext in order to properly construct an extent expression for the shape of an array constructor whose elements (nested in an implied DO loop) were not scalars.

Fixes https://github.com/llvm/llvm-project/issues/124191.

>From c22c34956167c911b39082c580e751d8ebf0c711 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler at nvidia.com>
Date: Fri, 24 Jan 2025 10:53:12 -0800
Subject: [PATCH] [flang] Fix spurious error due to bad expression shape
 calculation

GetShape() needed to be called with a FoldingContext in order to
properly construct an extent expression for the shape of an
array constructor whose elements (nested in an implied DO loop)
were not scalars.

Fixes https://github.com/llvm/llvm-project/issues/124191.
---
 flang/include/flang/Evaluate/shape.h | 25 +++++++++++++++++--------
 flang/lib/Evaluate/shape.cpp         | 13 ++++++++++---
 flang/test/Evaluate/bug124191.f90    |  6 ++++++
 3 files changed, 33 insertions(+), 11 deletions(-)
 create mode 100644 flang/test/Evaluate/bug124191.f90

diff --git a/flang/include/flang/Evaluate/shape.h b/flang/include/flang/Evaluate/shape.h
index 3e42ec691158bc..f0505cfcdf2d7f 100644
--- a/flang/include/flang/Evaluate/shape.h
+++ b/flang/include/flang/Evaluate/shape.h
@@ -71,6 +71,9 @@ template <typename A>
 std::optional<Shape> GetShape(
     FoldingContext &, const A &, bool invariantOnly = true);
 template <typename A>
+std::optional<Shape> GetShape(
+    FoldingContext *, const A &, bool invariantOnly = true);
+template <typename A>
 std::optional<Shape> GetShape(const A &, bool invariantOnly = true);
 
 // The dimension argument to these inquiries is zero-based,
@@ -149,6 +152,8 @@ inline MaybeExtentExpr GetSize(const std::optional<Shape> &maybeShape) {
 // Utility predicate: does an expression reference any implied DO index?
 bool ContainsAnyImpliedDoIndex(const ExtentExpr &);
 
+// GetShape()
+
 class GetShapeHelper
     : public AnyTraverse<GetShapeHelper, std::optional<Shape>> {
 public:
@@ -261,23 +266,27 @@ class GetShapeHelper
 
 template <typename A>
 std::optional<Shape> GetShape(
-    FoldingContext &context, const A &x, bool invariantOnly) {
-  if (auto shape{GetShapeHelper{&context, invariantOnly}(x)}) {
-    return Fold(context, std::move(shape));
+    FoldingContext *context, const A &x, bool invariantOnly) {
+  if (auto shape{GetShapeHelper{context, invariantOnly}(x)}) {
+    if (context) {
+      return Fold(*context, std::move(shape));
+    } else {
+      return shape;
+    }
   } else {
     return std::nullopt;
   }
 }
 
 template <typename A>
-std::optional<Shape> GetShape(const A &x, bool invariantOnly) {
-  return GetShapeHelper{/*context=*/nullptr, invariantOnly}(x);
+std::optional<Shape> GetShape(
+    FoldingContext &context, const A &x, bool invariantOnly) {
+  return GetShape(&context, x, invariantOnly);
 }
 
 template <typename A>
-std::optional<Shape> GetShape(
-    FoldingContext *context, const A &x, bool invariantOnly = true) {
-  return GetShapeHelper{context, invariantOnly}(x);
+std::optional<Shape> GetShape(const A &x, bool invariantOnly) {
+  return GetShape(/*context=*/nullptr, x, invariantOnly);
 }
 
 template <typename A>
diff --git a/flang/lib/Evaluate/shape.cpp b/flang/lib/Evaluate/shape.cpp
index 58b824d9b8e644..fa957cfc084951 100644
--- a/flang/lib/Evaluate/shape.cpp
+++ b/flang/lib/Evaluate/shape.cpp
@@ -16,6 +16,7 @@
 #include "flang/Evaluate/tools.h"
 #include "flang/Evaluate/type.h"
 #include "flang/Parser/message.h"
+#include "flang/Semantics/semantics.h"
 #include "flang/Semantics/symbol.h"
 #include <functional>
 
@@ -23,6 +24,10 @@ using namespace std::placeholders; // _1, _2, &c. for std::bind()
 
 namespace Fortran::evaluate {
 
+FoldingContext &GetFoldingContextFrom(const Symbol &symbol) {
+  return symbol.owner().context().foldingContext();
+}
+
 bool IsImpliedShape(const Symbol &original) {
   const Symbol &symbol{ResolveAssociations(original)};
   const auto *details{symbol.detailsIf<semantics::ObjectEntityDetails>()};
@@ -483,7 +488,7 @@ static MaybeExtentExpr GetAssociatedExtent(
     const Symbol &symbol, int dimension) {
   if (const auto *assoc{symbol.detailsIf<semantics::AssocEntityDetails>()};
       assoc && !assoc->rank()) { // not SELECT RANK case
-    if (auto shape{GetShape(assoc->expr())};
+    if (auto shape{GetShape(GetFoldingContextFrom(symbol), assoc->expr())};
         shape && dimension < static_cast<int>(shape->size())) {
       if (auto &extent{shape->at(dimension)};
           // Don't return a non-constant extent, as the variables that
@@ -519,7 +524,8 @@ MaybeExtentExpr GetExtent(
   }
   if (const auto *details{symbol.detailsIf<semantics::ObjectEntityDetails>()}) {
     if (IsImpliedShape(symbol) && details->init()) {
-      if (auto shape{GetShape(symbol, invariantOnly)}) {
+      if (auto shape{
+              GetShape(GetFoldingContextFrom(symbol), symbol, invariantOnly)}) {
         if (dimension < static_cast<int>(shape->size())) {
           return std::move(shape->at(dimension));
         }
@@ -568,7 +574,8 @@ MaybeExtentExpr GetExtent(const Subscript &subscript, const NamedEntity &base,
                 MaybeExtentExpr{triplet.stride()});
           },
           [&](const IndirectSubscriptIntegerExpr &subs) -> MaybeExtentExpr {
-            if (auto shape{GetShape(subs.value())};
+            if (auto shape{GetShape(
+                    GetFoldingContextFrom(base.GetLastSymbol()), subs.value())};
                 shape && GetRank(*shape) == 1) {
               // vector-valued subscript
               return std::move(shape->at(0));
diff --git a/flang/test/Evaluate/bug124191.f90 b/flang/test/Evaluate/bug124191.f90
new file mode 100644
index 00000000000000..27d08032efa2f0
--- /dev/null
+++ b/flang/test/Evaluate/bug124191.f90
@@ -0,0 +1,6 @@
+! RUN: %flang_fc1 -fsyntax-only -pedantic %s 2>&1 | FileCheck --allow-empty %s
+! CHECK-NOT: error:
+! Regression test for https://github.com/llvm/llvm-project/issues/124191
+character(3) :: arr(5) = ['aa.', 'bb.', 'cc.', 'dd.', 'ee.']
+arr([(mod(iachar(arr(i:i-1:-1)(1:1)),5)+1, i=2,5,3)]) = arr(5:2:-1)
+end