[flang] Fold MERGE() of derived type values

[klausler]

Summary

Generalize FoldMerge() to accommodate derived type arguments and results,
rename it into Folder<T>::MERGE(), and remove it from the various
FoldIntrinsicFunction() routines for intrinsic types.

Fixes llvm-test-suite/Fortran/gfortran/regression/merge_init_expr_2.f90.

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Evaluate/fold-character.cpp

if (name == "achar" || name == "char") {
     return FoldMINorMAX(context, std::move(funcRef), Ordering::Greater);
   } else if (name == "maxval") {
     SingleCharType least{0};
     if (auto identity{Identity<T>(
             StringType{least}, GetConstantLength(context, funcRef, 0))}) {
       return FoldMaxvalMinval<T>(
           context, std::move(funcRef), RelationalOperator::GT, *identity);
     }
-  } else if (name == "merge") {
-    return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "min") {
     return FoldMINorMAX(context, std::move(funcRef), Ordering::Less);
   } else if (name == "minval") {
     // Collating sequences correspond to positive integers (3.31)
     SingleCharType most{0x7fffffff >> (8 * (4 - KIND))};
     if (auto identity{Identity<T>(
             StringType{most}, GetConstantLength(context, funcRef, 0))}) {
       return FoldMaxvalMinval<T>(

flang/lib/Evaluate/fold-complex.cpp

if (args.size() > 0 && args[0].has_value()) {
         return Fold(context,
             Expr<T>{
                 ComplexConstructor<KIND>{ToReal<KIND>(context, std::move(re)),
                     ToReal<KIND>(context, std::move(im))}});
       }
     }
   } else if (name == "dot_product") {
     return FoldDotProduct<T>(context, std::move(funcRef));
-  } else if (name == "merge") {
-    return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "product") {
     auto one{Scalar<Part>::FromInteger(value::Integer<8>{1}).value};
     return FoldProduct<T>(context, std::move(funcRef), Scalar<T>{one});
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   }
   // TODO: matmul
   return Expr<T>{std::move(funcRef)};

flang/lib/Evaluate/fold-implementation.h

std::optional<Constant<T>> GetConstantComponent(
       Component &, const std::vector<Constant<SubscriptInteger>> * = nullptr);
   std::optional<Constant<T>> Folding(ArrayRef &);
   std::optional<Constant<T>> Folding(DataRef &);
   Expr<T> Folding(Designator<T> &&);
   Constant<T> *Folding(std::optional<ActualArgument> &);
 
   Expr<T> CSHIFT(FunctionRef<T> &&);
   Expr<T> EOSHIFT(FunctionRef<T> &&);
+  Expr<T> MERGE(FunctionRef<T> &&);
   Expr<T> PACK(FunctionRef<T> &&);
   Expr<T> RESHAPE(FunctionRef<T> &&);
   Expr<T> SPREAD(FunctionRef<T> &&);
   Expr<T> TRANSPOSE(FunctionRef<T> &&);
   Expr<T> UNPACK(FunctionRef<T> &&);
 
   Expr<T> TRANSFER(FunctionRef<T> &&);
 
return common::visit(
       std::move(designator.u));
 }
 
 // Apply type conversion and re-folding if necessary.
 // This is where BOZ arguments are converted.
 template <typename T>
 Constant<T> *Folder<T>::Folding(std::optional<ActualArgument> &arg) {
   if (auto *expr{UnwrapExpr<Expr<SomeType>>(arg)}) {
+    if constexpr (T::category != TypeCategory::Derived) {
       if (!UnwrapExpr<Expr<T>>(*expr)) {
         if (auto converted{ConvertToType(T::GetType(), std::move(*expr))}) {
           *expr = Fold(context_, std::move(*converted));
         }
       }
+    }
     return UnwrapConstantValue<T>(*expr);
   }
   return nullptr;
 }
 
 template <typename... A, std::size_t... I>
 std::optional<std::tuple<const Constant<A> *...>> GetConstantArgumentsHelper(
     FoldingContext &context, ActualArguments &arguments,
     std::index_sequence<I...>) {
-  static_assert(
-      (... && IsSpecificIntrinsicType<A>)); // TODO derived types for MERGE?
   static_assert(sizeof...(A) > 0);
   std::tuple<const Constant<A> *...> args{
       Folder<A>{context}.Folding(arguments.at(I))...};
   if ((... && (std::get<I>(args)))) {
     return args;
   } else {
     return std::nullopt;
   }
for (unsigned int i{0}; i < sizeof...(TA); ++i) {
             context.messages().Say(
                 "Arguments in elemental intrinsic function are not conformable"_err_en_US);
             return Expr<TR>{std::move(funcRef)};
           }
         }
       }
     }
     CHECK(rank == GetRank(shape));
-
     // Compute all the scalar values of the results
     std::vector<Scalar<TR>> results;
     if (TotalElementCount(shape) > 0) {
       ConstantBounds bounds{shape};
       ConstantSubscripts resultIndex(rank, 1);
       ConstantSubscripts argIndex[]{std::get<I>(*args)->lbounds()...};
       do {
         if constexpr (std::is_same_v<WrapperType<TR, TA...>,
                           ScalarFuncWithContext<TR, TA...>>) {
           results.emplace_back(
               func(context, std::get<I>(*args)->At(argIndex[I])...));
         } else if constexpr (std::is_same_v<WrapperType<TR, TA...>,
                                  ScalarFunc<TR, TA...>>) {
           results.emplace_back(func(std::get<I>(*args)->At(argIndex[I])...));
         }
         (std::get<I>(*args)->IncrementSubscripts(argIndex[I]), ...);
       } while (bounds.IncrementSubscripts(resultIndex));
     }
     // Build and return constant result
     if constexpr (TR::category == TypeCategory::Character) {
       auto len{static_cast<ConstantSubscript>(
           results.empty() ? 0 : results[0].length())};
       return Expr<TR>{Constant<TR>{len, std::move(results), std::move(shape)}};
+    } else if constexpr (TR::category == TypeCategory::Derived) {
+      if (!results.empty()) {
+        return Expr<TR>{rank == 0
+                ? Constant<TR>{results.front()}
+                : Constant<TR>{results.front().derivedTypeSpec(),
+                      std::move(results), std::move(shape)}};
+      }
     } else {
       return Expr<TR>{Constant<TR>{std::move(results), std::move(shape)}};
     }
   }
   return Expr<TR>{std::move(funcRef)};
 }
 
 template <typename TR, typename... TA>
if (ok) {
       return Expr<T>{PackageConstant<T>(
           std::move(resultElements), *array, array->shape())};
     }
   }
   // Invalid, prevent re-folding
   return MakeInvalidIntrinsic(std::move(funcRef));
 }
 
+template <typename T> Expr<T> Folder<T>::MERGE(FunctionRef<T> &&funcRef) {
+  return FoldElementalIntrinsic<T, T, T, LogicalResult>(context_,
+      std::move(funcRef),
+      ScalarFunc<T, T, T, LogicalResult>(
+          [](const Scalar<T> &ifTrue, const Scalar<T> &ifFalse,
+              const Scalar<LogicalResult> &predicate) -> Scalar<T> {
+            return predicate.IsTrue() ? ifTrue : ifFalse;
+          }));
+}
+
 template <typename T> Expr<T> Folder<T>::PACK(FunctionRef<T> &&funcRef) {
   auto args{funcRef.arguments()};
   CHECK(args.size() == 3);
   const auto *array{UnwrapConstantValue<T>(args[0])};
   const auto *vector{UnwrapConstantValue<T>(args[2])};
   auto convertedMask{Fold(context_,
       ConvertToType<LogicalResult>(
           Expr<SomeLogical>{DEREF(UnwrapExpr<Expr<SomeLogical>>(args[1]))}))};
for (std::optional<ActualArgument> &arg : args) {
     }
   }
   if (auto *intrinsic{std::get_if<SpecificIntrinsic>(&funcRef.proc().u)}) {
     const std::string name{intrinsic->name};
     if (name == "cshift") {
       return Folder<T>{context}.CSHIFT(std::move(funcRef));
     } else if (name == "eoshift") {
       return Folder<T>{context}.EOSHIFT(std::move(funcRef));
+    } else if (name == "merge") {
+      return Folder<T>{context}.MERGE(std::move(funcRef));
     } else if (name == "pack") {
       return Folder<T>{context}.PACK(std::move(funcRef));
     } else if (name == "reshape") {
       return Folder<T>{context}.RESHAPE(std::move(funcRef));
     } else if (name == "spread") {
       return Folder<T>{context}.SPREAD(std::move(funcRef));
     } else if (name == "transfer") {
       return Folder<T>{context}.TRANSFER(std::move(funcRef));
     } else if (name == "transpose") {
       return Folder<T>{context}.TRANSPOSE(std::move(funcRef));
     } else if (name == "unpack") {
       return Folder<T>{context}.UNPACK(std::move(funcRef));
     }
     // TODO: extends_type_of, same_type_as
     if constexpr (!std::is_same_v<T, SomeDerived>) {
       return FoldIntrinsicFunction(context, std::move(funcRef));
     }
   }
   return Expr<T>{std::move(funcRef)};
 }
 
-template <typename T>
-Expr<T> FoldMerge(FoldingContext &context, FunctionRef<T> &&funcRef) {
-  return FoldElementalIntrinsic<T, T, T, LogicalResult>(context,
-      std::move(funcRef),
-      ScalarFunc<T, T, T, LogicalResult>(
-          [](const Scalar<T> &ifTrue, const Scalar<T> &ifFalse,
-              const Scalar<LogicalResult> &predicate) -> Scalar<T> {
-            return predicate.IsTrue() ? ifTrue : ifFalse;
-          }));
-}
-
 Expr<ImpliedDoIndex::Result> FoldOperation(FoldingContext &, ImpliedDoIndex &&);
 
 // Array constructor folding
 template <typename T> class ArrayConstructorFolder {
 public:
   explicit ArrayConstructorFolder(FoldingContext &c) : context_{c} {}
 
   Expr<T> FoldArray(ArrayConstructor<T> &&array) {

flang/lib/Evaluate/fold-integer.cpp

if (auto *sx{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
           },
           sx->u);
     }
   } else if (name == "maxloc") {
     return FoldLocation<WhichLocation::Maxloc, T>(context, std::move(funcRef));
   } else if (name == "maxval") {
     return FoldMaxvalMinval<T>(context, std::move(funcRef),
         RelationalOperator::GT, T::Scalar::Least());
-  } else if (name == "merge") {
-    return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "merge_bits") {
     return FoldElementalIntrinsic<T, T, T, T>(
         context, std::move(funcRef), &Scalar<T>::MERGE_BITS);
   } else if (name == "min") {
     return FoldMINorMAX(context, std::move(funcRef), Ordering::Less);
   } else if (name == "min0" || name == "min1") {
     return RewriteSpecificMINorMAX(context, std::move(funcRef));
   } else if (name == "minexponent") {

flang/lib/Evaluate/fold-logical.cpp

if (cx0 && cx1) {
                                             : RelationalOperator::LT,
                   ConvertToType<Ascii>(std::move(*cx0)),
                   ConvertToType<Ascii>(std::move(*cx1)))));
     }
   } else if (name == "logical") {
     if (auto *expr{UnwrapExpr<Expr<SomeLogical>>(args[0])}) {
       return Fold(context, ConvertToType<T>(std::move(*expr)));
     }
-  } else if (name == "merge") {
-    return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "parity") {
     return FoldAllAnyParity(
         context, std::move(funcRef), &Scalar<T>::NEQV, Scalar<T>{false});
   } else if (name == "same_type_as") {
     // Type equality testing with SAME_TYPE_AS() ignores any type parameters.
     // Returns a constant truth value when the result is known now.
     if (args[0] && args[1]) {
       auto t0{args[0]->GetType()};

flang/lib/Evaluate/fold-real.cpp

return FoldElementalIntrinsic<T, T, T>(context, std::move(funcRef),
               }
               return result.value;
             }));
   } else if (name == "max") {
     return FoldMINorMAX(context, std::move(funcRef), Ordering::Greater);
   } else if (name == "maxval") {
     return FoldMaxvalMinval<T>(context, std::move(funcRef),
         RelationalOperator::GT, T::Scalar::HUGE().Negate());
-  } else if (name == "merge") {
-    return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "min") {
     return FoldMINorMAX(context, std::move(funcRef), Ordering::Less);
   } else if (name == "minval") {
     return FoldMaxvalMinval<T>(
         context, std::move(funcRef), RelationalOperator::LT, T::Scalar::HUGE());
   } else if (name == "mod") {
     CHECK(args.size() == 2);
     return FoldElementalIntrinsic<T, T, T>(context, std::move(funcRef),

flang/test/Evaluate/fold-merge.f90

+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of MERGE
+module m
+  type t
+    integer n
+  end type
+  logical, parameter :: test_01 = all(merge([1,2,3],4,[.true.,.false.,.true.]) == [1,4,3])
+  logical, parameter :: test_02 = all(merge([1,2,3],4,.true.) == [1,2,3])
+  logical, parameter :: test_03 = all(merge([1,2,3],4,.false.) == [4,4,4])
+  logical, parameter :: test_04 = all(merge(1,4,[.true.,.false.,.true.,.false.]) == [1,4,1,4])
+  type(t), parameter :: dt00a = merge(t(1),t(2),.true.)
+  logical, parameter :: test_05 = dt00a%n == 1
+  type(t), parameter :: dt00b = merge(t(1),t(2),.false.)
+  logical, parameter :: test_06 = dt00b%n == 2
+  type(t), parameter :: dt01(*) = merge([t(1),t(2)],[t(3),t(4)],[.false.,.true.])
+  logical, parameter :: test_07 = all(dt01%n == [3,2])
+  type(t), parameter :: dt02(*) = merge(t(1),[t(3),t(4)],.true.)
+  logical, parameter :: test_08 = all(dt02%n == [1,1])
+  type(t), parameter :: dt03(*) = merge([t(1),t(2)],t(3),[.true.,.false.])
+  logical, parameter :: test_09 = all(dt03%n == [1,3])
+  logical, parameter :: test_10 = merge('ab','cd',.true.) == 'ab'
+end