[flang] Enforce constraint: defined ass't in WHERE must be elemental

[klausler]

Summary

A defined assignment subroutine invoked in the context of a WHERE
statement or construct must necessarily be elemental (C1032).

[klausler]

ping

[jeanPerier]

LGTM

flang/include/flang/Semantics/expression.h

 
 class ExpressionAnalyzer {
 public:
   using MaybeExpr = std::optional<Expr<SomeType>>;
 
   explicit ExpressionAnalyzer(semantics::SemanticsContext &sc) : context_{sc} {}
   ExpressionAnalyzer(semantics::SemanticsContext &sc, FoldingContext &fc)
       : context_{sc}, foldingContext_{fc} {}
-  ExpressionAnalyzer(ExpressionAnalyzer &) = default;
+  ExpressionAnalyzer(const ExpressionAnalyzer &) = default;
 
   semantics::SemanticsContext &context() const { return context_; }
+  bool inWhereBody() const { return inWhereBody_; }
+  void set_inWhereBody(bool yes = true) { inWhereBody_ = yes; }
 
   FoldingContext &GetFoldingContext() const { return foldingContext_; }
 
   parser::ContextualMessages &GetContextualMessages() {
     return foldingContext_.messages();
   }
 
   template <typename... A> parser::Message *Say(A &&...args) {
private:
   }
   bool CheckIsValidForwardReference(const semantics::DerivedTypeSpec &);
 
   semantics::SemanticsContext &context_;
   FoldingContext &foldingContext_{context_.foldingContext()};
   std::map<parser::CharBlock, int> impliedDos_; // values are INTEGER kinds
   bool isWholeAssumedSizeArrayOk_{false};
   bool useSavedTypedExprs_{true};
+  bool inWhereBody_{false};
   friend class ArgumentAnalyzer;
 };
 
 inline bool AreConformable(int leftRank, int rightRank) {
   return leftRank == 0 || rightRank == 0 || leftRank == rightRank;
 }
 
 template <typename L, typename R>
std::optional<evaluate::Expr<evaluate::SomeType>> AnalyzeExpr(
   return evaluate::ExpressionAnalyzer{context}.Analyze(expr);
 }
 
 // Semantic analysis of an intrinsic type's KIND parameter expression.
 evaluate::Expr<evaluate::SubscriptInteger> AnalyzeKindSelector(
     SemanticsContext &, common::TypeCategory,
     const std::optional<parser::KindSelector> &);
 
-void AnalyzeCallStmt(SemanticsContext &, const parser::CallStmt &);
-const evaluate::Assignment *AnalyzeAssignmentStmt(
-    SemanticsContext &, const parser::AssignmentStmt &);
-const evaluate::Assignment *AnalyzePointerAssignmentStmt(
-    SemanticsContext &, const parser::PointerAssignmentStmt &);
-
 // Semantic analysis of all expressions in a parse tree, which becomes
 // decorated with typed representations for top-level expressions.
 class ExprChecker {
 public:
   explicit ExprChecker(SemanticsContext &);
 
   template <typename A> bool Pre(const A &) { return true; }
   template <typename A> void Post(const A &) {}
public:
   }
   bool Pre(const parser::PointerObject &x) {
     exprAnalyzer_.Analyze(x);
     return false;
   }
   bool Pre(const parser::DataImpliedDo &);
 
   bool Pre(const parser::CallStmt &x) {
-    AnalyzeCallStmt(context_, x);
+    exprAnalyzer_.Analyze(x);
     return false;
   }
   bool Pre(const parser::AssignmentStmt &x) {
-    AnalyzeAssignmentStmt(context_, x);
+    exprAnalyzer_.Analyze(x);
     return false;
   }
   bool Pre(const parser::PointerAssignmentStmt &x) {
-    AnalyzePointerAssignmentStmt(context_, x);
+    exprAnalyzer_.Analyze(x);
     return false;
   }
 
+  // Track whether we're in a WHERE statement or construct body
+  bool Pre(const parser::WhereStmt &) {
+    ++whereDepth_;
+    exprAnalyzer_.set_inWhereBody(InWhereBody());
+    return true;
+  }
+  void Post(const parser::WhereStmt &) {
+    --whereDepth_;
+    exprAnalyzer_.set_inWhereBody(InWhereBody());
+  }
+  bool Pre(const parser::WhereBodyConstruct &) {
+    ++whereDepth_;
+    exprAnalyzer_.set_inWhereBody(InWhereBody());
+    return true;
+  }
+  void Post(const parser::WhereBodyConstruct &) {
+    --whereDepth_;
+    exprAnalyzer_.set_inWhereBody(InWhereBody());
+  }
+
   template <typename A> bool Pre(const parser::Scalar<A> &x) {
     exprAnalyzer_.Analyze(x);
     return false;
   }
   template <typename A> bool Pre(const parser::Constant<A> &x) {
     exprAnalyzer_.Analyze(x);
     return false;
   }
   template <typename A> bool Pre(const parser::Integer<A> &x) {
     exprAnalyzer_.Analyze(x);
     return false;
   }
   template <typename A> bool Pre(const parser::Logical<A> &x) {
     exprAnalyzer_.Analyze(x);
     return false;
   }
   template <typename A> bool Pre(const parser::DefaultChar<A> &x) {
     exprAnalyzer_.Analyze(x);
     return false;
   }
 
 private:
+  bool InWhereBody() const { return whereDepth_ > 0; }
+
   SemanticsContext &context_;
   evaluate::ExpressionAnalyzer exprAnalyzer_{context_};
+  int whereDepth_{0}; // nesting of WHERE statements & constructs
 };
 } // namespace Fortran::semantics
 #endif // FORTRAN_SEMANTICS_EXPRESSION_H_

flang/lib/Semantics/expression.cpp

std::optional<ProcedureRef> ArgumentAnalyzer::TryDefinedAssignment() {
   if (isDefined == Tristate::No) {
     if (lhsType && rhsType) {
       AddAssignmentConversion(*lhsType, *rhsType);
     }
     return std::nullopt; // user-defined assignment not allowed for these args
   }
   auto restorer{context_.GetContextualMessages().SetLocation(source_)};
   if (std::optional<ProcedureRef> procRef{GetDefinedAssignmentProc()}) {
+    if (context_.inWhereBody() && !procRef->proc().IsElemental()) { // C1032
+      context_.Say(
+          "Defined assignment in WHERE must be elemental, but '%s' is not"_err_en_US,
+          DEREF(procRef->proc().GetSymbol()).name());
+    }
     context_.CheckCall(source_, procRef->proc(), procRef->arguments());
     return std::move(*procRef);
   }
   if (isDefined == Tristate::Yes) {
     if (!lhsType || !rhsType || (lhsRank != rhsRank && rhsRank != 0) ||
         !OkLogicalIntegerAssignment(lhsType->category(), rhsType->category())) {
       SayNoMatch("ASSIGNMENT(=)", true);
     }
evaluate::Expr<evaluate::SubscriptInteger> AnalyzeKindSelector(
     SemanticsContext &context, common::TypeCategory category,
     const std::optional<parser::KindSelector> &selector) {
   evaluate::ExpressionAnalyzer analyzer{context};
   auto restorer{
       analyzer.GetContextualMessages().SetLocation(context.location().value())};
   return analyzer.AnalyzeKindSelector(category, selector);
 }
 
-void AnalyzeCallStmt(SemanticsContext &context, const parser::CallStmt &call) {
-  evaluate::ExpressionAnalyzer{context}.Analyze(call);
-}
-
-const evaluate::Assignment *AnalyzeAssignmentStmt(
-    SemanticsContext &context, const parser::AssignmentStmt &stmt) {
-  return evaluate::ExpressionAnalyzer{context}.Analyze(stmt);
-}
-const evaluate::Assignment *AnalyzePointerAssignmentStmt(
-    SemanticsContext &context, const parser::PointerAssignmentStmt &stmt) {
-  return evaluate::ExpressionAnalyzer{context}.Analyze(stmt);
-}
-
 ExprChecker::ExprChecker(SemanticsContext &context) : context_{context} {}
 
 bool ExprChecker::Pre(const parser::DataImpliedDo &ido) {
   parser::Walk(std::get<parser::DataImpliedDo::Bounds>(ido.t), *this);
   const auto &bounds{std::get<parser::DataImpliedDo::Bounds>(ido.t)};
   auto name{bounds.name.thing.thing};
   int kind{evaluate::ResultType<evaluate::ImpliedDoIndex>::kind};
   if (const auto dynamicType{evaluate::DynamicType::From(*name.symbol)}) {

flang/test/Semantics/assign04.f90

subroutine sub(arg1, arg2, arg3)
       local1 = local2
       local2 = local3
       !ERROR: No intrinsic or user-defined ASSIGNMENT(=) matches operand types TYPE(dtype(k1=2_4,k2=4_4,l1=1_4,l2=3_4)) and TYPE(dtype(k1=2_4,k2=4_4,l1=9_4,l2=3_4))
       local1 = local4 ! mismatched constant KIND type parameter
       !ERROR: No intrinsic or user-defined ASSIGNMENT(=) matches operand types TYPE(dtype(k1=2_4,k2=4_4,l1=1_4,l2=3_4)) and TYPE(dtype(k1=9_4,k2=4_4,l1=1_4,l2=3_4))
       local1 = local5 ! mismatched constant LEN type parameter
     end subroutine sub
 end subroutine s12
+
+subroutine s13()
+  interface assignment(=)
+    procedure :: cToR, cToRa, cToI
+  end interface
+  real :: x(1)
+  integer :: n(1)
+  x='0' ! fine
+  n='0' ! fine
+  !ERROR: Defined assignment in WHERE must be elemental, but 'ctora' is not
+  where ([1==1]) x='*'
+  where ([1==1]) n='*' ! fine
+  forall (j=1:1)
+    where (j==1)
+      !ERROR: Defined assignment in WHERE must be elemental, but 'ctor' is not
+      x(j)='?'
+      n(j)='?' ! fine
+    elsewhere (.false.)
+      !ERROR: Defined assignment in WHERE must be elemental, but 'ctor' is not
+      x(j)='1'
+      n(j)='1' ! fine
+    elsewhere
+      !ERROR: Defined assignment in WHERE must be elemental, but 'ctor' is not
+      x(j)='9'
+      n(j)='9' ! fine
+    end where
+  end forall
+  x='0' ! still fine
+  n='0' ! still fine
+ contains
+  subroutine cToR(x, c)
+    real, intent(out) :: x
+    character, intent(in) :: c
+  end subroutine
+  subroutine cToRa(x, c)
+    real, intent(out) :: x(:)
+    character, intent(in) :: c
+  end subroutine
+  elemental subroutine cToI(n, c)
+    integer, intent(out) :: n
+    character, intent(in) :: c
+  end subroutine
+end subroutine s13