[ConstraintSystem] Improve contextual mismatch diagnostics for for ... in <expr> loop

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -447,6 +447,13 @@ ERROR(cannot_convert_assign_protocol,none,
 ERROR(cannot_convert_assign_nil,none,
       "'nil' cannot be assigned to type %0", (Type))
 
+// for ... in expression
+ERROR(cannot_convert_sequence_element_value,none,
+      "cannot convert sequence element type %0 to expected type %1",
+      (Type, Type))
+ERROR(cannot_convert_sequence_element_protocol,none,
+      "sequence element type %0 does not conform to expected type %1",
+      (Type, Type))
 
 ERROR(throws_functiontype_mismatch,none,
       "invalid conversion from throwing function of type %0 to "

lib/Sema/CSDiagnostics.cpp

@@ -2847,6 +2847,7 @@ bool ExtraneousReturnFailure::diagnoseAsError() {
 
 bool CollectionElementContextualFailure::diagnoseAsError() {
   auto *anchor = getAnchor();
+  auto *locator = getLocator();
 
   auto eltType = getFromType();
   auto contextualType = getToType();
@@ -2859,8 +2860,7 @@ bool CollectionElementContextualFailure::diagnoseAsError() {
   }
 
   if (isa<DictionaryExpr>(getRawAnchor())) {
-    auto *locator = getLocator();
-    const auto eltLoc = locator->getPath().back();
+    const auto &eltLoc = locator->getPath().back();
 
     switch (eltLoc.getValue()) {
     case 0: // key
@@ -2880,6 +2880,15 @@ bool CollectionElementContextualFailure::diagnoseAsError() {
     }
   }
 
+  if (locator->isForSequenceElementType()) {
+    diagnostic.emplace(
+        emitDiagnostic(anchor->getLoc(),
+                       contextualType->isExistentialType()
+                           ? diag::cannot_convert_sequence_element_protocol
+                           : diag::cannot_convert_sequence_element_value,
+                       eltType, contextualType));
+  }
+
   if (!diagnostic)
     return false;
 

lib/Sema/CSSimplify.cpp

@@ -2193,6 +2193,12 @@ bool ConstraintSystem::repairFailures(
     break;
   }
 
+  case ConstraintLocator::SequenceElementType: {
+    conversionsOrFixes.push_back(CollectionElementContextualMismatch::create(
+        *this, lhs, rhs, getConstraintLocator(locator)));
+    break;
+  }
+
   default:
     break;
   }

lib/Sema/ConstraintLocator.cpp

@@ -175,6 +175,12 @@ bool ConstraintLocator::isForGenericParameter() const {
          path.back().getKind() == ConstraintLocator::GenericParameter;
 }
 
+bool ConstraintLocator::isForSequenceElementType() const {
+  auto path = getPath();
+  return !path.empty() &&
+         path.back().getKind() == ConstraintLocator::SequenceElementType;
+}
+
 void ConstraintLocator::dump(SourceManager *sm) {
   dump(sm, llvm::errs());
   llvm::errs() << "\n";

lib/Sema/ConstraintLocator.h

@@ -567,6 +567,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// Determine whether this locator points to the generic parameter.
   bool isForGenericParameter() const;
 
+  /// Determine whether this locator points to the element type of a
+  /// sequence in a for ... in ... loop.
+  bool isForSequenceElementType() const;
+
   /// Produce a profile of this locator, for use in a folding set.
   static void Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
                       ArrayRef<PathElement> path);

test/stmt/foreach.swift

@@ -53,12 +53,15 @@ struct GoodTupleIterator: Sequence, IteratorProtocol {
   func makeIterator() -> GoodTupleIterator {}
 }
 
+protocol ElementProtocol {}
+
 func patterns(gir: GoodRange<Int>, gtr: GoodTupleIterator) {
   var sum : Int
   var sumf : Float
   for i : Int in gir { sum = sum + i }
   for i in gir { sum = sum + i }
-  for f : Float in gir { sum = sum + f } // expected-error{{'Int' is not convertible to 'Float'}}
+  for f : Float in gir { sum = sum + f } // expected-error{{cannot convert sequence element type 'Int' to expected type 'Float'}}
+  for f : ElementProtocol in gir { } // expected-error {{sequence element type 'Int' does not conform to expected type 'ElementProtocol'}}
 
   for (i, f) : (Int, Float) in gtr { sum = sum + i }
 
@@ -70,7 +73,7 @@ func patterns(gir: GoodRange<Int>, gtr: GoodTupleIterator) {
 
   for (i, _) : (Int, Float) in gtr { sum = sum + i }
 
-  for (i, _) : (Int, Int) in gtr { sum = sum + i } // expected-error{{'GoodTupleIterator.Element' (aka '(Int, Float)') is not convertible to '(Int, Int)'}}
+  for (i, _) : (Int, Int) in gtr { sum = sum + i } // expected-error{{cannot convert sequence element type 'GoodTupleIterator.Element' (aka '(Int, Float)') to expected type '(Int, Int)'}}
 
   for (i, f) in gtr {}
 }