[flang] Allow "non-expandable" scalars in single-element contexts

[klausler]

Summary

When a scalar expression is not expandable -- i.e., it would have to be
evaluated once and saved in a temporary to avoid changing the semantics
of the program if it were to be evaluated more than once -- it affects
some aspects of folding and expression semantics. In cases where
scalar expansion would not cause multiple evaluations due to the shape
of the result having but a single element, however, these "non-expandable"
scalar expressions can be safely allowed.

flang/include/flang/Evaluate/characteristics.h

 // of Fortran 2018.
 
 #ifndef FORTRAN_EVALUATE_CHARACTERISTICS_H_
 #define FORTRAN_EVALUATE_CHARACTERISTICS_H_
 
 #include "common.h"
 #include "expression.h"
 #include "shape.h"
+#include "tools.h"
 #include "type.h"
 #include "flang/Common/Fortran-features.h"
 #include "flang/Common/Fortran.h"
 #include "flang/Common/enum-set.h"
 #include "flang/Common/idioms.h"
 #include "flang/Common/indirection.h"
 #include "flang/Parser/char-block.h"
 #include "flang/Semantics/symbol.h"

flang/include/flang/Evaluate/shape.h

 //===-- include/flang/Evaluate/shape.h --------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 // GetShape() analyzes an expression and determines its shape, if possible,
 // representing the result as a vector of scalar integer expressions.
 
 #ifndef FORTRAN_EVALUATE_SHAPE_H_
 #define FORTRAN_EVALUATE_SHAPE_H_
 
 #include "expression.h"
-#include "fold.h"
 #include "traverse.h"
 #include "variable.h"
 #include "flang/Common/indirection.h"
-#include "flang/Evaluate/tools.h"
 #include "flang/Evaluate/type.h"
 #include <optional>
 #include <variant>
 
 namespace Fortran::parser {
 class ContextualMessages;
 }
 
private:
   }
 
   template <typename T>
   MaybeExtentExpr GetArrayConstructorExtent(
       const ArrayConstructorValues<T> &values) const {
     ExtentExpr result{0};
     for (const auto &value : values) {
       if (MaybeExtentExpr n{GetArrayConstructorValueExtent(value)}) {
-        result = std::move(result) + std::move(*n);
-        if (context_) {
-          // Fold during expression creation to avoid creating an expression so
-          // large we can't evalute it without overflowing the stack.
-          result = Fold(*context_, std::move(result));
-        }
+        AccumulateExtent(result, std::move(*n));
       } else {
         return std::nullopt;
       }
     }
     return result;
   }
 
+  // Add an extent to another, with folding
+  void AccumulateExtent(ExtentExpr &, ExtentExpr &&) const;
+
   FoldingContext *context_{nullptr};
   bool useResultSymbolShape_{true};
 };
 
 template <typename A>
 std::optional<Shape> GetShape(FoldingContext &context, const A &x) {
   if (auto shape{GetShapeHelper{context}(x)}) {
     return Fold(context, std::move(shape));

flang/include/flang/Evaluate/tools.h

 #define FORTRAN_EVALUATE_TOOLS_H_
 
 #include "traverse.h"
 #include "flang/Common/idioms.h"
 #include "flang/Common/template.h"
 #include "flang/Common/unwrap.h"
 #include "flang/Evaluate/constant.h"
 #include "flang/Evaluate/expression.h"
+#include "flang/Evaluate/shape.h"
+#include "flang/Evaluate/type.h"
 #include "flang/Parser/message.h"
 #include "flang/Semantics/attr.h"
 #include "flang/Semantics/symbol.h"
 #include <array>
 #include <optional>
 #include <set>
 #include <type_traits>
 #include <utility>
public:
   }
   bool operator()(const CoarrayRef &) { return true; }
 
 private:
   bool admitPureCall_{false};
 };
 
 template <typename T>
-bool IsExpandableScalar(const Expr<T> &expr, bool admitPureCall = false) {
-  return !UnexpandabilityFindingVisitor{admitPureCall}(expr);
+bool IsExpandableScalar(const Expr<T> &expr, FoldingContext &context,
+    const Shape &shape, bool admitPureCall = false) {
+  if (UnexpandabilityFindingVisitor{admitPureCall}(expr)) {
+    auto extents{AsConstantExtents(context, shape)};
+    return extents && GetSize(*extents) == 1;
+  } else {
+    return true;
+  }
 }
 
 // Common handling for procedure pointer compatibility of left- and right-hand
 // sides.  Returns nullopt if they're compatible.  Otherwise, it returns a
 // message that needs to be augmented by the names of the left and right sides
 std::optional<parser::MessageFixedText> CheckProcCompatibility(bool isCall,
     const std::optional<characteristics::Procedure> &lhsProcedure,
     const characteristics::Procedure *rhsProcedure,

flang/lib/Evaluate/fold-implementation.h

if (std::optional<Shape> leftShape{GetShape(context, leftExpr)}) {
                     std::move(*right));
               } else {
                 return std::nullopt;
               }
               return MapOperation(context, std::move(f), *leftShape,
                   std::move(resultLength), std::move(*left), std::move(*right));
             }
           }
-        } else if (IsExpandableScalar(rightExpr)) {
+        } else if (IsExpandableScalar(rightExpr, context, *leftShape)) {
           return MapOperation(context, std::move(f), *leftShape,
               std::move(resultLength), std::move(*left), rightExpr);
         }
       }
     }
-  } else if (rightExpr.Rank() > 0 && IsExpandableScalar(leftExpr)) {
-    if (std::optional<Shape> shape{GetShape(context, rightExpr)}) {
+  } else if (rightExpr.Rank() > 0) {
+    if (std::optional<Shape> rightShape{GetShape(context, rightExpr)}) {
+      if (IsExpandableScalar(leftExpr, context, *rightShape)) {
         if (auto right{AsFlatArrayConstructor(rightExpr)}) {
-        return MapOperation(context, std::move(f), *shape,
+          return MapOperation(context, std::move(f), *rightShape,
               std::move(resultLength), leftExpr, std::move(*right));
         }
       }
     }
+  }
   return std::nullopt;
 }
 
 template <typename DERIVED, typename RESULT, typename LEFT, typename RIGHT>
 auto ApplyElementwise(
     FoldingContext &context, Operation<DERIVED, RESULT, LEFT, RIGHT> &operation)
     -> std::optional<Expr<RESULT>> {
   return ApplyElementwise(context, operation,

flang/lib/Evaluate/shape.cpp

if (call.Rank() == 0) {
     } else {
       // TODO: shapes of other non-elemental intrinsic results
     }
   }
   // The rank is always known even if the extents are not.
   return Shape(static_cast<std::size_t>(call.Rank()), MaybeExtentExpr{});
 }
 
+void GetShapeHelper::AccumulateExtent(
+    ExtentExpr &result, ExtentExpr &&n) const {
+  result = std::move(result) + std::move(n);
+  if (context_) {
+    // Fold during expression creation to avoid creating an expression so
+    // large we can't evalute it without overflowing the stack.
+    result = Fold(*context_, std::move(result));
+  }
+}
+
 // Check conformance of the passed shapes.
 std::optional<bool> CheckConformance(parser::ContextualMessages &messages,
     const Shape &left, const Shape &right, CheckConformanceFlags::Flags flags,
     const char *leftIs, const char *rightIs) {
   int n{GetRank(left)};
   if (n == 0 && (flags & CheckConformanceFlags::LeftScalarExpandable)) {
     return true;
   }

flang/lib/Semantics/expression.cpp

if (symbol) {
                     *symbol);
               } else {
                 auto checked{
                     CheckConformance(messages, *componentShape, *valueShape,
                         CheckConformanceFlags::RightIsExpandableDeferred,
                         "component", "value")};
                 if (checked && *checked && GetRank(*componentShape) > 0 &&
                     GetRank(*valueShape) == 0 &&
-                    !IsExpandableScalar(*converted, true /*admit PURE call*/)) {
+                    !IsExpandableScalar(*converted, GetFoldingContext(),
+                        *componentShape, true /*admit PURE call*/)) {
                   AttachDeclaration(
                       Say(expr.source,
                           "Scalar value cannot be expanded to shape of array component '%s'"_err_en_US,
                           symbol->name()),
                       *symbol);
                 }
                 if (checked.value_or(true)) {
                   result.Add(*symbol, std::move(*converted));