[CodeCompletion] Use Self type in extension for generic substitution when doing implicit member lookup on protocol type

This fixes the regression introduced in #75068 (comment).

Author: ahoppen

lib/IDE/CompletionLookup.cpp

@@ -720,6 +720,35 @@ Type CompletionLookup::getTypeOfMember(const ValueDecl *VD, Type ExprType) {
       if (MaybeNominalType->hasUnboundGenericType())
         return T;
 
+      // If we are doing implicit member lookup on a protocol and we have found
+      // a declaration in an extension, use the extension's `Self` type for the
+      // generic substitution.
+      // Eg in the following, the `Self` type returned by `qux` is
+      // `MyGeneric<Int>`, not `MyProto` because of the `Self` type restriction.
+      // ```
+      // protocol MyProto {}
+      // struct MyGeneric<T>: MyProto {}
+      // extension MyProto where Self == MyGeneric<Int> {
+      //   static func qux() -> Self { .init() }
+      // }
+      // func takeMyProto(_: any MyProto)  {}
+      // func test() {
+      //   takeMyProto(.#^COMPLETE^#)
+      // }
+      // ```
+      if (MaybeNominalType->isExistentialType()) {
+        Type SelfType;
+        if (auto ED = dyn_cast<ExtensionDecl>(VD->getDeclContext())) {
+          SelfType = ED->getGenericSignature()->getConcreteType(
+              ED->getSelfInterfaceType());
+        }
+        if (SelfType) {
+          MaybeNominalType = SelfType;
+        } else {
+          return T;
+        }
+      }
+
       // For everything else, substitute in the base type.
       auto Subs = MaybeNominalType->getMemberSubstitutionMap(VD);
 

test/IDE/complete_from_stdlib.swift

@@ -28,7 +28,7 @@ func protocolExtCollection1b(_ a: Collection) {
 }
 
 // FIXME(https://github.com/apple/swift/issues/65696): We should not be showing this because (1) it cannot be accessed on the existential (2) we don't have the syntax and features to represent the projected type sig anyway.
-// PRIVATE_NOMINAL_MEMBERS_2B-DAG: map({#(transform): (any Collection.Element) throws(Error) -> T##(any Collection.Element) throws(Error) -> T#})[' throws'][#[T]#]{{; name=.+}}
+// PRIVATE_NOMINAL_MEMBERS_2B-DAG: map({#(transform): (Collection.Element) throws(Error) -> T##(Collection.Element) throws(Error) -> T#})[' throws'][#[T]#]{{; name=.+}}
 // NEGATIVE_PRIVATE_NOMINAL_MEMBERS_2B-NOT: Decl{{.*}}: index({#before: any Comparable#})
 
 func protocolExtCollection2<C : Collection where C.Index : BidirectionalIndex>(_ a: C) {

test/IDE/complete_rdar129024996.swift

@@ -46,4 +46,4 @@ extension P where Self == B<A> {
 func bar() {}
 // COMPLETE:     Begin completions, 2 items
 // COMPLETE-DAG: Decl[StaticMethod]/Super/TypeRelation[Convertible]: foo()[#S#]; name=foo()
-// COMPLETE-DAG: Decl[StaticMethod]/Super/Flair[ExprSpecific]/TypeRelation[Convertible]: qux()[#any Q#]; name=qux()
+// COMPLETE-DAG: Decl[StaticMethod]/Super/TypeRelation[Convertible]: qux()[#B<S>#]; name=qux()

test/IDE/complete_value_expr.swift

@@ -1523,8 +1523,8 @@ func testUnusableProtExt(_ x: PWithT) {
   x.#^PROTOCOL_EXT_UNUSABLE_EXISTENTIAL^#
 }
 // FIXME(https://github.com/apple/swift/issues/65696): We should not be showing these because (1) they cannot be accessed on the existential (2) we don't have the syntax and features to represent the projected type sigs anyway.
-// PROTOCOL_EXT_UNUSABLE_EXISTENTIAL: Decl[InstanceMethod]/CurrNominal:   foo({#(x): any PWithT.T#})[#any PWithT.T#]{{; name=.+}}
-// PROTOCOL_EXT_UNUSABLE_EXISTENTIAL: Decl[InstanceMethod]/CurrNominal:   bar({#(x): any PWithT.T#})[#any PWithT.T#]{{; name=.+}}
+// PROTOCOL_EXT_UNUSABLE_EXISTENTIAL: Decl[InstanceMethod]/CurrNominal:   foo({#(x): PWithT.T#})[#PWithT.T#]{{; name=.+}}
+// PROTOCOL_EXT_UNUSABLE_EXISTENTIAL: Decl[InstanceMethod]/CurrNominal:   bar({#(x): PWithT.T#})[#PWithT.T#]{{; name=.+}}
 
 protocol dedupP {
   associatedtype T
@@ -1559,10 +1559,10 @@ func testDeDuped(_ x: dedupS) {
 func testDeDuped2(_ x: dedupP) {
   x#^PROTOCOL_EXT_DEDUP_2^#
 // PROTOCOL_EXT_DEDUP_2: Begin completions, 5 items
-// PROTOCOL_EXT_DEDUP_2-DAG: Decl[InstanceMethod]/CurrNominal:   .foo()[#any dedupP.T#]; name=foo()
-// PROTOCOL_EXT_DEDUP_2-DAG: Decl[InstanceVar]/CurrNominal:      .bar[#any dedupP.T#]; name=bar
+// PROTOCOL_EXT_DEDUP_2-DAG: Decl[InstanceMethod]/CurrNominal:   .foo()[#dedupP.T#]; name=foo()
+// PROTOCOL_EXT_DEDUP_2-DAG: Decl[InstanceVar]/CurrNominal:      .bar[#dedupP.T#]; name=bar
 // FIXME(https://github.com/apple/swift/issues/65696): We should not be showing this because (1) it cannot be accessed on the existential (2) we don't have the syntax and features to represent the projected type sig anyway.
-// PROTOCOL_EXT_DEDUP_2-DAG: Decl[Subscript]/CurrNominal:        [{#(x): any dedupP.T#}][#any dedupP.T#]; name=[:]
+// PROTOCOL_EXT_DEDUP_2-DAG: Decl[Subscript]/CurrNominal:        [{#(x): dedupP.T#}][#dedupP.T#]; name=[:]
 // PROTOCOL_EXT_DEDUP_2-DAG: Keyword[self]/CurrNominal:          .self[#any dedupP#]; name=self
 }
 func testDeDuped3<T : dedupP where T.T == Int>(_ x: T) {