[ConstraintSystem] Use witnesses for makeIterator and next refs in for-in context

include/swift/AST/ASTContext.h

@@ -602,6 +602,12 @@ class ASTContext final {
   /// Get AsyncSequence.makeAsyncIterator().
   FuncDecl *getAsyncSequenceMakeAsyncIterator() const;
 
+  /// Get IteratorProtocol.next().
+  FuncDecl *getIteratorNext() const;
+
+  /// Get AsyncIteratorProtocol.next().
+  FuncDecl *getAsyncIteratorNext() const;
+
   /// Check whether the standard library provides all the correct
   /// intrinsic support for Optional<T>.
   ///

include/swift/Sema/Constraint.h

@@ -137,6 +137,11 @@ enum class ConstraintKind : char {
   /// name, and the type of that member, when referenced as a value, is the
   /// second type.
   UnresolvedValueMember,
+  /// The first type conforms to the protocol in which the member requirement
+  /// resides. Once the conformance is resolved, the value witness will be
+  /// determined, and the type of that witness, when referenced as a value,
+  /// will be bound to the second type.
+  ValueWitness,
   /// The first type can be defaulted to the second (which currently
   /// cannot be dependent).  This is more like a type property than a
   /// relational constraint.
@@ -406,11 +411,18 @@ class Constraint final : public llvm::ilist_node<Constraint>,
       /// The type of the member.
       Type Second;
 
-      /// If non-null, the name of a member of the first type is that
-      /// being related to the second type.
-      ///
-      /// Used for ValueMember an UnresolvedValueMember constraints.
-      DeclNameRef Name;
+      union {
+        /// If non-null, the name of a member of the first type is that
+        /// being related to the second type.
+        ///
+        /// Used for ValueMember an UnresolvedValueMember constraints.
+        DeclNameRef Name;
+
+        /// If non-null, the member being referenced.
+        ///
+        /// Used for ValueWitness constraints.
+        ValueDecl *Ref;
+      } Member;
 
       /// The DC in which the use appears.
       DeclContext *UseDC;
@@ -525,6 +537,12 @@ class Constraint final : public llvm::ilist_node<Constraint>,
                                   FunctionRefKind functionRefKind,
                                   ConstraintLocator *locator);
 
+  /// Create a new value witness constraint.
+  static Constraint *createValueWitness(
+      ConstraintSystem &cs, ConstraintKind kind, Type first, Type second,
+      ValueDecl *requirement, DeclContext *useDC,
+      FunctionRefKind functionRefKind, ConstraintLocator *locator);
+
   /// Create an overload-binding constraint.
   static Constraint *createBindOverload(ConstraintSystem &cs, Type type, 
                                         OverloadChoice choice, 
@@ -672,6 +690,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
     case ConstraintKind::ValueMember:
     case ConstraintKind::UnresolvedValueMember:
+    case ConstraintKind::ValueWitness:
     case ConstraintKind::PropertyWrapper:
       return ConstraintClassification::Member;
 
@@ -711,6 +730,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
     case ConstraintKind::ValueMember:
     case ConstraintKind::UnresolvedValueMember:
+    case ConstraintKind::ValueWitness:
       return Member.First;
 
     case ConstraintKind::SyntacticElement:
@@ -732,6 +752,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
     case ConstraintKind::ValueMember:
     case ConstraintKind::UnresolvedValueMember:
+    case ConstraintKind::ValueWitness:
       return Member.Second;
 
     default:
@@ -757,13 +778,20 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   DeclNameRef getMember() const {
     assert(Kind == ConstraintKind::ValueMember ||
            Kind == ConstraintKind::UnresolvedValueMember);
-    return Member.Name;
+    return Member.Member.Name;
+  }
+
+  /// Retrieve the requirement being referenced by a value witness constraint.
+  ValueDecl *getRequirement() const {
+    assert(Kind == ConstraintKind::ValueWitness);
+    return Member.Member.Ref;
   }
 
   /// Determine the kind of function reference we have for a member reference.
   FunctionRefKind getFunctionRefKind() const {
     if (Kind == ConstraintKind::ValueMember ||
-        Kind == ConstraintKind::UnresolvedValueMember)
+        Kind == ConstraintKind::UnresolvedValueMember ||
+        Kind == ConstraintKind::ValueWitness)
       return static_cast<FunctionRefKind>(TheFunctionRefKind);
 
     // Conservative answer: drop all of the labels.
@@ -823,7 +851,8 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   /// Retrieve the DC in which the member was used.
   DeclContext *getMemberUseDC() const {
     assert(Kind == ConstraintKind::ValueMember ||
-           Kind == ConstraintKind::UnresolvedValueMember);
+           Kind == ConstraintKind::UnresolvedValueMember ||
+           Kind == ConstraintKind::ValueWitness);
     return Member.UseDC;
   }
 

include/swift/Sema/ConstraintSystem.h

@@ -4354,6 +4354,26 @@ class ConstraintSystem {
     }
   }
 
+  /// Add a value witness constraint to the constraint system.
+  void addValueWitnessConstraint(
+      Type baseTy, ValueDecl *requirement, Type memberTy, DeclContext *useDC,
+      FunctionRefKind functionRefKind, ConstraintLocatorBuilder locator) {
+    assert(baseTy);
+    assert(memberTy);
+    assert(requirement);
+    assert(useDC);
+    switch (simplifyValueWitnessConstraint(
+        ConstraintKind::ValueWitness, baseTy, requirement, memberTy, useDC,
+        functionRefKind, TMF_GenerateConstraints, locator)) {
+    case SolutionKind::Unsolved:
+      llvm_unreachable("Unsolved result when generating constraints!");
+
+    case SolutionKind::Solved:
+    case SolutionKind::Error:
+      break;
+    }
+  }
+
   /// Add an explicit conversion constraint (e.g., \c 'x as T').
   ///
   /// \param fromType The type of the expression being converted.

lib/AST/ASTContext.cpp

@@ -217,6 +217,12 @@ struct ASTContext::Implementation {
   /// The declaration of 'AsyncSequence.makeAsyncIterator()'.
   FuncDecl *MakeAsyncIterator = nullptr;
 
+  /// The declaration of 'IteratorProtocol.next()'.
+  FuncDecl *IteratorNext = nullptr;
+
+  /// The declaration of 'AsyncIteratorProtocol.next()'.
+  FuncDecl *AsyncIteratorNext = nullptr;
+
   /// The declaration of Swift.Optional<T>.Some.
   EnumElementDecl *OptionalSomeDecl = nullptr;
 
@@ -806,31 +812,40 @@ FuncDecl *ASTContext::getPlusFunctionOnString() const {
   return getImpl().PlusFunctionOnString;
 }
 
-FuncDecl *ASTContext::getSequenceMakeIterator() const {
-  if (getImpl().MakeIterator) {
-    return getImpl().MakeIterator;
-  }
-
-  auto proto = getProtocol(KnownProtocolKind::Sequence);
-  if (!proto)
-    return nullptr;
-
-  for (auto result : proto->lookupDirect(Id_makeIterator)) {
+static FuncDecl *lookupRequirement(ProtocolDecl *proto,
+                                   Identifier requirement) {
+  for (auto result : proto->lookupDirect(requirement)) {
     if (result->getDeclContext() != proto)
       continue;
 
     if (auto func = dyn_cast<FuncDecl>(result)) {
       if (func->getParameters()->size() != 0)
         continue;
 
-      getImpl().MakeIterator = func;
       return func;
     }
   }
 
   return nullptr;
 }
 
+FuncDecl *ASTContext::getSequenceMakeIterator() const {
+  if (getImpl().MakeIterator) {
+    return getImpl().MakeIterator;
+  }
+
+  auto proto = getProtocol(KnownProtocolKind::Sequence);
+  if (!proto)
+    return nullptr;
+
+  if (auto *func = lookupRequirement(proto, Id_makeIterator)) {
+    getImpl().MakeIterator = func;
+    return func;
+  }
+
+  return nullptr;
+}
+
 FuncDecl *ASTContext::getAsyncSequenceMakeAsyncIterator() const {
   if (getImpl().MakeAsyncIterator) {
     return getImpl().MakeAsyncIterator;
@@ -840,17 +855,43 @@ FuncDecl *ASTContext::getAsyncSequenceMakeAsyncIterator() const {
   if (!proto)
     return nullptr;
 
-  for (auto result : proto->lookupDirect(Id_makeAsyncIterator)) {
-    if (result->getDeclContext() != proto)
-      continue;
+  if (auto *func = lookupRequirement(proto, Id_makeAsyncIterator)) {
+    getImpl().MakeAsyncIterator = func;
+    return func;
+  }
 
-    if (auto func = dyn_cast<FuncDecl>(result)) {
-      if (func->getParameters()->size() != 0)
-        continue;
+  return nullptr;
+}
 
-      getImpl().MakeAsyncIterator = func;
-      return func;
-    }
+FuncDecl *ASTContext::getIteratorNext() const {
+  if (getImpl().IteratorNext) {
+    return getImpl().IteratorNext;
+  }
+
+  auto proto = getProtocol(KnownProtocolKind::IteratorProtocol);
+  if (!proto)
+    return nullptr;
+
+  if (auto *func = lookupRequirement(proto, Id_next)) {
+    getImpl().IteratorNext = func;
+    return func;
+  }
+
+  return nullptr;
+}
+
+FuncDecl *ASTContext::getAsyncIteratorNext() const {
+  if (getImpl().AsyncIteratorNext) {
+    return getImpl().AsyncIteratorNext;
+  }
+
+  auto proto = getProtocol(KnownProtocolKind::AsyncIteratorProtocol);
+  if (!proto)
+    return nullptr;
+
+  if (auto *func = lookupRequirement(proto, Id_next)) {
+    getImpl().AsyncIteratorNext = func;
+    return func;
   }
 
   return nullptr;

lib/Sema/CSBindings.cpp

@@ -1473,6 +1473,7 @@ void PotentialBindings::infer(Constraint *constraint) {
 
   case ConstraintKind::ValueMember:
   case ConstraintKind::UnresolvedValueMember:
+  case ConstraintKind::ValueWitness:
   case ConstraintKind::PropertyWrapper: {
     // If current type variable represents a member type of some reference,
     // it would be bound once member is resolved either to a actual member

lib/Sema/CSGen.cpp

@@ -4057,6 +4057,21 @@ generateForEachStmtConstraints(
           AwaitExpr::createImplicit(ctx, /*awaitLoc=*/SourceLoc(), nextCall);
     }
 
+    // The iterator type must conform to IteratorProtocol.
+    {
+      ProtocolDecl *iteratorProto = TypeChecker::getProtocol(
+          cs.getASTContext(), stmt->getForLoc(),
+          isAsync ? KnownProtocolKind::AsyncIteratorProtocol
+                  : KnownProtocolKind::IteratorProtocol);
+      if (!iteratorProto)
+        return None;
+
+      cs.setContextualType(
+          nextRef->getBase(),
+          TypeLoc::withoutLoc(iteratorProto->getDeclaredInterfaceType()),
+          CTP_ForEachSequence);
+    }
+
     SolutionApplicationTarget nextTarget(nextCall, dc, CTP_Unused,
                                          /*contextualType=*/Type(),
                                          /*isDiscarded=*/false);