[CodeCompletion] Use TypeContextInfo to get expected return types

If the type check return type from the context has unresolved types, it
can't be used for checking convertibility. Fallback to get
'TypeContextInfo' of the closure position.

rdar://problem/66002497

Author: rintaro

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
+}