[Completion] Better handle merging of lookup base types

For unresolved member completion, we were preferring
the more general type, when we ought to be preferring
the more specific type. Additionally, for both
unresolved member and postfix completion we were
opening archetypes, which doesn't work as expected
since we don't compare requirements. Factor out
the logic that deals with merging base types for
lookup, and have it prefer either the subtype, or
the optional type in the case of optional promotion.

rdar://126168123

include/swift/IDE/PostfixCompletion.h

@@ -66,13 +66,9 @@ class PostfixCompletionCallback : public TypeCheckCompletionCallback {
     llvm::DenseMap<AbstractClosureExpr *, ActorIsolation>
         ClosureActorIsolations;
 
-    /// Checks whether this result has the same \c BaseTy and \c BaseDecl as
-    /// \p Other and if the two can thus be merged to be one value lookup in
-    /// \c deliverResults.
-    bool canBeMergedWith(const Result &Other, DeclContext &DC) const;
-
-    /// Merge this result with \p Other. Assumes that they can be merged.
-    void merge(const Result &Other, DeclContext &DC);
+    /// Merge this result with \p Other, returning \c true if
+    /// successful, else \c false.
+    bool tryMerge(const Result &Other, DeclContext *DC);
   };
 
   CodeCompletionExpr *CompletionExpr;

include/swift/IDE/UnresolvedMemberCompletion.h

@@ -33,13 +33,9 @@ class UnresolvedMemberTypeCheckCompletionCallback
     /// functions is supported.
     bool IsInAsyncContext;
 
-    /// Checks whether this result has the same \c BaseTy and \c BaseDecl as
-    /// \p Other and if the two can thus be merged to be one value lookup in
-    /// \c deliverResults.
-    bool canBeMergedWith(const Result &Other, DeclContext &DC) const;
-
-    /// Merge this result with \p Other. Assumes that they can be merged.
-    void merge(const Result &Other, DeclContext &DC);
+    /// Attempts to merge this result with \p Other, returning \c true if
+    /// successful, else \c false.
+    bool tryMerge(const Result &Other, DeclContext *DC);
   };
 
   CodeCompletionExpr *CompletionExpr;

include/swift/Sema/IDETypeChecking.h

@@ -60,11 +60,20 @@ namespace swift {
   /// Typecheck binding initializer at \p bindingIndex.
   void typeCheckPatternBinding(PatternBindingDecl *PBD, unsigned bindingIndex);
 
+  /// Attempt to merge two types for the purposes of completion lookup. In
+  /// general this means preferring a subtype over a supertype, but can also e.g
+  /// prefer an optional over a non-optional. If the two types are incompatible,
+  /// null is returned.
+  Type tryMergeBaseTypeForCompletionLookup(Type ty1, Type ty2, DeclContext *dc);
+
   /// Check if T1 is convertible to T2.
   ///
   /// \returns true on convertible, false on not.
   bool isConvertibleTo(Type T1, Type T2, bool openArchetypes, DeclContext &DC);
 
+  /// Check whether \p T1 is a subtype of \p T2.
+  bool isSubtypeOf(Type T1, Type T2, DeclContext *DC);
+
   void collectDefaultImplementationForProtocolMembers(ProtocolDecl *PD,
                         llvm::SmallDenseMap<ValueDecl*, ValueDecl*> &DefaultMap);
 

include/swift/Sema/IDETypeCheckingRequests.h

@@ -86,6 +86,7 @@ public:
 //----------------------------------------------------------------------------//
 enum class TypeRelation: uint8_t {
   ConvertTo,
+  SubtypeOf
 };
 
 struct TypePair {
@@ -153,6 +154,7 @@ struct TypeRelationCheckInput {
     switch(owner.Relation) {
 #define CASE(NAME) case TypeRelation::NAME: out << #NAME << " "; break;
     CASE(ConvertTo)
+    CASE(SubtypeOf)
 #undef CASE
     }
   }

lib/IDE/IDETypeChecking.cpp

@@ -941,6 +941,45 @@ bool swift::isMemberDeclApplied(const DeclContext *DC, Type BaseTy,
     IsDeclApplicableRequest(DeclApplicabilityOwner(DC, BaseTy, VD)), false);
 }
 
+Type swift::tryMergeBaseTypeForCompletionLookup(Type ty1, Type ty2,
+                                                DeclContext *dc) {
+  // Easy case, equivalent so just pick one.
+  if (ty1->isEqual(ty2))
+    return ty1;
+
+  // Check to see if one is an optional of another. In that case, prefer the
+  // optional since we can unwrap a single level when doing a lookup.
+  {
+    SmallVector<Type, 4> ty1Optionals;
+    SmallVector<Type, 4> ty2Optionals;
+    auto ty1Unwrapped = ty1->lookThroughAllOptionalTypes(ty1Optionals);
+    auto ty2Unwrapped = ty2->lookThroughAllOptionalTypes(ty2Optionals);
+
+    if (ty1Unwrapped->isEqual(ty2Unwrapped)) {
+      // We currently only unwrap a single level of optional, so if the
+      // difference is greater, don't merge.
+      if (ty1Optionals.size() == 1 && ty2Optionals.empty())
+        return ty1;
+      if (ty2Optionals.size() == 1 && ty1Optionals.empty())
+        return ty2;
+    }
+    // We don't want to consider subtyping for optional mismatches since
+    // optional promotion is modelled as a subtype, which isn't useful for us
+    // (i.e if we have T? and U, preferring U would miss members on T?).
+    if (ty1Optionals.size() != ty2Optionals.size())
+      return Type();
+  }
+
+  // In general we want to prefer a subtype over a supertype.
+  if (isSubtypeOf(ty1, ty2, dc))
+    return ty1;
+  if (isSubtypeOf(ty2, ty1, dc))
+    return ty2;
+
+  // Incomparable, return null.
+  return Type();
+}
+
 bool swift::isConvertibleTo(Type T1, Type T2, bool openArchetypes,
                             DeclContext &DC) {
   return evaluateOrDefault(DC.getASTContext().evaluator,
@@ -948,6 +987,12 @@ bool swift::isConvertibleTo(Type T1, Type T2, bool openArchetypes,
       TypeRelation::ConvertTo, openArchetypes)), false);
 }
 
+bool swift::isSubtypeOf(Type T1, Type T2, DeclContext *DC) {
+  return evaluateOrDefault(DC->getASTContext().evaluator,
+    TypeRelationCheckRequest(TypeRelationCheckInput(DC, T1, T2,
+      TypeRelation::SubtypeOf, /*openArchetypes*/ false)), false);
+}
+
 Type swift::getRootTypeOfKeypathDynamicMember(SubscriptDecl *SD) {
   return evaluateOrDefault(SD->getASTContext().evaluator,
     RootTypeOfKeypathDynamicMemberRequest{SD}, Type());

lib/IDE/PostfixCompletion.cpp

@@ -21,31 +21,20 @@ using namespace swift;
 using namespace swift::constraints;
 using namespace swift::ide;
 
-bool PostfixCompletionCallback::Result::canBeMergedWith(const Result &Other,
-                                                        DeclContext &DC) const {
-  if (BaseDecl != Other.BaseDecl) {
+bool PostfixCompletionCallback::Result::tryMerge(const Result &Other,
+                                                 DeclContext *DC) {
+  if (BaseDecl != Other.BaseDecl)
     return false;
-  }
-  if (!BaseTy->isEqual(Other.BaseTy) &&
-      !isConvertibleTo(BaseTy, Other.BaseTy, /*openArchetypes=*/true, DC) &&
-      !isConvertibleTo(Other.BaseTy, BaseTy, /*openArchetypes=*/true, DC)) {
-    return false;
-  }
-  return true;
-}
 
-void PostfixCompletionCallback::Result::merge(const Result &Other,
-                                              DeclContext &DC) {
-  assert(canBeMergedWith(Other, DC));
   // These properties should match if we are talking about the same BaseDecl.
   assert(IsBaseDeclUnapplied == Other.IsBaseDeclUnapplied);
   assert(BaseIsStaticMetaType == Other.BaseIsStaticMetaType);
 
-  if (!BaseTy->isEqual(Other.BaseTy) &&
-      isConvertibleTo(Other.BaseTy, BaseTy, /*openArchetypes=*/true, DC)) {
-    // Pick the more specific base type as it will produce more solutions.
-    BaseTy = Other.BaseTy;
-  }
+  auto baseTy = tryMergeBaseTypeForCompletionLookup(BaseTy, Other.BaseTy, DC);
+  if (!baseTy)
+    return false;
+
+  BaseTy = baseTy;
 
   // There could be multiple results that have different actor isolations if the
   // closure is an argument to a function that has multiple overloads with
@@ -66,18 +55,15 @@ void PostfixCompletionCallback::Result::merge(const Result &Other,
   ExpectsNonVoid &= Other.ExpectsNonVoid;
   IsImpliedResult |= Other.IsImpliedResult;
   IsInAsyncContext |= Other.IsInAsyncContext;
+  return true;
 }
 
 void PostfixCompletionCallback::addResult(const Result &Res) {
-  auto ExistingRes =
-      llvm::find_if(Results, [&Res, DC = DC](const Result &ExistingResult) {
-        return ExistingResult.canBeMergedWith(Res, *DC);
-      });
-  if (ExistingRes != Results.end()) {
-    ExistingRes->merge(Res, *DC);
-  } else {
-    Results.push_back(Res);
+  for (auto idx : indices(Results)) {
+    if (Results[idx].tryMerge(Res, DC))
+      return;
   }
+  Results.push_back(Res);
 }
 
 void PostfixCompletionCallback::fallbackTypeCheck(DeclContext *DC) {

lib/IDE/UnresolvedMemberCompletion.cpp

@@ -21,43 +21,27 @@ using namespace swift;
 using namespace swift::constraints;
 using namespace swift::ide;
 
-bool UnresolvedMemberTypeCheckCompletionCallback::Result::canBeMergedWith(
-    const Result &Other, DeclContext &DC) const {
-  if (!isConvertibleTo(ExpectedTy, Other.ExpectedTy, /*openArchetypes=*/true,
-                       DC) &&
-      !isConvertibleTo(Other.ExpectedTy, ExpectedTy, /*openArchetypes=*/true,
-                       DC)) {
+bool UnresolvedMemberTypeCheckCompletionCallback::Result::tryMerge(
+    const Result &Other, DeclContext *DC) {
+  auto expectedTy = tryMergeBaseTypeForCompletionLookup(ExpectedTy,
+                                                        Other.ExpectedTy, DC);
+  if (!expectedTy)
     return false;
-  }
-  return true;
-}
 
-void UnresolvedMemberTypeCheckCompletionCallback::Result::merge(
-    const Result &Other, DeclContext &DC) {
-  assert(canBeMergedWith(Other, DC));
-  if (!ExpectedTy->isEqual(Other.ExpectedTy) &&
-      isConvertibleTo(ExpectedTy, Other.ExpectedTy, /*openArchetypes=*/true,
-                      DC)) {
-    // ExpectedTy is more general than Other.ExpectedTy. Complete based on the
-    // more general type because it offers more completion options.
-    ExpectedTy = Other.ExpectedTy;
-  }
+  ExpectedTy = expectedTy;
 
   IsImpliedResult |= Other.IsImpliedResult;
   IsInAsyncContext |= Other.IsInAsyncContext;
+  return true;
 }
 
 void UnresolvedMemberTypeCheckCompletionCallback::addExprResult(
     const Result &Res) {
-  auto ExistingRes =
-      llvm::find_if(ExprResults, [&Res, DC = DC](const Result &ExistingResult) {
-        return ExistingResult.canBeMergedWith(Res, *DC);
-      });
-  if (ExistingRes != ExprResults.end()) {
-    ExistingRes->merge(Res, *DC);
-  } else {
-    ExprResults.push_back(Res);
+  for (auto idx : indices(ExprResults)) {
+    if (ExprResults[idx].tryMerge(Res, DC))
+      return;
   }
+  ExprResults.push_back(Res);
 }
 
 void UnresolvedMemberTypeCheckCompletionCallback::sawSolutionImpl(

lib/Sema/IDETypeCheckingRequests.cpp

@@ -246,10 +246,14 @@ IsDeclApplicableRequest::evaluate(Evaluator &evaluator,
 bool
 TypeRelationCheckRequest::evaluate(Evaluator &evaluator,
                                    TypeRelationCheckInput Owner) const {
-  std::optional<constraints::ConstraintKind> CKind;
+  using namespace constraints;
+  std::optional<ConstraintKind> CKind;
   switch (Owner.Relation) {
   case TypeRelation::ConvertTo:
-    CKind = constraints::ConstraintKind::Conversion;
+    CKind = ConstraintKind::Conversion;
+    break;
+  case TypeRelation::SubtypeOf:
+    CKind = ConstraintKind::Subtype;
     break;
   }
   assert(CKind.has_value());

test/IDE/complete_cgfloat_double.swift

@@ -0,0 +1,14 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -batch-code-completion -source-filename %s -filecheck %raw-FileCheck -completion-output-dir %t
+
+// REQUIRES: objc_interop
+
+import Foundation
+
+func foo(_ x: CGFloat) {}
+func foo(_ x: Double) {}
+
+// Make sure we suggest completions for both CGFloat and Double.
+foo(.#^FOO^#)
+// FOO-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Convertible]: init()[#CGFloat#]; name=init()
+// FOO-DAG: Decl[Constructor]/CurrNominal/IsSystem/TypeRelation[Convertible]: init()[#Double#]; name=init()

test/IDE/complete_rdar126168123.swift

@@ -0,0 +1,22 @@
+// RUN: %empty-directory(%t)
+// RUN: %swift-ide-test -batch-code-completion -source-filename %s -filecheck %raw-FileCheck -completion-output-dir %t
+
+// rdar://126168123
+
+protocol MyProto {}
+protocol MyProto2 {}
+
+struct MyStruct : MyProto {}
+
+extension MyProto where Self == MyStruct {
+  static var automatic: MyStruct { fatalError() }
+}
+
+func use<T: MyProto>(_ someT: T) {}
+func use<T: MyProto2>(_ someT: T) {}
+
+func test() {
+  use(.#^COMPLETE^#)
+}
+
+// COMPLETE: Decl[StaticVar]/CurrNominal/TypeRelation[Convertible]: automatic[#MyStruct#]; name=automatic

test/IDE/complete_subtype_overload.swift

@@ -0,0 +1,61 @@
+// RUN: %empty-directory(%t)
+// RUN: %swift-ide-test -batch-code-completion -source-filename %s -filecheck %raw-FileCheck -completion-output-dir %t
+
+class C {
+  static func cMethod() -> C {}
+}
+class D : C {
+  static func dMethod() -> D {}
+}
+
+func test1(_ x: C) {}
+func test1(_ x: D) {}
+
+// We prefer the subtype here, so we show completions for D.
+test1(.#^TEST1^#)
+// TEST1-DAG: Decl[StaticMethod]/Super:           cMethod()[#C#]; name=cMethod()
+// TEST1-DAG: Decl[StaticMethod]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: dMethod()[#D#]; name=dMethod()
+// TEST1-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Convertible]: init()[#D#]; name=init()
+
+func test2(_ x: C?) {}
+func test2(_ x: D?) {}
+
+test2(.#^TEST2^#)
+// TEST2-DAG: Decl[StaticMethod]/Super:           cMethod()[#C#]; name=cMethod()
+// TEST2-DAG: Decl[StaticMethod]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: dMethod()[#D#]; name=dMethod()
+// TEST2-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Convertible]: init()[#D#]; name=init()
+// TEST2-DAG: Decl[EnumElement]/CurrNominal/IsSystem/TypeRelation[Convertible]: none[#Optional<D>#]; name=none
+// TEST2-DAG: Decl[EnumElement]/CurrNominal/IsSystem/TypeRelation[Convertible]: some({#D#})[#Optional<D>#]; name=some()
+
+func test3(_ x: C?) {}
+func test3(_ x: D) {}
+
+// We can still provide both C and D completions here.
+test3(.#^TEST3^#)
+// TEST3-DAG: Decl[StaticMethod]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: cMethod()[#C#]; name=cMethod()
+// TEST3-DAG: Decl[StaticMethod]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: dMethod()[#D#]; name=dMethod()
+// TEST3-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Convertible]: init()[#D#]; name=init()
+// TEST3-DAG: Decl[EnumElement]/CurrNominal/IsSystem/TypeRelation[Convertible]: none[#Optional<C>#]; name=none
+// TEST3-DAG: Decl[EnumElement]/CurrNominal/IsSystem/TypeRelation[Convertible]: some({#C#})[#Optional<C>#]; name=some()
+
+func test4(_ x: Int) {}
+func test4(_ x: AnyHashable) {}
+
+// Make sure we show Int completions.
+test4(.#^TEST4^#)
+// TEST4: Decl[StaticVar]/Super/Flair[ExprSpecific]/IsSystem/TypeRelation[Convertible]: zero[#Int#]; name=zero
+
+protocol P {}
+extension P {
+  func pMethod() {}
+}
+struct S : P {
+  func sMethod() {}
+}
+
+func test5() -> any P {}
+func test5() -> S {}
+
+test5().#^TEST5^#
+// TEST5-DAG: Decl[InstanceMethod]/CurrNominal:   pMethod()[#Void#]; name=pMethod()
+// TEST5-DAG: Decl[InstanceMethod]/CurrNominal:   sMethod()[#Void#]; name=sMethod()