[flang] Take result length into account in ApplyElementwise folding

[jeanPerier]

Summary

ApplyElementwise on character operation was always creating a result
ArrayConstructor with the length of the left operand. This is not
correct for concatenation and SetLength operations.

Compute and thread the length to the spot creating the ArrayConstructor
so that the length is correct for those character operations.

[klausler]

LGTM, thanks!

flang/lib/Evaluate/fold-implementation.h

1048

braces, please

[jeanPerier]

Use braces.

flang/lib/Evaluate/fold-implementation.h

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));
 }
 
+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,
     std::function<Expr<RESULT>(Expr<LEFT> &&, Expr<RIGHT> &&)> &&f,
-    const Shape &shape, Expr<LEFT> &&leftValues, Expr<RIGHT> &&rightValues) {
-  ArrayConstructor<RESULT> result{leftValues};
+    const Shape &shape, std::optional<Expr<SubscriptInteger>> &&length,
+    Expr<LEFT> &&leftValues, Expr<RIGHT> &&rightValues) {
+  auto result{ArrayConstructorFromMold<RESULT>(leftValues, std::move(length))};
   auto &leftArrConst{std::get<ArrayConstructor<LEFT>>(leftValues.u)};
   if constexpr (common::HasMember<RIGHT, AllIntrinsicCategoryTypes>) {
     std::visit(
         [&](auto &&kindExpr) {
           using kindType = ResultType<decltype(kindExpr)>;
 
           auto &rightArrConst{std::get<ArrayConstructor<kindType>>(kindExpr.u)};
           auto rightIter{rightArrConst.begin()};
Expr<RESULT> MapOperation(FoldingContext &context,
   return FromArrayConstructor(
       context, std::move(result), AsConstantExtents(context, shape));
 }
 
 // array * scalar case
 template <typename RESULT, typename LEFT, typename RIGHT>
 Expr<RESULT> MapOperation(FoldingContext &context,
     std::function<Expr<RESULT>(Expr<LEFT> &&, Expr<RIGHT> &&)> &&f,
-    const Shape &shape, Expr<LEFT> &&leftValues,
-    const Expr<RIGHT> &rightScalar) {
-  ArrayConstructor<RESULT> result{leftValues};
+    const Shape &shape, std::optional<Expr<SubscriptInteger>> &&length,
+    Expr<LEFT> &&leftValues, const Expr<RIGHT> &rightScalar) {
+  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));
 }
 
 // scalar * array case
 template <typename RESULT, typename LEFT, typename RIGHT>
 Expr<RESULT> MapOperation(FoldingContext &context,
     std::function<Expr<RESULT>(Expr<LEFT> &&, Expr<RIGHT> &&)> &&f,
-    const Shape &shape, const Expr<LEFT> &leftScalar,
-    Expr<RIGHT> &&rightValues) {
-  ArrayConstructor<RESULT> result{leftScalar};
+    const Shape &shape, std::optional<Expr<SubscriptInteger>> &&length,
+    const Expr<LEFT> &leftScalar, Expr<RIGHT> &&rightValues) {
+  auto result{ArrayConstructorFromMold<RESULT>(leftScalar, std::move(length))};
   if constexpr (common::HasMember<RIGHT, AllIntrinsicCategoryTypes>) {
     std::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,
for (auto &rightValue : rightArrConst) {
       result.Push(
           Fold(context, f(Expr<LEFT>{leftScalar}, std::move(rightScalar))));
     }
   }
   return FromArrayConstructor(
       context, std::move(result), AsConstantExtents(context, 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();
+  }
+  return std::nullopt;
+}
+
 // ApplyElementwise() recursively folds the operand expression(s) of an
 // operation, then attempts to apply the operation to the (corresponding)
 // scalar element(s) of those operands.  Returns std::nullopt for scalars
 // or unlinearizable operands.
 template <typename DERIVED, typename RESULT, typename OPERAND>
 auto ApplyElementwise(FoldingContext &context,
     Operation<DERIVED, RESULT, OPERAND> &operation,
     std::function<Expr<RESULT>(Expr<OPERAND> &&)> &&f)
return ApplyElementwise(context, operation,
           }});
 }
 
 template <typename DERIVED, typename RESULT, typename LEFT, typename RIGHT>
 auto ApplyElementwise(FoldingContext &context,
     Operation<DERIVED, RESULT, LEFT, RIGHT> &operation,
     std::function<Expr<RESULT>(Expr<LEFT> &&, Expr<RIGHT> &&)> &&f)
     -> std::optional<Expr<RESULT>> {
+  auto resultLength{ComputeResultLength(operation)};
   auto &leftExpr{operation.left()};
   leftExpr = Fold(context, std::move(leftExpr));
   auto &rightExpr{operation.right()};
   rightExpr = Fold(context, std::move(rightExpr));
   if (leftExpr.Rank() > 0) {
     if (std::optional<Shape> leftShape{GetShape(context, leftExpr)}) {
       if (auto left{AsFlatArrayConstructor(leftExpr)}) {
         if (rightExpr.Rank() > 0) {
           if (std::optional<Shape> rightShape{GetShape(context, rightExpr)}) {
             if (auto right{AsFlatArrayConstructor(rightExpr)}) {
               if (CheckConformance(context.messages(), *leftShape, *rightShape,
                       CheckConformanceFlags::EitherScalarExpandable)
                       .value_or(false /*fail if not known now to conform*/)) {
                 return MapOperation(context, std::move(f), *leftShape,
-                    std::move(*left), std::move(*right));
+                    std::move(resultLength), std::move(*left),
+                    std::move(*right));
               } else {
                 return std::nullopt;
               }
               return MapOperation(context, std::move(f), *leftShape,
-                  std::move(*left), std::move(*right));
+                  std::move(resultLength), std::move(*left), std::move(*right));
             }
           }
         } else if (IsExpandableScalar(rightExpr)) {
-          return MapOperation(
-              context, std::move(f), *leftShape, std::move(*left), rightExpr);
+          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)}) {
       if (auto right{AsFlatArrayConstructor(rightExpr)}) {
-        return MapOperation(
-            context, std::move(f), *shape, leftExpr, std::move(*right));
+        return MapOperation(context, std::move(f), *shape,
+            std::move(resultLength), leftExpr, std::move(*right));
       }
     }
   }
   return std::nullopt;
 }
 
 template <typename DERIVED, typename RESULT, typename LEFT, typename RIGHT>
 auto ApplyElementwise(

flang/test/Evaluate/folding22.f90

+! RUN: %S/test_folding.sh %s %t %flang_fc1
+! REQUIRES: shell
+
+! Test character concatenation folding
+
+logical, parameter :: test_scalar_scalar =  ('ab' // 'cde').eq.('abcde')
+
+character(2), parameter :: scalar_array(2) =  ['1','2'] // 'a'
+logical, parameter :: test_scalar_array = all(scalar_array.eq.(['1a', '2a']))
+
+character(2), parameter :: array_scalar(2) =  '1' // ['a', 'b']
+logical, parameter :: test_array_scalar = all(array_scalar.eq.(['1a', '1b']))
+
+character(2), parameter :: array_array(2) =  ['1','2'] // ['a', 'b']
+logical, parameter :: test_array_array = all(array_array.eq.(['1a', '2b']))
+
+
+character(1), parameter :: input(2) = ['x', 'y']
+character(*), parameter :: zero_sized(*) = input(2:1:1) // 'abcde'
+logical, parameter :: test_zero_sized = len(zero_sized).eq.6
+
+end
+