[CSSimplify] Member ref decays into value witness for makeIterator in for-in

Reference to `makeIterator` in for-in loop context needs to be
treated as a reference to a witness of `Sequence#makeIterator`
requirement, otherwise it could lead to problems with retroactive
conformances.

Author: xedin

lib/Sema/CSSimplify.cpp

@@ -9260,6 +9260,40 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
     }
   }
 
+  // Special handling of injected references to `makeIterator` in for-in loops.
+  {
+    auto memberRef = getAsExpr<UnresolvedDotExpr>(locator->getAnchor());
+    if (memberRef && memberRef->isImplicit() &&
+        locator->isLastElement<LocatorPathElt::Member>()) {
+      // Cannot simplify this constraint yet since we don't know whether
+      // the base type is going to be existential or not.
+      if (baseObjTy->isTypeVariableOrMember())
+        return formUnsolved();
+
+      auto *sequenceExpr = memberRef->getBase();
+      // If base type is an existential, member lookup is fine because
+      // it would return a witness.
+      if (!baseObjTy->isExistentialType() &&
+          getContextualTypePurpose(sequenceExpr) == CTP_ForEachSequence) {
+        auto &ctx = getASTContext();
+
+        auto *sequenceProto = cast<ProtocolDecl>(
+            getContextualType(sequenceExpr, /*forConstraint=*/false)
+                ->getAnyNominal());
+        bool isAsync = sequenceProto ==
+                       TypeChecker::getProtocol(
+                           ctx, SourceLoc(), KnownProtocolKind::AsyncSequence);
+
+        auto *makeIterator = isAsync ? ctx.getAsyncSequenceMakeAsyncIterator()
+                                     : ctx.getSequenceMakeIterator();
+
+        return simplifyValueWitnessConstraint(
+            ConstraintKind::ValueWitness, baseTy, makeIterator, memberTy, DC,
+            FunctionRefKind::Compound, flags, locator);
+      }
+    }
+  }
+
   MemberLookupResult result =
       performMemberLookup(kind, member, baseTy, functionRefKind, locator,
                           /*includeInaccessibleMembers*/ shouldAttemptFixes());
@@ -9696,6 +9730,17 @@ ConstraintSystem::simplifyValueWitnessConstraint(
     return SolutionKind::Unsolved;
   };
 
+  auto fail = [&] {
+    // The constraint failed, so mark the member type as a "hole".
+    // We cannot do anything further here.
+    if (!shouldAttemptFixes())
+      return SolutionKind::Error;
+
+    recordAnyTypeVarAsPotentialHole(memberType);
+
+    return SolutionKind::Solved;
+  };
+
   // Resolve the base type, if we can. If we can't resolve the base type,
   // then we can't solve this constraint.
   Type baseObjectType = getFixedTypeRecursive(
@@ -9712,16 +9757,8 @@ ConstraintSystem::simplifyValueWitnessConstraint(
   assert(proto && "Value witness constraint for a non-requirement");
   auto conformance = useDC->getParentModule()->lookupConformance(
       baseObjectType, proto);
-  if (!conformance) {
-    // The conformance failed, so mark the member type as a "hole". We cannot
-    // do anything further here.
-    if (!shouldAttemptFixes())
-      return SolutionKind::Error;
-
-    recordAnyTypeVarAsPotentialHole(memberType);
-
-    return SolutionKind::Solved;
-  }
+  if (!conformance)
+    return fail();
 
   // Reference the requirement.
   Type resolvedBaseType = simplifyType(baseType, flags);
@@ -9731,7 +9768,7 @@ ConstraintSystem::simplifyValueWitnessConstraint(
   auto witness =
       conformance.getWitnessByName(baseObjectType, requirement->getName());
   if (!witness)
-    return SolutionKind::Error;
+    return fail();
 
   auto choice = OverloadChoice(resolvedBaseType, witness.getDecl(), functionRefKind);
   resolveOverload(getConstraintLocator(locator), memberType, choice,

validation-test/Sema/issue-60514.swift

@@ -0,0 +1,54 @@
+// RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s
+
+// https://github.com/apple/swift/issues/60514
+// Make sure that `makeIterator` witness is picked over the non-witness.
+
+public protocol VectorSectionReader: Sequence where Element == Result<Item, Error> {
+  associatedtype Item
+  var count: UInt32 { get }
+  mutating func read() throws -> Item
+}
+
+public struct VectorSectionIterator<Reader: VectorSectionReader>: IteratorProtocol {
+  private(set) var reader: Reader
+  private(set) var left: UInt32
+
+  init(reader: Reader, count: UInt32) {
+      self.reader = reader
+      self.left = count
+  }
+
+  private var end: Bool = false
+  public mutating func next() -> Reader.Element? {
+    guard !end else { return nil }
+    guard left != 0 else { return nil }
+    let result = Result(catching: { try reader.read() })
+    left -= 1
+    switch result {
+    case .success: return result
+    case .failure:
+      end = true
+      return result
+    }
+  }
+}
+
+extension VectorSectionReader {
+  __consuming public func makeIterator() -> VectorSectionIterator<Self> {
+    VectorSectionIterator(reader: self, count: count)
+  }
+
+  // CHECK: sil [ossa] @$s4main19VectorSectionReaderPAAE7collectSay4ItemQzGyKF
+  public func collect() throws -> [Item] {
+    var items: [Item] = []
+    items.reserveCapacity(Int(count))
+    for result in self {
+      // CHECK: [[ITERATOR:%.*]] = project_box {{.*}} : $<τ_0_0 where τ_0_0 : VectorSectionReader> { var τ_0_0.Iterator } <Self>, 0
+      // CHECK-NEXT: [[SELF:%.*]] = alloc_stack $Self
+      // CHECK: [[MAKE_ITERATOR_REF:%.*]] = witness_method $Self, #Sequence.makeIterator : <Self where Self : Sequence> (__owned Self) -> () -> Self.Iterator : $@convention(witness_method: Sequence) <τ_0_0 where τ_0_0 : Sequence> (@in τ_0_0) -> @out τ_0_0.Iterator
+      // CHECK-NEXT: apply [[MAKE_ITERATOR_REF]]<Self>([[ITERATOR]], [[SELF]]) : $@convention(witness_method: Sequence) <τ_0_0 where τ_0_0 : Sequence> (@in τ_0_0) -> @out τ_0_0.Iterator
+      try items.append(result.get())
+    }
+    return items
+  }
+}

validation-test/Sema/issue-60958.swift

@@ -0,0 +1,36 @@
+// RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s
+
+// https://github.com/apple/swift/issues/60958
+// Make sure that `makeIterator` witness is picked over the non-witness.
+
+struct S: Sequence {
+  private var _storage: [UInt8] = []
+
+  func makeIterator() -> Array<UInt8>.Iterator {
+    _storage.makeIterator()
+  }
+
+  typealias Element = UInt8
+
+  class Iterator: IteratorProtocol {
+    func next() -> UInt8? { 0 }
+    typealias Element = UInt8
+  }
+
+  func makeIterator() -> Iterator {
+    Iterator()
+  }
+}
+
+extension S {
+  // CHECK: sil hidden [ossa] @$s4main1SV1fyyF
+  func f() {
+    for elt in self {
+      // CHECK: [[ITERATOR_VAR:%.*]] = project_box {{.*}} : ${ var S.Iterator }, 0
+      // CHECK: [[MAKE_ITERATOR_REF:%.*]] = function_ref @$s4main1SV12makeIteratorAC0C0CyF : $@convention(method) (@guaranteed S) -> @owned S.Iterator
+      // CHECK-NEXT: [[ITERATOR:%.*]] = apply [[MAKE_ITERATOR_REF]](%0) : $@convention(method) (@guaranteed S) -> @owned S.Iterator
+      // CHECK-NEXT: store [[ITERATOR]] to [init] [[ITERATOR_VAR]] : $*S.Iterator
+      print(elt)
+    }
+  }
+}