[flang] Fold complex component references

[klausler]

Summary

Complex component references (z%RE, z%IM) of complex named constants
should be evaluated at compilation time.

[jeanPerier]

Looks good

flang/include/flang/Evaluate/tools.h

common::IfNoLvalue<Expr<SomeKind<ResultType<A>::category>>, A> AsCategoryExpr(
     A &&x) {
   return Expr<SomeKind<ResultType<A>::category>>{AsExpr(std::move(x))};
 }
 
 Expr<SomeType> Parenthesize(Expr<SomeType> &&);
 
 Expr<SomeReal> GetComplexPart(
     const Expr<SomeComplex> &, bool isImaginary = false);
+Expr<SomeReal> GetComplexPart(Expr<SomeComplex> &&, bool isImaginary = false);
 
 template <int KIND>
 Expr<SomeComplex> MakeComplex(Expr<Type<TypeCategory::Real, KIND>> &&re,
     Expr<Type<TypeCategory::Real, KIND>> &&im) {
   return AsCategoryExpr(ComplexConstructor<KIND>{std::move(re), std::move(im)});
 }
 
 template <typename A> constexpr bool IsNumericCategoryExpr() {

flang/include/flang/Evaluate/type.h

 #define EXPAND_FOR_EACH_INTEGER_KIND(M, P, S) \
   M(P, S, 1) M(P, S, 2) M(P, S, 4) M(P, S, 8) M(P, S, 16)
 #define EXPAND_FOR_EACH_REAL_KIND(M, P, S) \
   M(P, S, 2) M(P, S, 3) M(P, S, 4) M(P, S, 8) M(P, S, 10) M(P, S, 16)
 #define EXPAND_FOR_EACH_COMPLEX_KIND(M, P, S) EXPAND_FOR_EACH_REAL_KIND(M, P, S)
 #define EXPAND_FOR_EACH_CHARACTER_KIND(M, P, S) M(P, S, 1) M(P, S, 2) M(P, S, 4)
 #define EXPAND_FOR_EACH_LOGICAL_KIND(M, P, S) \
   M(P, S, 1) M(P, S, 2) M(P, S, 4) M(P, S, 8)
-#define TEMPLATE_INSTANTIATION(P, S, ARG) P<ARG> S;
 
 #define FOR_EACH_INTEGER_KIND_HELP(PREFIX, SUFFIX, K) \
   PREFIX<Type<TypeCategory::Integer, K>> SUFFIX;
 #define FOR_EACH_REAL_KIND_HELP(PREFIX, SUFFIX, K) \
   PREFIX<Type<TypeCategory::Real, K>> SUFFIX;
 #define FOR_EACH_COMPLEX_KIND_HELP(PREFIX, SUFFIX, K) \
   PREFIX<Type<TypeCategory::Complex, K>> SUFFIX;
 #define FOR_EACH_CHARACTER_KIND_HELP(PREFIX, SUFFIX, K) \

flang/include/flang/Evaluate/variable.h

 // R915 complex-part-designator
 // In the F2018 standard, complex parts of array sections are parsed as
 // variants of sections instead.
 class ComplexPart {
 public:
   ENUM_CLASS(Part, RE, IM)
   CLASS_BOILERPLATE(ComplexPart)
   ComplexPart(DataRef &&z, Part p) : complex_{std::move(z)}, part_{p} {}
+  DataRef &complex() { return complex_; }
   const DataRef &complex() const { return complex_; }
   Part part() const { return part_; }
   int Rank() const;
   const Symbol &GetFirstSymbol() const { return complex_.GetFirstSymbol(); }
   const Symbol &GetLastSymbol() const { return complex_.GetLastSymbol(); }
   bool operator==(const ComplexPart &) const;
   llvm::raw_ostream &AsFortran(llvm::raw_ostream &) const;
 

flang/lib/Evaluate/fold-implementation.h

public:
   std::optional<Constant<T>> ApplySubscripts(const Constant<T> &array,
       const std::vector<Constant<SubscriptInteger>> &subscripts);
   std::optional<Constant<T>> ApplyComponent(Constant<SomeDerived> &&,
       const Symbol &component,
       const std::vector<Constant<SubscriptInteger>> * = nullptr);
   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> PACK(FunctionRef<T> &&);
   Expr<T> RESHAPE(FunctionRef<T> &&);
   Expr<T> SPREAD(FunctionRef<T> &&);

 NamedEntity FoldOperation(FoldingContext &, NamedEntity &&);
 Triplet FoldOperation(FoldingContext &, Triplet &&);
 Subscript FoldOperation(FoldingContext &, Subscript &&);
 ArrayRef FoldOperation(FoldingContext &, ArrayRef &&);
 CoarrayRef FoldOperation(FoldingContext &, CoarrayRef &&);
 DataRef FoldOperation(FoldingContext &, DataRef &&);
 Substring FoldOperation(FoldingContext &, Substring &&);
 ComplexPart FoldOperation(FoldingContext &, ComplexPart &&);
-
 template <typename T>
-Expr<T> FoldOperation(FoldingContext &context, FunctionRef<T> &&);
-template <int KIND>
-Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
-    FoldingContext &context, FunctionRef<Type<TypeCategory::Integer, KIND>> &&);
-template <int KIND>
-Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
-    FoldingContext &context, FunctionRef<Type<TypeCategory::Real, KIND>> &&);
-template <int KIND>
-Expr<Type<TypeCategory::Complex, KIND>> FoldIntrinsicFunction(
-    FoldingContext &context, FunctionRef<Type<TypeCategory::Complex, KIND>> &&);
-template <int KIND>
-Expr<Type<TypeCategory::Logical, KIND>> FoldIntrinsicFunction(
-    FoldingContext &context, FunctionRef<Type<TypeCategory::Logical, KIND>> &&);
-
+Expr<T> FoldOperation(FoldingContext &, FunctionRef<T> &&);
 template <typename T>
 Expr<T> FoldOperation(FoldingContext &context, Designator<T> &&designator) {
   return Folder<T>{context}.Folding(std::move(designator));
 }
-
 Expr<TypeParamInquiry::Result> FoldOperation(
     FoldingContext &, TypeParamInquiry &&);
 Expr<ImpliedDoIndex::Result> FoldOperation(
     FoldingContext &context, ImpliedDoIndex &&);
 template <typename T>
 Expr<T> FoldOperation(FoldingContext &, ArrayConstructor<T> &&);
 Expr<SomeDerived> FoldOperation(FoldingContext &, StructureConstructor &&);
 
return GetConstantComponent(*component, &subscripts);
                  GetNamedConstant(aRef.base().GetLastSymbol())}) {
     return ApplySubscripts(*array, subscripts);
   } else {
     return std::nullopt;
   }
 }
 
 template <typename T>
+std::optional<Constant<T>> Folder<T>::Folding(DataRef &ref) {
+  return common::visit(
+      common::visitors{
+          [this](SymbolRef &sym) { return GetNamedConstant(*sym); },
+          [this](Component &comp) {
+            comp = FoldOperation(context_, std::move(comp));
+            return GetConstantComponent(comp);
+          },
+          [this](ArrayRef &aRef) {
+            aRef = FoldOperation(context_, std::move(aRef));
+            return Folding(aRef);
+          },
+          [](CoarrayRef &) { return std::optional<Constant<T>>{}; },
+      },
+      ref.u);
+}
+
+// TODO: This would be more natural as a member function of Constant<T>.
+template <typename T>
 std::optional<Constant<T>> Folder<T>::ApplySubscripts(const Constant<T> &array,
     const std::vector<Constant<SubscriptInteger>> &subscripts) {
   const auto &shape{array.shape()};
   const auto &lbounds{array.lbounds()};
   int rank{GetRank(shape)};
   CHECK(rank == static_cast<int>(subscripts.size()));
   std::size_t elements{1};
   ConstantSubscripts resultShape;
if (auto *substring{common::Unwrap<Substring>(designator.u)}) {
         }
       }
       if (auto length{ToInt64(Fold(context_, substring->LEN()))}) {
         if (*length == 0) {
           return Expr<T>{Constant<T>{Scalar<T>{}}};
         }
       }
     }
+  } else if constexpr (T::category == TypeCategory::Real) {
+    if (auto *zPart{std::get_if<ComplexPart>(&designator.u)}) {
+      *zPart = FoldOperation(context_, std::move(*zPart));
+      using ComplexT = Type<TypeCategory::Complex, T::kind>;
+      if (auto zConst{Folder<ComplexT>{context_}.Folding(zPart->complex())}) {
+        return Fold(context_,
+            Expr<T>{ComplexComponent<T::kind>{
+                zPart->part() == ComplexPart::Part::IM,
+                Expr<ComplexT>{std::move(*zConst)}}});
+      } else {
+        return Expr<T>{Designator<T>{std::move(*zPart)}};
+      }
+    }
   }
   return common::visit(
       common::visitors{
           [&](SymbolRef &&symbol) {
             if (auto constant{GetNamedConstant(*symbol)}) {
               return Expr<T>{std::move(*constant)};
             }
             return Expr<T>{std::move(designator)};
Expr<T> RewriteSpecificMINorMAX(
   }};
   if (auto *sx{UnwrapExpr<Expr<SomeReal>>(*resultTypeArg)}) {
     return common::visit(insertConversion, sx->u);
   }
   auto &sx{DEREF(UnwrapExpr<Expr<SomeInteger>>(*resultTypeArg))};
   return common::visit(insertConversion, sx.u);
 }
 
+// FoldIntrinsicFunction()
+template <int KIND>
+Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
+    FoldingContext &context, FunctionRef<Type<TypeCategory::Integer, KIND>> &&);
+template <int KIND>
+Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
+    FoldingContext &context, FunctionRef<Type<TypeCategory::Real, KIND>> &&);
+template <int KIND>
+Expr<Type<TypeCategory::Complex, KIND>> FoldIntrinsicFunction(
+    FoldingContext &context, FunctionRef<Type<TypeCategory::Complex, KIND>> &&);
+template <int KIND>
+Expr<Type<TypeCategory::Logical, KIND>> FoldIntrinsicFunction(
+    FoldingContext &context, FunctionRef<Type<TypeCategory::Logical, KIND>> &&);
+
 template <typename T>
 Expr<T> FoldOperation(FoldingContext &context, FunctionRef<T> &&funcRef) {
   ActualArguments &args{funcRef.arguments()};
   for (std::optional<ActualArgument> &arg : args) {
     if (auto *expr{UnwrapExpr<Expr<SomeType>>(arg)}) {
       *expr = Fold(context, std::move(*expr));
     }
   }
common::visit(
         } else {
           common::die("ToReal: bad argument expression");
         }
       },
       std::move(expr.u));
   return result.value();
 }
 
+// REAL(z) and AIMAG(z)
+template <int KIND>
+Expr<Type<TypeCategory::Real, KIND>> FoldOperation(
+    FoldingContext &context, ComplexComponent<KIND> &&x) {
+  using Operand = Type<TypeCategory::Complex, KIND>;
+  using Result = Type<TypeCategory::Real, KIND>;
+  if (auto array{ApplyElementwise(context, x,
+          std::function<Expr<Result>(Expr<Operand> &&)>{
+              [=](Expr<Operand> &&operand) {
+                return Expr<Result>{ComplexComponent<KIND>{
+                    x.isImaginaryPart, std::move(operand)}};
+              }})}) {
+    return *array;
+  }
+  auto &operand{x.left()};
+  if (auto value{GetScalarConstantValue<Operand>(operand)}) {
+    if (x.isImaginaryPart) {
+      return Expr<Result>{Constant<Result>{value->AIMAG()}};
+    } else {
+      return Expr<Result>{Constant<Result>{value->REAL()}};
+    }
+  }
+  return Expr<Result>{std::move(x)};
+}
+
 template <typename T>
 Expr<T> ExpressionBase<T>::Rewrite(FoldingContext &context, Expr<T> &&expr) {
   return common::visit(
       [&](auto &&x) -> Expr<T> {
         if constexpr (IsSpecificIntrinsicType<T>) {
           return FoldOperation(context, std::move(x));
         } else if constexpr (std::is_same_v<T, SomeDerived>) {
           return FoldOperation(context, std::move(x));
         } else if constexpr (common::HasMember<decltype(x),
                                  TypelessExpression>) {
           return std::move(expr);
         } else {
           return Expr<T>{Fold(context, std::move(x))};
         }
       },
       std::move(expr.u));
 }
 
 FOR_EACH_TYPE_AND_KIND(extern template class ExpressionBase, )
-
 } // namespace Fortran::evaluate
 #endif // FORTRAN_EVALUATE_FOLD_IMPLEMENTATION_H_

flang/lib/Evaluate/fold-real.cpp

if (auto *x{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
       return FoldElementalIntrinsic<T, ComplexT>(context, std::move(funcRef),
           ScalarFunc<T, ComplexT>([](const Scalar<ComplexT> &z) -> Scalar<T> {
             return z.ABS().value;
           }));
     } else {
       common::die(" unexpected argument type inside abs");
     }
   } else if (name == "aimag") {
-    return FoldElementalIntrinsic<T, ComplexT>(
-        context, std::move(funcRef), &Scalar<ComplexT>::AIMAG);
+    if (auto *zExpr{UnwrapExpr<Expr<ComplexT>>(args[0])}) {
+      return Fold(context, Expr<T>{ComplexComponent{true, std::move(*zExpr)}});
+    }
   } else if (name == "aint" || name == "anint") {
     // ANINT rounds ties away from zero, not to even
     common::RoundingMode mode{name == "aint"
             ? common::RoundingMode::ToZero
             : common::RoundingMode::TiesAwayFromZero};
     return FoldElementalIntrinsic<T, T>(context, std::move(funcRef),
         ScalarFunc<T, T>([&name, &context, mode](
                              const Scalar<T> &x) -> Scalar<T> {
return FoldElementalIntrinsic<T, T>(context, std::move(funcRef),
           }
           return result.value;
         }));
   }
   // TODO: dot_product, fraction, matmul, norm2, set_exponent, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
-template <int KIND>
-Expr<Type<TypeCategory::Real, KIND>> FoldOperation(
-    FoldingContext &context, ComplexComponent<KIND> &&x) {
-  using Operand = Type<TypeCategory::Complex, KIND>;
-  using Result = Type<TypeCategory::Real, KIND>;
-  if (auto array{ApplyElementwise(context, x,
-          std::function<Expr<Result>(Expr<Operand> &&)>{
-              [=](Expr<Operand> &&operand) {
-                return Expr<Result>{ComplexComponent<KIND>{
-                    x.isImaginaryPart, std::move(operand)}};
-              }})}) {
-    return *array;
-  }
-  using Part = Type<TypeCategory::Real, KIND>;
-  auto &operand{x.left()};
-  if (auto value{GetScalarConstantValue<Operand>(operand)}) {
-    if (x.isImaginaryPart) {
-      return Expr<Part>{Constant<Part>{value->AIMAG()}};
-    } else {
-      return Expr<Part>{Constant<Part>{value->REAL()}};
-    }
-  }
-  return Expr<Part>{std::move(x)};
-}
-
 #ifdef _MSC_VER // disable bogus warning about missing definitions
 #pragma warning(disable : 4661)
 #endif
 FOR_EACH_REAL_KIND(template class ExpressionBase, )
 template class ExpressionBase<SomeReal>;
 } // namespace Fortran::evaluate

flang/lib/Evaluate/tools.cpp

Expr<SomeReal> GetComplexPart(const Expr<SomeComplex> &z, bool isImaginary) {
   return common::visit(
       [&](const auto &zk) {
         static constexpr int kind{ResultType<decltype(zk)>::kind};
         return AsCategoryExpr(ComplexComponent<kind>{isImaginary, zk});
       },
       z.u);
 }
 
+Expr<SomeReal> GetComplexPart(Expr<SomeComplex> &&z, bool isImaginary) {
+  return common::visit(
+      [&](auto &&zk) {
+        static constexpr int kind{ResultType<decltype(zk)>::kind};
+        return AsCategoryExpr(
+            ComplexComponent<kind>{isImaginary, std::move(zk)});
+      },
+      z.u);
+}
+
 // Convert REAL to COMPLEX of the same kind. Preserving the real operand kind
 // and then applying complex operand promotion rules allows the result to have
 // the highest precision of REAL and COMPLEX operands as required by Fortran
 // 2018 10.9.1.3.
 Expr<SomeComplex> PromoteRealToComplex(Expr<SomeReal> &&someX) {
   return common::visit(
       [](auto &&x) {
         using RT = ResultType<decltype(x)>;

flang/lib/Semantics/expression.cpp

MaybeExpr ExpressionAnalyzer::Analyze(const parser::StructureComponent &sc) {
   } else if (auto *details{sym->detailsIf<semantics::MiscDetails>()}) {
     // special part-ref: %re, %im, %kind, %len
     // Type errors on the base of %re/%im/%len are detected and
     // reported in name resolution.
     using MiscKind = semantics::MiscDetails::Kind;
     MiscKind kind{details->kind()};
     if (kind == MiscKind::ComplexPartRe || kind == MiscKind::ComplexPartIm) {
       if (auto *zExpr{std::get_if<Expr<SomeComplex>>(&base->u)}) {
-        if (std::optional<DataRef> dataRef{ExtractDataRef(std::move(*zExpr))}) {
+        if (std::optional<DataRef> dataRef{ExtractDataRef(*zExpr)}) {
+          // Represent %RE/%IM as a designator
           Expr<SomeReal> realExpr{common::visit(
               [&](const auto &z) {
                 using PartType = typename ResultType<decltype(z)>::Part;
                 auto part{kind == MiscKind::ComplexPartRe
                         ? ComplexPart::Part::RE
                         : ComplexPart::Part::IM};
                 return AsCategoryExpr(Designator<PartType>{
                     ComplexPart{std::move(*dataRef), part}});

flang/test/Evaluate/fold-re-im.f90

+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of complex components
+module m
+  complex, parameter :: z = (1., 2.)
+  logical, parameter :: test_1 = z%re == 1.
+  logical, parameter :: test_2 = z%im == 2.
+  logical, parameter :: test_3 = real(z+z) == 2.
+  logical, parameter :: test_4 = aimag(z+z) == 4.
+  type :: t
+    complex :: z
+  end type
+  type(t), parameter :: tz(*) = [t((3., 4.)), t((5., 6.))]
+  logical, parameter :: test_5 = all(tz%z%re == [3., 5.])
+  logical, parameter :: test_6 = all(tz%z%im == [4., 6.])
+end module