[TypeChecker] Simplify for ... in ... type checking

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

@@ -64,8 +64,7 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
     case ClosureResult:
     case ParentType:
     case InstanceType:
-    case SequenceIteratorProtocol:
-    case GeneratorElementType:
+    case SequenceElementType:
     case AutoclosureResult:
     case GenericArgument:
     case NamedTupleElement:
@@ -176,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";
@@ -257,8 +262,8 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) {
       out << "function result";
       break;
 
-    case GeneratorElementType:
-      out << "generator element type";
+    case SequenceElementType:
+      out << "sequence element type";
       break;
 
     case GenericArgument:
@@ -301,10 +306,6 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) {
       out << "rvalue adjustment";
       break;
 
-    case SequenceIteratorProtocol:
-      out << "sequence iterator type";
-      break;
-
     case SubscriptMember:
       out << "subscript member";
       break;

lib/Sema/ConstraintLocator.h

@@ -98,10 +98,8 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     ParentType,
     /// The instance of a metatype type.
     InstanceType,
-    /// The generic type of a sequence.
-    SequenceIteratorProtocol,
-    /// The element type of a generator.
-    GeneratorElementType,
+    /// The element type of a sequence in a for ... in ... loop.
+    SequenceElementType,
     /// An argument passed in an autoclosure parameter
     /// position, which must match the autoclosure return type.
     AutoclosureResult,
@@ -161,8 +159,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ClosureResult:
     case ParentType:
     case InstanceType:
-    case SequenceIteratorProtocol:
-    case GeneratorElementType:
+    case SequenceElementType:
     case AutoclosureResult:
     case Requirement:
     case Witness:
@@ -211,8 +208,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ApplyArgument:
     case ApplyFunction:
     case ApplyArgToParam:
-    case SequenceIteratorProtocol:
-    case GeneratorElementType:
+    case SequenceElementType:
     case ClosureResult:
     case ConstructorMember:
     case InstanceType:
@@ -571,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);

lib/Sema/ConstraintSystem.cpp

@@ -2194,6 +2194,8 @@ Type simplifyTypeImpl(ConstraintSystem &cs, Type type, Fn getFixedTypeFn) {
       // FIXME: It's kind of weird in general that we have to look
       // through lvalue, inout and IUO types here
       Type lookupBaseType = newBase->getWithoutSpecifierType();
+      if (auto selfType = lookupBaseType->getAs<DynamicSelfType>())
+        lookupBaseType = selfType->getSelfType();
 
       if (lookupBaseType->mayHaveMembers() ||
           lookupBaseType->is<DynamicSelfType>()) {

lib/Sema/TypeCheckConstraints.cpp

@@ -2881,85 +2881,29 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
         return true;
       }
 
+      auto elementAssocType =
+        cast<AssociatedTypeDecl>(
+          sequenceProto->lookupDirect(tc.Context.Id_Element).front());
+
       SequenceType = cs.createTypeVariable(Locator, TVO_CanBindToNoEscape);
       cs.addConstraint(ConstraintKind::Conversion, cs.getType(expr),
                        SequenceType, Locator);
       cs.addConstraint(ConstraintKind::ConformsTo, SequenceType,
                        sequenceProto->getDeclaredType(), Locator);
 
-      auto iteratorLocator =
-        cs.getConstraintLocator(Locator,
-                                ConstraintLocator::SequenceIteratorProtocol);
       auto elementLocator =
-        cs.getConstraintLocator(iteratorLocator,
-                                ConstraintLocator::GeneratorElementType);
+        cs.getConstraintLocator(Locator,
+                                ConstraintLocator::SequenceElementType);
 
       // Collect constraints from the element pattern.
       auto pattern = Stmt->getPattern();
       InitType = cs.generateConstraints(pattern, elementLocator);
       if (!InitType)
         return true;
-
-      // Manually search for the iterator witness. If no iterator/element pair
-      // exists, solve for them.
-      Type iteratorType;
-      Type elementType;
-
-      NameLookupOptions lookupOptions = defaultMemberTypeLookupOptions;
-      if (isa<AbstractFunctionDecl>(cs.DC))
-        lookupOptions |= NameLookupFlags::KnownPrivate;
-
-      auto sequenceType = cs.getType(expr)->getRValueType();
-
-      // Look through one level of optional; this improves recovery but doesn't
-      // change the result.
-      if (auto sequenceObjectType = sequenceType->getOptionalObjectType())
-        sequenceType = sequenceObjectType;
-
-      // If the sequence type is an existential, we should not attempt to
-      // look up the member type at all, since we cannot represent associated
-      // types of existentials.
-      //
-      // We will diagnose it later.
-      if (!sequenceType->isExistentialType() &&
-          (sequenceType->mayHaveMembers() ||
-           sequenceType->isTypeVariableOrMember())) {
-        ASTContext &ctx = tc.Context;
-        auto iteratorAssocType =
-          cast<AssociatedTypeDecl>(
-            sequenceProto->lookupDirect(ctx.Id_Iterator).front());
-
-        auto subs = sequenceType->getContextSubstitutionMap(
-          cs.DC->getParentModule(),
-          sequenceProto);
-        iteratorType = iteratorAssocType->getDeclaredInterfaceType()
-          .subst(subs);
-
-        if (iteratorType) {
-          auto iteratorProto =
-            tc.getProtocol(Stmt->getForLoc(),
-                           KnownProtocolKind::IteratorProtocol);
-          if (!iteratorProto)
-            return true;
-
-          auto elementAssocType =
-            cast<AssociatedTypeDecl>(
-              iteratorProto->lookupDirect(ctx.Id_Element).front());
-
-          elementType = iteratorType->getTypeOfMember(
-                          cs.DC->getParentModule(),
-                          elementAssocType,
-                          elementAssocType->getDeclaredInterfaceType());
-        }
-      }
-
-      if (elementType.isNull()) {
-        elementType = cs.createTypeVariable(elementLocator,
-                                            TVO_CanBindToNoEscape);
-      }
 
       // Add a conversion constraint between the element type of the sequence
       // and the type of the element pattern.
+      auto elementType = DependentMemberType::get(SequenceType, elementAssocType);
       cs.addConstraint(ConstraintKind::Conversion, elementType, InitType,
                        elementLocator);
 

test/decl/func/dynamic_self.swift

@@ -419,3 +419,20 @@ func useSelfOperator() {
   let s = SelfOperator()
   _ = s + s
 }
+
+// for ... in loops
+
+struct DummyIterator : IteratorProtocol {
+  func next() -> Int? { return nil }
+}
+
+class Iterable : Sequence {
+  func returnsSelf() -> Self {
+    for _ in self {}
+    return self
+  }
+
+  func makeIterator() -> DummyIterator {
+    return DummyIterator()
+  }
+}

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 {}
 }