Sema: Two associated type inference fixes

lib/Sema/AssociatedTypeInference.cpp

@@ -1017,6 +1017,10 @@ class AssociatedTypeInference {
   std::pair<Type, TypeDecl *>
   computeDerivedTypeWitness(AssociatedTypeDecl *assocType);
 
+  /// See if we have a generic parameter named the same as this associated
+  /// type.
+  Type computeGenericParamWitness(AssociatedTypeDecl *assocType) const;
+
   /// Compute a type witness without using a specific potential witness.
   llvm::Optional<AbstractTypeWitness>
   computeAbstractTypeWitness(AssociatedTypeDecl *assocType);
@@ -2389,6 +2393,7 @@ Type AssociatedTypeInference::computeFixedTypeWitness(
     // any fix this associated type to a concrete type.
     for (auto conformance : getPeerConformances(conformance)) {
       auto *conformedProto = conformance->getProtocol();
+
       auto sig = conformedProto->getGenericSignature();
 
       // FIXME: The RequirementMachine will assert on re-entrant construction.
@@ -2656,6 +2661,28 @@ AssociatedTypeInference::computeAbstractTypeWitness(
   return llvm::None;
 }
 
+/// Look for a generic parameter that matches the name of the
+/// associated type.
+Type AssociatedTypeInference::computeGenericParamWitness(
+    AssociatedTypeDecl *assocType) const {
+  if (auto genericSig = dc->getGenericSignatureOfContext()) {
+    // Ignore the generic parameters for AsyncIteratorProtocol.Failure and
+    // AsyncSequence.Failure.
+    if (!isAsyncIteratorProtocolFailure(assocType)) {
+      for (auto *gp : genericSig.getInnermostGenericParams()) {
+        // Packs cannot witness associated type requirements.
+        if (gp->isParameterPack())
+          continue;
+
+        if (gp->getName() == assocType->getName())
+          return dc->mapTypeIntoContext(gp);
+      }
+    }
+  }
+
+  return Type();
+}
+
 void AssociatedTypeInference::collectAbstractTypeWitnesses(
     TypeWitnessSystem &system,
     ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes) const {
@@ -2704,35 +2731,15 @@ void AssociatedTypeInference::collectAbstractTypeWitnesses(
     if (system.hasResolvedTypeWitness(assocType->getName()))
       continue;
 
-    // If we find a default type definition, feed it to the system.
-    if (const auto &typeWitness = computeDefaultTypeWitness(assocType)) {
+    if (auto gpType = computeGenericParamWitness(assocType)) {
+      system.addTypeWitness(assocType->getName(), gpType, /*preferred=*/true);
+    } else if (const auto &typeWitness = computeDefaultTypeWitness(assocType)) {
       bool preferred = (typeWitness->getDefaultedAssocType()->getDeclContext()
                         == conformance->getProtocol());
       system.addDefaultTypeWitness(typeWitness->getType(),
                                    typeWitness->getDefaultedAssocType(),
                                    preferred);
-    } else {
-      // As a last resort, look for a generic parameter that matches the name
-      // of the associated type.
-      if (auto genericSig = dc->getGenericSignatureOfContext()) {
-        // Ignore the generic parameters for AsyncIteratorProtocol.Failure and
-        // AsyncSequence.Failure.
-        if (!isAsyncIteratorProtocolFailure(assocType)) {
-          for (auto *gp : genericSig.getInnermostGenericParams()) {
-            // Packs cannot witness associated type requirements.
-            if (gp->isParameterPack())
-              continue;
-
-            if (gp->getName() == assocType->getName()) {
-              system.addTypeWitness(assocType->getName(),
-                                    dc->mapTypeIntoContext(gp),
-                                    /*preferred=*/true);
-            }
-          }
-        }
-      }
     }
-
   }
 }
 
@@ -3150,8 +3157,15 @@ AssociatedTypeDecl *AssociatedTypeInference::inferAbstractTypeWitnesses(
 
     // If simplification failed, give up.
     if (type->hasTypeParameter()) {
-      LLVM_DEBUG(llvm::dbgs() << "-- Simplification failed: " << type << "\n");
-      return assocType;
+      if (auto gpType = computeGenericParamWitness(assocType)) {
+        LLVM_DEBUG(llvm::dbgs() << "-- Found generic parameter as last resort: "
+                                << gpType << "\n");
+        type = gpType;
+        typeWitnesses.insert(assocType, {type, reqDepth});
+      } else {
+        LLVM_DEBUG(llvm::dbgs() << "-- Simplification failed: " << type << "\n");
+        return assocType;
+      }
     }
 
     if (const auto failed =

test/decl/protocol/req/associated_type_generic_param.swift

@@ -1,11 +1,10 @@
 // RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
-// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+// RUN: not %target-typecheck-verify-swift -disable-experimental-associated-type-inference
 
 protocol P {
   associatedtype A = Int
 }
 
 struct S<A>: P {}
 
-// This is unfortunate but it is the behavior of Swift 5.10.
-let x: Int.Type = S<String>.A.self
+let x: String.Type = S<String>.A.self

test/decl/protocol/req/associated_type_inference_fixed_type_experimental_inference.swift

@@ -212,7 +212,7 @@ do {
   // CHECK-NEXT: }
   struct Conformer1: P17a {} // expected-error {{type 'Conformer1' does not conform to protocol 'P17a'}}
   // CHECK-LABEL: Abstract type witness system for conformance of Conformer2<A> to P17b: {
-  // CHECK-NEXT: A => (unresolved), [[EQUIV_CLASS:0x[0-9a-f]+]]
+  // CHECK-NEXT: A => A (preferred), [[EQUIV_CLASS:0x[0-9a-f]+]]
   // CHECK-NEXT: B => (unresolved)
   // CHECK-NEXT: }
   struct Conformer2<A>: P17b {} // expected-error {{type 'Conformer2<A>' does not conform to protocol 'P17b'}}
@@ -221,7 +221,7 @@ do {
   // CHECK-NEXT: }
   struct Conformer3: P17c {}
   // CHECK-LABEL: Abstract type witness system for conformance of Conformer4<A> to P17d: {
-  // CHECK-NEXT: A => Int (preferred), [[EQUIV_CLASS:0x[0-9a-f]+]]
+  // CHECK-NEXT: A => A (preferred), [[EQUIV_CLASS:0x[0-9a-f]+]]
   // CHECK-NEXT: B => Int (preferred), [[EQUIV_CLASS:0x[0-9a-f]+]]
   // CHECK-NEXT: }
   struct Conformer4<A>: P17d {}

test/decl/protocol/req/associated_type_inference_foundation_1.swift

@@ -0,0 +1,33 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+
+// REQUIRES: objc_interop
+
+import Foundation
+
+public struct CustomCollection<T>: RandomAccessCollection {
+    public typealias Indices = Range<Int>
+
+    public var startIndex: Int { fatalError() }
+    public var endIndex: Int { fatalError() }
+    public var count: Int { fatalError() }
+
+    public subscript(position: Int) -> T {
+        get { fatalError() }
+        set { fatalError() }
+    }
+
+    public var underestimatedCount: Int { fatalError() }
+}
+
+extension CustomCollection: ContiguousBytes where Element == UInt8 {
+    public func withUnsafeBytes<R>(_ body: (UnsafeRawBufferPointer) throws -> R) rethrows -> R {
+        fatalError()
+    }
+}
+
+extension CustomCollection: DataProtocol where Element == UInt8 {
+    public var regions: CollectionOfOne<CustomCollection<UInt8>> {
+        fatalError()
+    }
+}

test/decl/protocol/req/associated_type_inference_foundation_2.swift

@@ -0,0 +1,45 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+
+// REQUIRES: objc_interop
+
+import Foundation
+
+public struct CustomCollection<Element>: RandomAccessCollection, MutableCollection, ExpressibleByArrayLiteral {
+    public typealias Indices = Range<Int>
+    public typealias Index = Int
+
+    public init() {}
+
+    public init(arrayLiteral elements: Element ...) {
+        fatalError()
+    }
+}
+
+extension CustomCollection {
+    public var startIndex: Int { fatalError() }
+    public var endIndex: Int { fatalError() }
+
+    public subscript(position: Int) -> Element {
+        get { fatalError() }
+        set { fatalError() }
+    }
+}
+
+extension CustomCollection: RangeReplaceableCollection {
+    public mutating func append(_ newElement: Element) { }
+    public mutating func append<S: Sequence>(contentsOf newElements: S) where S.Element == Element { }
+    public mutating func reserveCapacity(_ minimumCapacity: Int) { }
+    public mutating func removeAll(keepingCapacity keepCapacity: Bool = false) { }
+    public mutating func replaceSubrange<C: Collection>(_ subRange: Range<Int>, with newElements: C) where C.Element == Element { }
+}
+
+extension CustomCollection {
+    public func withUnsafeBytes<R>(_ body: (UnsafeRawBufferPointer) throws -> R) rethrows -> R { fatalError() }
+}
+
+extension CustomCollection: ContiguousBytes where Element == UInt8 { }
+
+extension CustomCollection: DataProtocol where Element == UInt8 {
+    public var regions: CollectionOfOne<CustomCollection<UInt8>> { fatalError() }
+}

test/decl/protocol/req/associated_type_inference_stdlib_3.swift

@@ -1,8 +1,6 @@
-// RUN: not %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
 // RUN: not %target-typecheck-verify-swift -disable-experimental-associated-type-inference
 
-// FIXME: Get this passing with -enable-experimental-associated-type-inference again.
-
 struct FooIterator<T: Sequence>: IteratorProtocol {
   typealias Element = T.Element
 

test/decl/protocol/req/rdar123261282.swift

@@ -0,0 +1,21 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+
+struct CustomCollection<T>: RandomAccessCollection {
+    struct CustomIndices: RandomAccessCollection {
+        var count: Int { fatalError() }
+
+        var startIndex: Int { fatalError() }
+        var endIndex: Int { fatalError() }
+
+        subscript(index: Int) -> Int { fatalError() }
+    }
+
+    var count: Int { fatalError() }
+    var indices: CustomIndices { fatalError() }
+    var startIndex: Int { fatalError() }
+    var endIndex: Int { fatalError() }
+    func index(before i: Int) -> Int { fatalError() }
+    func index(after i: Int) -> Int { fatalError() }
+    subscript(index: Int) -> T { fatalError() }
+}

test/decl/protocol/req/rdar123262178.swift

@@ -0,0 +1,14 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+
+public protocol P {
+  associatedtype A = Never
+}
+
+public struct G<A>: P {}
+
+public struct ConcreteA {}
+
+public protocol Q: P where Self.A == ConcreteA {}
+
+extension G: Q where A == ConcreteA {}

test/decl/protocol/req/rdar123270042.swift

@@ -0,0 +1,28 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: not %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+
+func f<T: Sequence>(_: T) -> T.Element {}
+
+let x: Int = f(CustomCollection<Int>())
+
+struct CustomCollection<Element>: RandomAccessCollection, MutableCollection {
+    typealias SubSequence = ArraySlice<Element>
+    typealias Index = Int
+    typealias Indices = Range<Int>
+
+    var startIndex: Int { fatalError() }
+    var endIndex: Int { fatalError() }
+    var first: Element? { fatalError() }
+    var last: Element? { fatalError() }
+
+    subscript(position: Index) -> Element {
+        get { fatalError() }
+        set { fatalError() }
+    }
+
+    subscript(bounds: Indices) -> SubSequence {
+        get { fatalError() }
+        set { fatalError() }
+    }
+}
+

test/decl/protocol/req/rdar123345520.swift

@@ -0,0 +1,28 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
+
+struct G1<T> {}
+
+struct G2<A, B>: AP {
+    func f1(_: G1<(B) -> A>, _: G1<B>) -> G1<A> { fatalError() }
+    func f2<C>(_: (A) -> C) -> G1<C> { fatalError() }
+}
+
+protocol OP: EP {
+    associatedtype L
+    associatedtype R
+
+    func f1(_: G1<L>, _: G1<R>) -> G1<A>
+}
+
+extension OP {
+    func f1(_: G1<L>?, _: G1<R>?) -> G1<A> { fatalError() }
+}
+
+protocol AP: OP where L == (B) -> A, R == B {}
+
+protocol EP {
+    associatedtype A
+    associatedtype B
+    func f2<C>(_: (A) -> C) -> G1<C>
+}