[CSRanking] Adjust ranking of makeIterator overloads in for-in context

lib/Sema/CSRanking.cpp

@@ -875,6 +875,99 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
     auto decl2 = choice2.getDecl();
     auto dc2 = decl2->getDeclContext();
 
+    // Prefer an overload of `makeIterator` that satisfies
+    // `Sequence#makeIterator` requirement when `.makeIterator()`
+    // is implicitly injected after sequence expression in `for-in`
+    // loop context.
+    //
+    // This helps to avoid issues related to conditional conformances
+    // to `Sequence` that add `makeIterator` preferred by overloading
+    // rules that could change behavior i.e.:
+    //
+    // \code
+    // extension Array where Element == Double {
+    //   func makeIterator() -> Array<Int>.Iterator {
+    //     return [1, 2, 3].makeIterator()
+    //   }
+    // }
+    //
+    // for i in [4.0, 5.0, 6.0] {
+    //   print(i)
+    // }
+    // \endcode
+    //
+    // `makeIterator` declared in the `Array` extension is not a witness
+    // to `Sequence#makeIterator` and shouldn't be used because that would
+    // change runtime behavior.
+    if (auto *locator = overload.locator) {
+      auto memberRef = getAsExpr<UnresolvedDotExpr>(locator->getAnchor());
+      // `<sequence-expr>.makeIterator()` in `for-in` loop.
+      if (memberRef && memberRef->isImplicit() &&
+          locator->isLastElement<LocatorPathElt::Member>()) {
+        auto *sequenceExpr = memberRef->getBase();
+        if (cs.getContextualTypePurpose(sequenceExpr) == CTP_ForEachSequence) {
+          auto &ctx = cs.getASTContext();
+
+          assert(decl1->getBaseIdentifier() == ctx.Id_makeIterator &&
+                 decl2->getBaseIdentifier() == ctx.Id_makeIterator);
+
+          auto *sequenceProto = cast<ProtocolDecl>(
+              cs.getContextualType(sequenceExpr, /*forConstraint=*/false)
+                  ->getAnyNominal());
+
+          auto getWitnessFor = [&](Type type, ProtocolDecl *protocol,
+                                   DeclName requirement) -> ValueDecl * {
+            auto conformance =
+                cs.DC->getParentModule()->lookupConformance(type, protocol);
+
+            if (!conformance)
+              return nullptr;
+
+            return conformance.getWitnessByName(type, requirement).getDecl();
+          };
+
+          auto getSequenceExprType =
+              [&](const Solution &solution) -> Optional<Type> {
+            // Check whether given solution has a fixed type for
+            // the sequence expression, inner partial solutions
+            // would not have the expr type (because there is no
+            // constraints that bring that type into their scope),
+            // so this check could only be performed on "complete"
+            // solutions that have both overload choice and base type.
+            auto type = cs.getType(sequenceExpr);
+            if (auto *typeVar = type->getAs<TypeVariableType>()) {
+              if (!solution.hasFixedType(typeVar))
+                return None;
+            }
+            return solution.getResolvedType(sequenceExpr);
+          };
+
+          auto baseTy1 = getSequenceExprType(solutions[idx1]);
+          auto baseTy2 = getSequenceExprType(solutions[idx2]);
+
+          if (baseTy1 && baseTy2) {
+            auto witness1 =
+                getWitnessFor(*baseTy1, sequenceProto, decl1->getName());
+
+            auto witness2 =
+                getWitnessFor(*baseTy2, sequenceProto, decl2->getName());
+
+            if (decl1 == witness1) {
+              identical = false;
+              score1 += 2 * weight;
+            }
+
+            if (decl2 == witness2) {
+              identical = false;
+              score2 += 2 * weight;
+            }
+
+            continue;
+          }
+        }
+      }
+    }
+
     // The two systems are not identical. If the decls in question are distinct
     // protocol members, let the checks below determine if the two choices are
     // 'identical' or not. This allows us to structurally unify disparate

validation-test/Sema/issue60514.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] @$s10issue6051419VectorSectionReaderPAAE7collectSay4ItemQzGyKF
+  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/issue60958.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] @$s10issue609581SV1fyyF
+  func f() {
+    for elt in self {
+      // CHECK: [[ITERATOR_VAR:%.*]] = project_box {{.*}} : ${ var S.Iterator }, 0
+      // CHECK: [[MAKE_ITERATOR_REF:%.*]] = function_ref @$s10issue609581SV12makeIteratorAC0C0CyF : $@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)
+    }
+  }
+}