[TypeChecker] Type-check where clauses for for-in statements separately

xedin · xedin · commit e7a9bcb303a5 · 2024-12-21T00:42:13.000-08:00
Instead of using `one-way` constraints, just like in closure contexts
for-in statements should type-check their `where` clauses separately.

This also unifies and simplifies for-in preamble handling in the
solver.
diff --git a/include/swift/Sema/SyntacticElementTarget.h b/include/swift/Sema/SyntacticElementTarget.h
@@ -161,10 +161,9 @@ class SyntacticElementTarget {
       ForEachStmt *stmt;
       DeclContext *dc;
       Pattern *pattern;
-      bool ignoreWhereClause;
       GenericEnvironment *packElementEnv;
       ForEachStmtInfo info;
-    } forEachStmt;
+    } forEachPreamble;
 
     PatternBindingDecl *patternBinding;
   };
@@ -242,13 +241,11 @@ class SyntacticElementTarget {
   }
 
   SyntacticElementTarget(ForEachStmt *stmt, DeclContext *dc,
-                         bool ignoreWhereClause,
                          GenericEnvironment *packElementEnv)
       : kind(Kind::forEachPreamble) {
-    forEachStmt.stmt = stmt;
-    forEachStmt.dc = dc;
-    forEachStmt.ignoreWhereClause = ignoreWhereClause;
-    forEachStmt.packElementEnv = packElementEnv;
+    forEachPreamble.stmt = stmt;
+    forEachPreamble.dc = dc;
+    forEachPreamble.packElementEnv = packElementEnv;
   }
 
   /// Form a target for the initialization of a pattern from an expression.
@@ -267,10 +264,10 @@ class SyntacticElementTarget {
   static SyntacticElementTarget
   forReturn(ReturnStmt *returnStmt, Type contextTy, DeclContext *dc);
 
-  /// Form a target for the preamble of a for-in loop, excluding its body.
+  /// Form a target for the preamble of a for-in loop, excluding its where
+  /// clause and body.
   static SyntacticElementTarget
   forForEachPreamble(ForEachStmt *stmt, DeclContext *dc,
-                     bool ignoreWhereClause = false,
                      GenericEnvironment *packElementEnv = nullptr);
 
   /// Form a target for a property with an attached property wrapper that is
@@ -372,7 +369,7 @@ class SyntacticElementTarget {
     }
 
     case Kind::forEachPreamble:
-      return forEachStmt.dc;
+      return forEachPreamble.dc;
     }
     llvm_unreachable("invalid decl context type");
   }
@@ -541,24 +538,19 @@ class SyntacticElementTarget {
     return expression.initialization.patternBindingIndex;
   }
 
-  bool ignoreForEachWhereClause() const {
-    assert(isForEachPreamble());
-    return forEachStmt.ignoreWhereClause;
-  }
-
   GenericEnvironment *getPackElementEnv() const {
     assert(isForEachPreamble());
-    return forEachStmt.packElementEnv;
+    return forEachPreamble.packElementEnv;
   }
 
   const ForEachStmtInfo &getForEachStmtInfo() const {
     assert(isForEachPreamble());
-    return forEachStmt.info;
+    return forEachPreamble.info;
   }
 
   ForEachStmtInfo &getForEachStmtInfo() {
     assert(isForEachPreamble());
-    return forEachStmt.info;
+    return forEachPreamble.info;
   }
 
   /// Whether this context infers an opaque return type.
@@ -585,7 +577,7 @@ class SyntacticElementTarget {
       return getInitializationPattern();
 
     if (kind == Kind::forEachPreamble)
-      return forEachStmt.pattern;
+      return forEachPreamble.pattern;
 
     return nullptr;
   }
@@ -598,7 +590,7 @@ class SyntacticElementTarget {
     }
 
     if (kind == Kind::forEachPreamble) {
-      forEachStmt.pattern = pattern;
+      forEachPreamble.pattern = pattern;
       return;
     }
 
@@ -729,7 +721,7 @@ class SyntacticElementTarget {
       return nullptr;
 
     case Kind::forEachPreamble:
-      return forEachStmt.stmt;
+      return forEachPreamble.stmt;
     }
     llvm_unreachable("invalid case label type");
   }
@@ -841,7 +833,7 @@ class SyntacticElementTarget {
 
     // For-in preamble target doesn't cover the body.
     case Kind::forEachPreamble:
-      auto *stmt = forEachStmt.stmt;
+      auto *stmt = forEachPreamble.stmt;
       SourceLoc startLoc = stmt->getForLoc();
       SourceLoc endLoc = stmt->getParsedSequence()->getEndLoc();
 
@@ -884,7 +876,7 @@ class SyntacticElementTarget {
     }
 
     case Kind::forEachPreamble:
-      return forEachStmt.stmt->getStartLoc();
+      return forEachPreamble.stmt->getStartLoc();
     }
     llvm_unreachable("invalid target type");
   }
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -9247,9 +9247,9 @@ applySolutionToInitialization(SyntacticElementTarget target, Expr *initializer,
 }
 
 static std::optional<SequenceIterationInfo>
-applySolutionToForEachStmt(ForEachStmt *stmt, SequenceIterationInfo info,
-                           DeclContext *dc,
-                           SyntacticElementTargetRewriter &rewriter) {
+applySolutionToForEachStmtPreamble(ForEachStmt *stmt,
+                                   SequenceIterationInfo info, DeclContext *dc,
+                                   SyntacticElementTargetRewriter &rewriter) {
   auto &solution = rewriter.getSolution();
   auto &cs = solution.getConstraintSystem();
   auto &ctx = cs.getASTContext();
@@ -9374,23 +9374,12 @@ applySolutionToForEachStmt(ForEachStmt *stmt, SequenceIterationInfo info,
          "Couldn't find sequence conformance");
   stmt->setSequenceConformance(type, sequenceConformance);
 
-  // Apply the solution to the filtering condition, if there is one.
-  if (auto *whereExpr = stmt->getWhere()) {
-    auto whereTarget = *cs.getTargetFor(whereExpr);
-
-    auto rewrittenTarget = rewriter.rewriteTarget(whereTarget);
-    if (!rewrittenTarget)
-      return std::nullopt;
-
-    stmt->setWhere(rewrittenTarget->getAsExpr());
-  }
-
   return info;
 }
 
 static std::optional<PackIterationInfo>
-applySolutionToForEachStmt(ForEachStmt *stmt, PackIterationInfo info,
-                           SyntacticElementTargetRewriter &rewriter) {
+applySolutionToForEachStmtPreamble(ForEachStmt *stmt, PackIterationInfo info,
+                                   SyntacticElementTargetRewriter &rewriter) {
   auto &solution = rewriter.getSolution();
   auto &cs = solution.getConstraintSystem();
   auto *sequenceExpr = stmt->getParsedSequence();
@@ -9412,16 +9401,16 @@ applySolutionToForEachStmt(ForEachStmt *stmt, PackIterationInfo info,
 ///
 /// \returns the resulting initialization expression.
 static std::optional<SyntacticElementTarget>
-applySolutionToForEachStmt(SyntacticElementTarget target,
-                           SyntacticElementTargetRewriter &rewriter) {
+applySolutionToForEachStmtPreamble(SyntacticElementTarget target,
+                                   SyntacticElementTargetRewriter &rewriter) {
   auto resultTarget = target;
   auto &forEachStmtInfo = resultTarget.getForEachStmtInfo();
   auto *stmt = target.getAsForEachStmt();
 
   Type rewrittenPatternType;
 
   if (auto *info = forEachStmtInfo.dyn_cast<SequenceIterationInfo>()) {
-    auto resultInfo = applySolutionToForEachStmt(
+    auto resultInfo = applySolutionToForEachStmtPreamble(
         stmt, *info, target.getDeclContext(), rewriter);
     if (!resultInfo) {
       return std::nullopt;
@@ -9430,7 +9419,7 @@ applySolutionToForEachStmt(SyntacticElementTarget target,
     forEachStmtInfo = *resultInfo;
     rewrittenPatternType = resultInfo->initType;
   } else {
-    auto resultInfo = applySolutionToForEachStmt(
+    auto resultInfo = applySolutionToForEachStmtPreamble(
         stmt, forEachStmtInfo.get<PackIterationInfo>(), rewriter);
     if (!resultInfo) {
       return std::nullopt;
@@ -9664,11 +9653,12 @@ ExprWalker::rewriteTarget(SyntacticElementTarget target) {
 
     return std::nullopt;
   } else if (auto *forEach = target.getAsForEachStmt()) {
-    auto forEachResultTarget = applySolutionToForEachStmt(target, *this);
-    if (!forEachResultTarget)
+    auto forEachPreambleResultTarget =
+        applySolutionToForEachStmtPreamble(target, *this);
+    if (!forEachPreambleResultTarget)
       return std::nullopt;
 
-    result = *forEachResultTarget;
+    result = *forEachPreambleResultTarget;
   } else {
     auto fn = *target.getAsFunction();
     if (rewriteFunction(fn))
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -3660,8 +3660,7 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
 /// expression that conforms to `Swift.Sequence`.
 static std::optional<SequenceIterationInfo>
 generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
-                               ForEachStmt *stmt, Pattern *typeCheckedPattern,
-                               bool ignoreForEachWhereClause) {
+                               ForEachStmt *stmt, Pattern *typeCheckedPattern) {
   ASTContext &ctx = cs.getASTContext();
   bool isAsync = stmt->getAwaitLoc().isValid();
   auto *sequenceExpr = stmt->getParsedSequence();
@@ -3835,24 +3834,6 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
                      elementLocator);
   }
 
-  // Generate constraints for the "where" expression, if there is one.
-  auto *whereExpr = stmt->getWhere();
-  if (whereExpr && !ignoreForEachWhereClause) {
-    Type boolType = dc->getASTContext().getBoolType();
-    if (!boolType)
-      return std::nullopt;
-
-    SyntacticElementTarget whereTarget(whereExpr, dc, CTP_Condition, boolType,
-                                       /*isDiscarded=*/false);
-    if (cs.generateConstraints(whereTarget, FreeTypeVariableBinding::Disallow))
-      return std::nullopt;
-
-    cs.setTargetFor(whereExpr, whereTarget);
-
-    ContextualTypeInfo contextInfo(boolType, CTP_Condition);
-    cs.setContextualInfo(whereExpr, contextInfo);
-  }
-
   // Populate all of the information for a for-each loop.
   sequenceIterationInfo.elementType = elementType;
   sequenceIterationInfo.initType = initType;
@@ -3904,8 +3885,8 @@ generateForEachPreambleConstraints(ConstraintSystem &cs,
 
     target.getForEachStmtInfo() = *packIterationInfo;
   } else {
-    auto sequenceIterationInfo = generateForEachStmtConstraints(
-        cs, dc, stmt, pattern, target.ignoreForEachWhereClause());
+    auto sequenceIterationInfo =
+        generateForEachStmtConstraints(cs, dc, stmt, pattern);
     if (!sequenceIterationInfo) {
       return std::nullopt;
     }
diff --git a/lib/Sema/CSSyntacticElement.cpp b/lib/Sema/CSSyntacticElement.cpp
@@ -686,13 +686,8 @@ class SyntacticElementConstraintGenerator
   ///
   /// - From sequence to pattern, when pattern has no type information.
   void visitForEachPattern(Pattern *pattern, ForEachStmt *forEachStmt) {
-    // The `where` clause should be ignored because \c visitForEachStmt
-    // records it as a separate conjunction element to allow for a more
-    // granular control over what contextual information is brought into
-    // the scope during pattern + sequence and `where` clause solving.
     auto target = SyntacticElementTarget::forForEachPreamble(
-        forEachStmt, context.getAsDeclContext(),
-        /*ignoreWhereClause=*/true);
+        forEachStmt, context.getAsDeclContext());
 
     if (cs.generateConstraints(target)) {
       hadError = true;
@@ -1898,10 +1893,20 @@ class SyntacticElementSolutionApplication
   ASTNode visitForEachStmt(ForEachStmt *forEachStmt) {
     ConstraintSystem &cs = solution.getConstraintSystem();
 
-    auto forEachTarget = rewriter.rewriteTarget(*cs.getTargetFor(forEachStmt));
-
-    if (!forEachTarget)
+    // Apply solution to the preamble first.
+    if (!rewriter.rewriteTarget(*cs.getTargetFor(forEachStmt))) {
       hadError = true;
+    }
+
+    // Then apply the solution to the filtering condition, if there is one.
+    if (auto *whereExpr = forEachStmt->getWhere()) {
+      auto whereTarget = *cs.getTargetFor(whereExpr);
+      if (auto rewrittenWhereTarget = rewriter.rewriteTarget(whereTarget)) {
+        forEachStmt->setWhere(rewrittenWhereTarget->getAsExpr());
+      } else {
+        hadError = true;
+      }
+    }
 
     auto body = visit(forEachStmt->getBody()).get<Stmt *>();
     forEachStmt->setBody(cast<BraceStmt>(body));
diff --git a/lib/Sema/SyntacticElementTarget.cpp b/lib/Sema/SyntacticElementTarget.cpp
@@ -192,9 +192,8 @@ SyntacticElementTarget::forReturn(ReturnStmt *returnStmt, Type contextTy,
 
 SyntacticElementTarget
 SyntacticElementTarget::forForEachPreamble(ForEachStmt *stmt, DeclContext *dc,
-                                           bool ignoreWhereClause,
                                            GenericEnvironment *packElementEnv) {
-  SyntacticElementTarget target(stmt, dc, ignoreWhereClause, packElementEnv);
+  SyntacticElementTarget target(stmt, dc, packElementEnv);
   return target;
 }
 
@@ -234,8 +233,8 @@ ContextualPattern SyntacticElementTarget::getContextualPattern() const {
   }
 
   if (isForEachPreamble()) {
-    return ContextualPattern::forRawPattern(forEachStmt.pattern,
-                                            forEachStmt.dc);
+    return ContextualPattern::forRawPattern(forEachPreamble.pattern,
+                                            forEachPreamble.dc);
   }
 
   auto ctp = getExprContextualTypePurpose();
@@ -400,7 +399,7 @@ SyntacticElementTarget::walk(ASTWalker &walker) const {
     break;
   }
   case Kind::forEachPreamble: {
-    // We need to skip the where clause if requested, and we currently do not
+    // We need to skip the where clause, and we currently do not
     // type-check a for loop's BraceStmt as part of the SyntacticElementTarget,
     // so we need to skip it here.
     // TODO: We ought to be able to fold BraceStmt checking into the constraint
@@ -421,8 +420,7 @@ SyntacticElementTarget::walk(ASTWalker &walker) const {
       }
 
       PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-        // Ignore where clause if needed.
-        if (Target.ignoreForEachWhereClause() && E == ForStmt->getWhere())
+        if (E == ForStmt->getWhere())
           return Action::SkipNode(E);
 
         E = E->walk(Walker);
@@ -458,7 +456,7 @@ SyntacticElementTarget::walk(ASTWalker &walker) const {
     ForEachWalker forEachWalker(walker, *this);
 
     if (auto *newStmt = getAsForEachStmt()->walk(forEachWalker)) {
-      result.forEachStmt.stmt = cast<ForEachStmt>(newStmt);
+      result.forEachPreamble.stmt = cast<ForEachStmt>(newStmt);
     } else {
       return std::nullopt;
     }
diff --git a/lib/Sema/TypeCheckConstraints.cpp b/lib/Sema/TypeCheckConstraints.cpp
@@ -890,11 +890,26 @@ bool TypeChecker::typeCheckForEachPreamble(DeclContext *dc, ForEachStmt *stmt,
     return true;
   };
 
-  auto target = SyntacticElementTarget::forForEachPreamble(
-      stmt, dc, /*ignoreWhereClause=*/false, packElementEnv);
+  auto target =
+      SyntacticElementTarget::forForEachPreamble(stmt, dc, packElementEnv);
   if (!typeCheckTarget(target))
     return failed();
 
+  if (auto *where = stmt->getWhere()) {
+    auto boolType = dc->getASTContext().getBoolType();
+    if (!boolType)
+      return failed();
+
+    SyntacticElementTarget whereClause(stmt->getWhere(), dc,
+                                       {boolType, CTP_Condition},
+                                       /*isDiscarded=*/false);
+    auto result = typeCheckTarget(whereClause);
+    if (!result)
+      return true;
+
+    stmt->setWhere(result->getAsExpr());
+  }
+
   // Check to see if the sequence expr is throwing (in async context),
   // if so require the stmt to have a `try`.
   if (diagnoseUnhandledThrowsInAsyncContext(dc, stmt))
