[CodeCompletion] Use TypeContextInfo to get expected return types

lib/IDE/CodeCompletion.cpp

@@ -6241,7 +6241,9 @@ void CodeCompletionCallbacksImpl::doneParsing() {
 
   case CompletionKind::ReturnStmtExpr : {
     SourceLoc Loc = P.Context.SourceMgr.getCodeCompletionLoc();
-    Lookup.setExpectedTypes(getReturnTypeFromContext(CurDeclContext),
+    SmallVector<Type, 2> possibleReturnTypes;
+    collectPossibleReturnTypesFromContext(CurDeclContext, possibleReturnTypes);
+    Lookup.setExpectedTypes(possibleReturnTypes,
                             /*isSingleExpressionBody*/ false);
     Lookup.getValueCompletionsInDeclContext(Loc);
     break;

lib/IDE/ExprContextAnalysis.cpp

@@ -285,40 +285,61 @@ bool swift::ide::removeCodeCompletionExpr(ASTContext &Ctx, Expr *&expr) {
 }
 
 //===----------------------------------------------------------------------===//
-// getReturnTypeFromContext(DeclContext)
+// collectPossibleReturnTypesFromContext(DeclContext, SmallVectorImpl<Type>)
 //===----------------------------------------------------------------------===//
 
-Type swift::ide::getReturnTypeFromContext(const DeclContext *DC) {
+void swift::ide::collectPossibleReturnTypesFromContext(
+    DeclContext *DC, SmallVectorImpl<Type> &candidates) {
   if (auto FD = dyn_cast<AbstractFunctionDecl>(DC)) {
     auto Ty = FD->getInterfaceType();
     if (FD->getDeclContext()->isTypeContext())
       Ty = FD->getMethodInterfaceType();
-    if (auto FT = Ty->getAs<AnyFunctionType>())
-      return DC->mapTypeIntoContext(FT->getResult());
-  } else if (auto ACE = dyn_cast<AbstractClosureExpr>(DC)) {
-    if (ACE->getType() && !ACE->getType()->hasError())
-      return ACE->getResultType();
+    if (auto FT = Ty->getAs<AnyFunctionType>()) {
+      candidates.push_back(DC->mapTypeIntoContext(FT->getResult()));
+    }
+  }
+
+  if (auto ACE = dyn_cast<AbstractClosureExpr>(DC)) {
+    // Use the type checked type if it has.
+    if (ACE->getType() && !ACE->getType()->hasError() &&
+        !ACE->getResultType()->hasUnresolvedType()) {
+      candidates.push_back(ACE->getResultType());
+      return;
+    }
+
     if (auto CE = dyn_cast<ClosureExpr>(ACE)) {
       if (CE->hasExplicitResultType()) {
+        // If the closure has a explicit return type, use it.
         if (auto ty = CE->getExplicitResultType()) {
-          return ty;
+          candidates.push_back(ty);
+          return;
+        } else {
+          auto typeLoc = TypeLoc{CE->getExplicitResultTypeRepr()};
+          if (!swift::performTypeLocChecking(
+                  DC->getASTContext(), typeLoc, /*isSILMode=*/false,
+                  /*isSILType=*/false, DC->getGenericEnvironmentOfContext(),
+                  const_cast<DeclContext *>(DC), /*diagnostics=*/false)) {
+            candidates.push_back(typeLoc.getType());
+            return;
+          }
         }
-
-        auto typeLoc = TypeLoc{CE->getExplicitResultTypeRepr()};
-        if (swift::performTypeLocChecking(DC->getASTContext(),
-                                          typeLoc,
-                                          /*isSILMode*/ false,
-                                          /*isSILType*/ false,
-                                          DC->getGenericEnvironmentOfContext(),
-                                          const_cast<DeclContext *>(DC),
-                                          /*diagnostics*/ false)) {
-          return Type();
+      } else {
+        // Otherwise, check the context type of the closure.
+        ExprContextInfo closureCtxInfo(CE->getParent(), CE);
+        for (auto closureTy : closureCtxInfo.getPossibleTypes()) {
+          if (auto funcTy = closureTy->getAs<AnyFunctionType>())
+            candidates.push_back(funcTy->getResult());
         }
-        return typeLoc.getType();
+        if (!candidates.empty())
+          return;
       }
     }
+
+    // Even if the type checked type has unresolved types, it's better than
+    // nothing.
+    if (ACE->getType() && !ACE->getType()->hasError())
+      candidates.push_back(ACE->getResultType());
   }
-  return Type();
 }
 
 //===----------------------------------------------------------------------===//
@@ -929,7 +950,10 @@ class ExprContextAnalyzer {
       auto *CE = cast<ClosureExpr>(Parent);
       assert(isSingleExpressionBodyForCodeCompletion(CE->getBody()));
       singleExpressionBody = true;
-      recordPossibleType(getReturnTypeFromContext(CE));
+      SmallVector<Type, 2> candidates;
+      collectPossibleReturnTypesFromContext(CE, candidates);
+      for (auto ty : candidates)
+        recordPossibleType(ty);
       break;
     }
     default:
@@ -939,9 +963,13 @@ class ExprContextAnalyzer {
 
   void analyzeStmt(Stmt *Parent) {
     switch (Parent->getKind()) {
-    case StmtKind::Return:
-      recordPossibleType(getReturnTypeFromContext(DC));
+    case StmtKind::Return: {
+      SmallVector<Type, 2> candidates;
+      collectPossibleReturnTypesFromContext(DC, candidates);
+      for (auto ty : candidates)
+        recordPossibleType(ty);
       break;
+    }
     case StmtKind::ForEach:
       if (auto SEQ = cast<ForEachStmt>(Parent)->getSequence()) {
         if (containsTarget(SEQ)) {
@@ -1004,7 +1032,10 @@ class ExprContextAnalyzer {
       if (auto *FD = dyn_cast<FuncDecl>(D)) {
         assert(isSingleExpressionBodyForCodeCompletion(FD->getBody()));
         singleExpressionBody = true;
-        recordPossibleType(getReturnTypeFromContext(FD));
+        SmallVector<Type, 2> candidates;
+        collectPossibleReturnTypesFromContext(DC, candidates);
+        for (auto ty : candidates)
+          recordPossibleType(ty);
         break;
       }
       llvm_unreachable("Unhandled decl kind.");

lib/IDE/ExprContextAnalysis.h

@@ -41,9 +41,10 @@ Expr *findParsedExpr(const DeclContext *DC, SourceRange TargetRange);
 ///       position.
 bool removeCodeCompletionExpr(ASTContext &Ctx, Expr *&expr);
 
-/// Returns expected return type of the given decl context.
+/// Collects possible expected return types of the given decl context.
 /// \p DC should be an \c AbstractFunctionDecl or an \c AbstractClosureExpr.
-Type getReturnTypeFromContext(const DeclContext *DC);
+void collectPossibleReturnTypesFromContext(DeclContext *DC,
+                                           SmallVectorImpl<Type> &candidates);
 
 struct FunctionTypeAndDecl {
   AnyFunctionType *Type;

test/IDE/complete_sr13271.swift

@@ -0,0 +1,66 @@
+// RUN: %swift-ide-test -code-completion -source-filename=%s -code-completion-token=A | %FileCheck %s --check-prefix=A
+// RUN: %swift-ide-test -code-completion -source-filename=%s -code-completion-token=B | %FileCheck %s --check-prefix=B
+// RUN: %swift-ide-test -code-completion -source-filename=%s -code-completion-token=D | %FileCheck %s --check-prefix=D
+
+// https://bugs.swift.org/browse/SR-13271
+// https://forums.swift.org/t/code-completion-enhancement-request/38677
+
+enum AIdentifier {
+  case a
+}
+
+enum BIdentifier {
+  case b
+}
+
+struct X { }
+struct Y { }
+
+struct A <T> {
+  private init(){}
+  static func foo (arg: Bool) -> A<X> { A<X>() }
+  static func bar (arg: Int) -> A<Y> { A<Y>() }
+}
+
+struct B {
+  static var baz: B { B() }
+}
+
+func C<T>(_ identifier: AIdentifier, _ a: ()->A<T>) -> D<T> { }
+func C(_ identifier: BIdentifier, _ b: ()->B) { }
+
+struct D <T> {
+    func sink (_ handler: @escaping (T)->()) { }
+}
+
+func test() {
+  C(.a) {
+    .#^A^#
+  }
+// A: Begin completions, 2 items
+// A-DAG: Decl[StaticMethod]/CurrNominal/TypeRelation[Convertible]: foo({#arg: Bool#})[#A<X>#];
+// A-DAG: Decl[StaticMethod]/CurrNominal/TypeRelation[Convertible]: bar({#arg: Int#})[#A<Y>#];
+// A: End completions
+}
+
+func test() {
+  C(.b) {
+    .#^B^#
+  }
+// B: Begin completions, 2 items
+// B-DAG: Decl[StaticVar]/CurrNominal/TypeRelation[Identical]: baz[#B#]; name=baz
+// B-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Identical]: init()[#B#]; name=init()
+// B: End completions
+}
+
+func test() {
+  C(.a) {
+    .foo(arg: true)
+  }
+  .sink { value in
+    value.#^D^#
+  }
+// D: Begin completions, 1 items
+// D-DAG: Keyword[self]/CurrNominal:          self[#X#];
+// D: End completions
+}

test/IDE/complete_unresolved_members.swift

@@ -122,6 +122,7 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TERNARY_5 | %FileCheck %s -check-prefix=UNRESOLVED_3_NOTIDEAL
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TERNARY_6 | %FileCheck %s -check-prefix=UNRESOLVED_3_NOTIDEAL
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TERNARY_CONDITION | %FileCheck %s -check-prefix=TERNARY_CONDITION
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=OVERLOADED_CLOSURE_RETURN | %FileCheck %s -check-prefix=OVERLOADED_CLOSURE_RETURN
 
 enum SomeEnum1 {
   case South
@@ -783,3 +784,17 @@ func testTernaryOperator2(cond: Bool) {
 // TERNARY_CONDITION-DAG: Decl[Constructor]/CurrNominal/IsSystem/TypeRelation[Identical]: init()[#Bool#]; name=init()
 // TERNARY_CONDITION: End completions
 }
+
+func overloadedClosureRcv(_: () -> SomeEnum1) {}
+func overloadedClosureRcv(_: () -> SomeEnum2) {}
+func testClosureReturnTypeForOverloaded() {
+  overloadedClosureRcv {
+    .#^OVERLOADED_CLOSURE_RETURN^#
+  }
+// OVERLOADED_CLOSURE_RETURN: Begin completions, 4 items
+// OVERLOADED_CLOSURE_RETURN-DAG: Decl[EnumElement]/CurrNominal/TypeRelation[Identical]: South[#SomeEnum1#];
+// OVERLOADED_CLOSURE_RETURN-DAG: Decl[EnumElement]/CurrNominal/TypeRelation[Identical]: North[#SomeEnum1#];
+// OVERLOADED_CLOSURE_RETURN-DAG: Decl[EnumElement]/CurrNominal/TypeRelation[Identical]: East[#SomeEnum2#];
+// OVERLOADED_CLOSURE_RETURN-DAG: Decl[EnumElement]/CurrNominal/TypeRelation[Identical]: West[#SomeEnum2#];
+// OVERLOADED_CLOSURE_RETURN: End completions
+}