[flang] Prevent a bad folding rewrite

[klausler]

Summary

When a subexpression does not have both constant elements and
a constant shape, folding is rewriting it into a vector of its
elements. This is of course wrong when the shape shows that
the result has rank greater than 1.

flang/lib/Evaluate/fold-implementation.h

return common::visit(
       expr.u);
 }
 
 // FromArrayConstructor is a subroutine for MapOperation() below.
 // Given a flat ArrayConstructor<T> and a shape, it wraps the array
 // into an Expr<T>, folds it, and returns the resulting wrapped
 // array constructor or constant array value.
 template <typename T>
-Expr<T> FromArrayConstructor(FoldingContext &context,
-    ArrayConstructor<T> &&values, std::optional<ConstantSubscripts> &&shape) {
+std::optional<Expr<T>> FromArrayConstructor(
+    FoldingContext &context, ArrayConstructor<T> &&values, const Shape &shape) {
+  if (auto constShape{AsConstantExtents(context, shape)}) {
     Expr<T> result{Fold(context, Expr<T>{std::move(values)})};
-  if (shape) {
     if (auto *constant{UnwrapConstantValue<T>(result)}) {
-      return Expr<T>{constant->Reshape(std::move(*shape))};
+      // Elements and shape are both constant.
+      return Expr<T>{constant->Reshape(std::move(*constShape))};
+    }
+    if (constShape->size() == 1) {
+      if (auto elements{GetShape(context, result)}) {
+        if (auto constElements{AsConstantExtents(context, *elements)}) {
+          if (constElements->size() == 1 &&
+              constElements->at(0) == constShape->at(0)) {
+            // Elements are not constant, but array constructor has
+            // the right known shape and can be simply returned as is.
+            return std::move(result);
+          }
         }
       }
-  return result;
+    }
+  }
+  return std::nullopt;
 }
 
 // MapOperation is a utility for various specializations of ApplyElementwise()
 // that follow.  Given one or two flat ArrayConstructor<OPERAND> (wrapped in an
 // Expr<OPERAND>) for some specific operand type(s), apply a given function f
 // to each of their corresponding elements to produce a flat
 // ArrayConstructor<RESULT> (wrapped in an Expr<RESULT>).
 // Preserves shape.
 
 // Unary case
 template <typename RESULT, typename OPERAND>
-Expr<RESULT> MapOperation(FoldingContext &context,
+std::optional<Expr<RESULT>> MapOperation(FoldingContext &context,
     std::function<Expr<RESULT>(Expr<OPERAND> &&)> &&f, const Shape &shape,
     Expr<OPERAND> &&values) {
   ArrayConstructor<RESULT> result{values};
   if constexpr (common::HasMember<OPERAND, AllIntrinsicCategoryTypes>) {
     common::visit(
         [&](auto &&kindExpr) {
           using kindType = ResultType<decltype(kindExpr)>;
           auto &aConst{std::get<ArrayConstructor<kindType>>(kindExpr.u)};
           for (auto &acValue : aConst) {
             auto &scalar{std::get<Expr<kindType>>(acValue.u)};
             result.Push(Fold(context, f(Expr<OPERAND>{std::move(scalar)})));
           }
         },
         std::move(values.u));
   } else {
     auto &aConst{std::get<ArrayConstructor<OPERAND>>(values.u)};
     for (auto &acValue : aConst) {
       auto &scalar{std::get<Expr<OPERAND>>(acValue.u)};
       result.Push(Fold(context, f(std::move(scalar))));
     }
   }
-  return FromArrayConstructor(
-      context, std::move(result), AsConstantExtents(context, shape));
+  return FromArrayConstructor(context, std::move(result), shape);
 }
 
 template <typename RESULT, typename A>
 ArrayConstructor<RESULT> ArrayConstructorFromMold(
     const A &prototype, std::optional<Expr<SubscriptInteger>> &&length) {
   if constexpr (RESULT::category == TypeCategory::Character) {
     return ArrayConstructor<RESULT>{
         std::move(length.value()), ArrayConstructorValues<RESULT>{}};
   } else {
     return ArrayConstructor<RESULT>{prototype};
   }
 }
 
 // array * array case
 template <typename RESULT, typename LEFT, typename RIGHT>
-Expr<RESULT> MapOperation(FoldingContext &context,
+auto MapOperation(FoldingContext &context,
     std::function<Expr<RESULT>(Expr<LEFT> &&, Expr<RIGHT> &&)> &&f,
     const Shape &shape, std::optional<Expr<SubscriptInteger>> &&length,
-    Expr<LEFT> &&leftValues, Expr<RIGHT> &&rightValues) {
+    Expr<LEFT> &&leftValues, Expr<RIGHT> &&rightValues)
+    -> std::optional<Expr<RESULT>> {
   auto result{ArrayConstructorFromMold<RESULT>(leftValues, std::move(length))};
   auto &leftArrConst{std::get<ArrayConstructor<LEFT>>(leftValues.u)};
   if constexpr (common::HasMember<RIGHT, AllIntrinsicCategoryTypes>) {
     common::visit(
         [&](auto &&kindExpr) {
           using kindType = ResultType<decltype(kindExpr)>;
 
           auto &rightArrConst{std::get<ArrayConstructor<kindType>>(kindExpr.u)};
for (auto &leftValue : leftArrConst) {
       CHECK(rightIter != rightArrConst.end());
       auto &leftScalar{std::get<Expr<LEFT>>(leftValue.u)};
       auto &rightScalar{std::get<Expr<RIGHT>>(rightIter->u)};
       result.Push(
           Fold(context, f(std::move(leftScalar), std::move(rightScalar))));
       ++rightIter;
     }
   }
-  return FromArrayConstructor(
-      context, std::move(result), AsConstantExtents(context, shape));
+  return FromArrayConstructor(context, std::move(result), shape);
 }
 
 // array * scalar case
 template <typename RESULT, typename LEFT, typename RIGHT>
-Expr<RESULT> MapOperation(FoldingContext &context,
+auto MapOperation(FoldingContext &context,
     std::function<Expr<RESULT>(Expr<LEFT> &&, Expr<RIGHT> &&)> &&f,
     const Shape &shape, std::optional<Expr<SubscriptInteger>> &&length,
-    Expr<LEFT> &&leftValues, const Expr<RIGHT> &rightScalar) {
+    Expr<LEFT> &&leftValues, const Expr<RIGHT> &rightScalar)
+    -> std::optional<Expr<RESULT>> {
   auto result{ArrayConstructorFromMold<RESULT>(leftValues, std::move(length))};
   auto &leftArrConst{std::get<ArrayConstructor<LEFT>>(leftValues.u)};
   for (auto &leftValue : leftArrConst) {
     auto &leftScalar{std::get<Expr<LEFT>>(leftValue.u)};
     result.Push(
         Fold(context, f(std::move(leftScalar), Expr<RIGHT>{rightScalar})));
   }
-  return FromArrayConstructor(
-      context, std::move(result), AsConstantExtents(context, shape));
+  return FromArrayConstructor(context, std::move(result), shape);
 }
 
 // scalar * array case
 template <typename RESULT, typename LEFT, typename RIGHT>
-Expr<RESULT> MapOperation(FoldingContext &context,
+auto MapOperation(FoldingContext &context,
     std::function<Expr<RESULT>(Expr<LEFT> &&, Expr<RIGHT> &&)> &&f,
     const Shape &shape, std::optional<Expr<SubscriptInteger>> &&length,
-    const Expr<LEFT> &leftScalar, Expr<RIGHT> &&rightValues) {
+    const Expr<LEFT> &leftScalar, Expr<RIGHT> &&rightValues)
+    -> std::optional<Expr<RESULT>> {
   auto result{ArrayConstructorFromMold<RESULT>(leftScalar, std::move(length))};
   if constexpr (common::HasMember<RIGHT, AllIntrinsicCategoryTypes>) {
     common::visit(
         [&](auto &&kindExpr) {
           using kindType = ResultType<decltype(kindExpr)>;
           auto &rightArrConst{std::get<ArrayConstructor<kindType>>(kindExpr.u)};
           for (auto &rightValue : rightArrConst) {
             auto &rightScalar{std::get<Expr<kindType>>(rightValue.u)};
             result.Push(Fold(context,
                 f(Expr<LEFT>{leftScalar},
                     Expr<RIGHT>{std::move(rightScalar)})));
           }
         },
         std::move(rightValues.u));
   } else {
     auto &rightArrConst{std::get<ArrayConstructor<RIGHT>>(rightValues.u)};
     for (auto &rightValue : rightArrConst) {
       auto &rightScalar{std::get<Expr<RIGHT>>(rightValue.u)};
       result.Push(
           Fold(context, f(Expr<LEFT>{leftScalar}, std::move(rightScalar))));
     }
   }
-  return FromArrayConstructor(
-      context, std::move(result), AsConstantExtents(context, shape));
+  return FromArrayConstructor(context, std::move(result), shape);
 }
 
 template <typename DERIVED, typename RESULT, typename LEFT, typename RIGHT>
 std::optional<Expr<SubscriptInteger>> ComputeResultLength(
     Operation<DERIVED, RESULT, LEFT, RIGHT> &operation) {
   if constexpr (RESULT::category == TypeCategory::Character) {
     return Expr<RESULT>{operation.derived()}.LEN();
   }