[CodeCompletion] Improve getReturnTypeFromContext()

* For methods, un-curry function interface type so we can get declared
return type.
* For closures, fall back to getting explicit result type in case we
cannot retrieve it from the type of closure itself.

Author: rintaro

lib/IDE/CodeCompletion.cpp

@@ -1509,13 +1509,20 @@ static bool isTopLevelContext(const DeclContext *DC) {
 static Type getReturnTypeFromContext(const DeclContext *DC) {
   if (auto FD = dyn_cast<AbstractFunctionDecl>(DC)) {
     if (FD->hasInterfaceType()) {
-      if (auto FT = FD->getInterfaceType()->getAs<FunctionType>()) {
+      auto Ty = FD->getInterfaceType();
+      if (FD->getDeclContext()->isTypeContext())
+        Ty = FD->getMethodInterfaceType();
+      if (auto FT = Ty->getAs<AnyFunctionType>())
         return FT->getResult();
-      }
     }
-  } else if (auto CE = dyn_cast<AbstractClosureExpr>(DC)) {
-    if (CE->getType()) {
-      return CE->getResultType();
+  } else if (auto ACE = dyn_cast<AbstractClosureExpr>(DC)) {
+    if (ACE->getType())
+      return ACE->getResultType();
+    if (auto CE = dyn_cast<ClosureExpr>(ACE)) {
+      if (CE->hasExplicitResultType())
+        return const_cast<ClosureExpr *>(CE)
+            ->getExplicitResultTypeLoc()
+            .getType();
     }
   }
   return Type();
@@ -5452,11 +5459,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
 
   case CompletionKind::ReturnStmtExpr : {
     SourceLoc Loc = P.Context.SourceMgr.getCodeCompletionLoc();
-    if (auto FD = dyn_cast<AbstractFunctionDecl>(CurDeclContext)) {
-      if (auto FT = FD->getInterfaceType()->getAs<FunctionType>()) {
-        Lookup.setExpectedTypes(FT->getResult());
-      }
-    }
+    Lookup.setExpectedTypes(getReturnTypeFromContext(CurDeclContext));
     Lookup.getValueCompletionsInDeclContext(Loc);
     break;
   }

test/IDE/complete_return.swift

@@ -7,6 +7,15 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR1 | %FileCheck %s -check-prefix=RETURN_TR1
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR2 | %FileCheck %s -check-prefix=RETURN_TR2
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR3 | %FileCheck %s -check-prefix=RETURN_TR3
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR1_METHOD | %FileCheck %s -check-prefix=RETURN_TR1
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR2_METHOD | %FileCheck %s -check-prefix=RETURN_TR2
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR3_METHOD | %FileCheck %s -check-prefix=RETURN_TR3
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR1_STATICMETHOD | %FileCheck %s -check-prefix=RETURN_TR1
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR2_STATICMETHOD | %FileCheck %s -check-prefix=RETURN_TR2
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR3_STATICMETHOD | %FileCheck %s -check-prefix=RETURN_TR3
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR1_CLOSURE | %FileCheck %s -check-prefix=RETURN_TR1
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR2_CLOSURE | %FileCheck %s -check-prefix=RETURN_TR2
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=RETURN_TR3_CLOSURE | %FileCheck %s -check-prefix=RETURN_TR3
 
 struct FooStruct {
   var instanceVar : Int
@@ -104,3 +113,47 @@ func testTR3(_ g : Gen) -> Int? {
 // RETURN_TR3-DAG: Decl[InstanceMethod]/CurrNominal:   InternalStringOpGen()[#String?#]{{; name=.+$}}
 // RETURN_TR3-DAG: Decl[InstanceMethod]/CurrNominal/NotRecommended/TypeRelation[Invalid]: InternalIntTaker({#(i1): Int#}, {#i2: Int#})[#Void#]{{; name=.+$}}
 // RETURN_TR3-DAG: Decl[InstanceMethod]/CurrNominal/NotRecommended/TypeRelation[Invalid]: InternalStringTaker({#(s1): String#}, {#s2: String#})[#Void#]{{; name=.+$}}
+
+struct TestStruct {
+  func testTR1_method() -> Int? {
+    var i : Int
+    var oi : Int?
+    var fs : FooStruct
+    return #^RETURN_TR1_METHOD^#
+  }
+  func testTR2_method(_ g : Gen) -> Int? {
+    return g.#^RETURN_TR2_METHOD^#
+  }
+  func testTR3_method(_ g : Gen) -> Int? {
+    return g.IG.#^RETURN_TR3_METHOD^#
+  }
+
+  static func testTR1_static() -> Int? {
+    var i : Int
+    var oi : Int?
+    var fs : FooStruct
+    return #^RETURN_TR1_STATICMETHOD^#
+  }
+  static func testTR2_static(_ g : Gen) -> Int? {
+    return g.#^RETURN_TR2_STATICMETHOD^#
+  }
+  static func testTR3_static(_ g : Gen) -> Int? {
+    return g.IG.#^RETURN_TR3_STATICMETHOD^#
+  }
+}
+
+func testClosures(_ g: Gen) {
+  var i : Int
+  var oi : Int?
+  var fs : FooStruct
+
+  _ = { () -> Int? in
+    return #^RETURN_TR1_CLOSURE^#
+  }
+  _ = { () -> Int? in
+    return g.#^RETURN_TR2_CLOSURE^#
+  }
+  _ = { () -> Int? in
+    return g.IG.#^RETURN_TR3_CLOSURE^#
+  }
+}