Merge pull request #68787 from xedin/separate-solving-for-for-loop-condition

[ConstraintSystem] Solve `where` clauses of for-in loops separately

Author: xedin

include/swift/Sema/SyntacticElementTarget.h

@@ -147,7 +147,7 @@ class SyntacticElementTarget {
       ForEachStmt *stmt;
       DeclContext *dc;
       Pattern *pattern;
-      bool bindPatternVarsOneWay;
+      bool ignoreWhereClause;
       ForEachStmtInfo info;
     } forEachStmt;
 
@@ -227,11 +227,11 @@ class SyntacticElementTarget {
   }
 
   SyntacticElementTarget(ForEachStmt *stmt, DeclContext *dc,
-                         bool bindPatternVarsOneWay)
+                         bool ignoreWhereClause)
       : kind(Kind::forEachStmt) {
     forEachStmt.stmt = stmt;
     forEachStmt.dc = dc;
-    forEachStmt.bindPatternVarsOneWay = bindPatternVarsOneWay;
+    forEachStmt.ignoreWhereClause = ignoreWhereClause;
   }
 
   /// Form a target for the initialization of a pattern from an expression.
@@ -249,7 +249,7 @@ class SyntacticElementTarget {
   /// Form a target for a for-in loop.
   static SyntacticElementTarget forForEachStmt(ForEachStmt *stmt,
                                                DeclContext *dc,
-                                               bool bindPatternVarsOneWay);
+                                               bool ignoreWhereClause = false);
 
   /// Form a target for a property with an attached property wrapper that is
   /// initialized out-of-line.
@@ -469,10 +469,6 @@ class SyntacticElementTarget {
   bool shouldBindPatternVarsOneWay() const {
     if (kind == Kind::expression)
       return expression.bindPatternVarsOneWay;
-
-    if (kind == Kind::forEachStmt)
-      return forEachStmt.bindPatternVarsOneWay;
-
     return false;
   }
 
@@ -521,6 +517,11 @@ class SyntacticElementTarget {
     return expression.initialization.patternBindingIndex;
   }
 
+  bool ignoreForEachWhereClause() const {
+    assert(isForEachStmt());
+    return forEachStmt.ignoreWhereClause;
+  }
+
   const ForEachStmtInfo &getForEachStmtInfo() const {
     assert(isForEachStmt());
     return forEachStmt.info;

lib/Sema/CSGen.cpp

@@ -4579,22 +4579,16 @@ generateForEachStmtConstraints(ConstraintSystem &cs,
                                            /*flags=*/0);
   {
     auto nextType = cs.getType(forEachStmtInfo.nextCall);
-    // Note that `OptionalObject` is not used here. This is due to inference
-    // behavior where it would bind `elementType` to the `initType` before
-    // resolving `optional object` constraint which is sometimes too eager.
-    cs.addConstraint(ConstraintKind::Conversion, nextType,
-                     OptionalType::get(elementType), elementTypeLoc);
+    cs.addConstraint(ConstraintKind::OptionalObject, nextType, elementType,
+                     elementTypeLoc);
     cs.addConstraint(ConstraintKind::Conversion, elementType, initType,
                      elementLocator);
   }
 
   // Generate constraints for the "where" expression, if there is one.
-  if (auto *whereExpr = stmt->getWhere()) {
-    auto *boolDecl = dc->getASTContext().getBoolDecl();
-    if (!boolDecl)
-      return llvm::None;
-
-    Type boolType = boolDecl->getDeclaredInterfaceType();
+  auto *whereExpr = stmt->getWhere();
+  if (whereExpr && !target.ignoreForEachWhereClause()) {
+    Type boolType = dc->getASTContext().getBoolType();
     if (!boolType)
       return llvm::None;
 

lib/Sema/CSSimplify.cpp

@@ -6303,10 +6303,22 @@ bool ConstraintSystem::repairFailures(
   }
 
   case ConstraintLocator::OptionalPayload: {
+    if (lhs->isPlaceholder() || rhs->isPlaceholder())
+      return true;
+
     if (repairViaOptionalUnwrap(*this, lhs, rhs, matchKind, conversionsOrFixes,
                                 locator))
       return true;
 
+    if (path.size() > 1) {
+      path.pop_back();
+      if (path.back().is<LocatorPathElt::SequenceElementType>()) {
+        conversionsOrFixes.push_back(
+            CollectionElementContextualMismatch::create(
+                *this, lhs, rhs, getConstraintLocator(anchor, path)));
+        return true;
+      }
+    }
     break;
   }
 

lib/Sema/CSSyntacticElement.cpp

@@ -629,9 +629,13 @@ 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::forForEachStmt(
         forEachStmt, context.getAsDeclContext(),
-        /*bindTypeVarsOneWay=*/false);
+        /*ignoreWhereClause=*/true);
 
     if (cs.generateConstraints(target)) {
       hadError = true;
@@ -961,10 +965,24 @@ class SyntacticElementConstraintGenerator
 
     // For-each pattern.
     //
-    // Note that we don't record a sequence or where clause here,
-    // they would be handled together with pattern because pattern can
-    // inform a type of sequence element e.g. `for i: Int8 in 0 ..< 8`
+    // Note that we don't record a sequence here, it would be handled
+    // together with pattern because pattern can inform a type of sequence
+    // element e.g. `for i: Int8 in 0 ..< 8`
     elements.push_back(makeElement(forEachStmt->getPattern(), stmtLoc));
+
+    // Where clause if any.
+    if (auto *where = forEachStmt->getWhere()) {
+      Type boolType = cs.getASTContext().getBoolType();
+      if (!boolType) {
+        hadError = true;
+        return;
+      }
+
+      ContextualTypeInfo context(boolType, CTP_Condition);
+      elements.push_back(
+          makeElement(where, stmtLoc, context, /*isDiscarded=*/false));
+    }
+
     // Body of the `for-in` loop.
     elements.push_back(makeElement(forEachStmt->getBody(), stmtLoc));
 

lib/Sema/SyntacticElementTarget.cpp

@@ -178,9 +178,8 @@ SyntacticElementTarget SyntacticElementTarget::forInitialization(
 
 SyntacticElementTarget
 SyntacticElementTarget::forForEachStmt(ForEachStmt *stmt, DeclContext *dc,
-                                       bool bindPatternVarsOneWay) {
-  SyntacticElementTarget target(
-      stmt, dc, bindPatternVarsOneWay || bool(stmt->getWhere()));
+                                       bool ignoreWhereClause) {
+  SyntacticElementTarget target(stmt, dc, ignoreWhereClause);
   return target;
 }
 

lib/Sema/TypeCheckConstraints.cpp

@@ -901,8 +901,7 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
     return true;
   };
 
-  auto target = SyntacticElementTarget::forForEachStmt(
-      stmt, dc, /*bindPatternVarsOneWay=*/false);
+  auto target = SyntacticElementTarget::forForEachStmt(stmt, dc);
   if (!typeCheckTarget(target))
     return failed();
 

test/Constraints/rdar107651291.swift

@@ -7,6 +7,6 @@ func foo(xs: [String: [String]], ys: [String: [String]]) {
     for (a, b) in zip(xs, ys) {}
     // expected-error@-1 {{type 'Dictionary<String, [String]>.Element' (aka '(key: String, value: Array<String>)') cannot conform to 'Sequence'}}
     // expected-note@-2 {{only concrete types such as structs, enums and classes can conform to protocols}}
-    // expected-note@-3 {{required by referencing instance method 'next()'}}
+    // expected-note@-3 {{required by global function 'zip' where 'Sequence2' = 'Dictionary<String, [String]>.Element' (aka '(key: String, value: Array<String>)')}}
   }
 }

test/stmt/foreach.swift

@@ -273,3 +273,25 @@ do {
   // https://github.com/apple/swift/issues/65650 - Make sure we say 'String', not 'Any'.
   for (x, y) in [""] {} // expected-error {{tuple pattern cannot match values of non-tuple type 'String'}}
 }
+
+do {
+  class Base : Hashable {
+    static func ==(_: Base, _: Base) -> Bool { false }
+
+    func hash(into hasher: inout Hasher) {}
+  }
+
+  class Child : Base {
+    var value: Int = 0
+  }
+
+  struct Range {
+    func contains(_: Base) -> Bool { return false }
+  }
+
+  func test(data: Set<Child>, _ range: Range) {
+    for v in data where range.contains(v) {
+      _ = v.value // Ok (`v` is inferred from `data` and not from `range`)
+    }
+  }
+}