[CodeCompletion] Show 'return' completion by default when appropriate

lib/IDE/CodeCompletion.cpp

@@ -1361,7 +1361,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   void completeReturnStmt(CodeCompletionExpr *E) override;
   void completeAfterPound(CodeCompletionExpr *E, StmtKind ParentKind) override;
   void completeGenericParams(TypeLoc TL) override;
-  void addKeywords(CodeCompletionResultSink &Sink);
+  void addKeywords(CodeCompletionResultSink &Sink, bool MaybeFuncBody);
 
   void doneParsing() override;
 
@@ -4451,8 +4451,11 @@ static void addDeclKeywords(CodeCompletionResultSink &Sink) {
   AddCSKeyword("convenience");
 }
 
-static void addStmtKeywords(CodeCompletionResultSink &Sink) {
+static void addStmtKeywords(CodeCompletionResultSink &Sink, bool MaybeFuncBody) {
   auto AddKeyword = [&](StringRef Name, CodeCompletionKeywordKind Kind) {
+    if (!MaybeFuncBody && Kind == CodeCompletionKeywordKind::kw_return)
+      return;
+
     CodeCompletionResultBuilder Builder(
         Sink, CodeCompletionResult::ResultKind::Keyword,
         SemanticContextKind::None, {});
@@ -4505,7 +4508,8 @@ static void addExprKeywords(CodeCompletionResultSink &Sink) {
 }
 
 
-void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink) {
+void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
+                                              bool MaybeFuncBody) {
   switch (Kind) {
   case CompletionKind::None:
   case CompletionKind::DotExpr:
@@ -4523,7 +4527,7 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink) {
 
   case CompletionKind::StmtOrExpr:
     addDeclKeywords(Sink);
-    addStmtKeywords(Sink);
+    addStmtKeywords(Sink, MaybeFuncBody);
     SWIFT_FALLTHROUGH;
   case CompletionKind::AssignmentRHS:
   case CompletionKind::ReturnStmtExpr:
@@ -4778,8 +4782,15 @@ void CodeCompletionCallbacksImpl::doneParsing() {
     return;
   }
 
+  bool MaybeFuncBody = true;
+  if (CurDeclContext) {
+    auto *CD = CurDeclContext->getLocalContext();
+    if (!CD || CD->getContextKind() == DeclContextKind::Initializer ||
+        CD->getContextKind() == DeclContextKind::TopLevelCodeDecl)
+      MaybeFuncBody = false;
+  }
   // Add keywords even if type checking fails completely.
-  addKeywords(CompletionContext.getResultSink());
+  addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
 
   if (!typecheckContext())
     return;

test/IDE/complete_keywords.swift

@@ -1,10 +1,47 @@
-// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TOP_LEVEL_1 | FileCheck %s -check-prefix=KW_DECL_STMT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TOP_LEVEL_1 > %t.top1
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.top1
+// RUN: FileCheck %s -check-prefix=KW_NO_RETURN < %t.top1
 
-// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_1 | FileCheck %s -check-prefix=KW_DECL_STMT
-// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_2 | FileCheck %s -check-prefix=KW_DECL_STMT
-// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_3 | FileCheck %s -check-prefix=KW_DECL_STMT
-// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_4 | FileCheck %s -check-prefix=KW_DECL_STMT
-// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_5 | FileCheck %s -check-prefix=KW_DECL_STMT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TOP_LEVEL_2 > %t.top2
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.top2
+// RUN: FileCheck %s -check-prefix=KW_NO_RETURN < %t.top2
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_1 > %t.func1
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.func1
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.func1
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_2 > %t.func2
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.func2
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.func2
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_3 > %t.func3
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.func3
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.func3
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_4 > %t.func4
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.func4
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.func4
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_FUNC_BODY_5 > %t.func5
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.func5
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.func5
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_CLOSURE_1 > %t.clos1
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.clos1
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.clos1
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_CLOSURE_2 > %t.clos2
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.clos2
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.clos2
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_CLOSURE_3 > %t.clos3
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.clos3
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.clos3
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_CLOSURE_4 > %t.clos4
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.clos4
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.clos4
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_SUBSCRIPT_1 > %t.subs
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.subs
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.subs
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_INIT_1 > %t.init
+// RUN: FileCheck %s -check-prefix=KW_DECL_STMT < %t.init
+// RUN: FileCheck %s -check-prefix=KW_RETURN < %t.init
 
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_NOMINAL_DECL_1 | FileCheck %s -check-prefix=KW_DECL
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_NOMINAL_DECL_2 | FileCheck %s -check-prefix=KW_DECL
@@ -41,6 +78,9 @@
 // RUN: FileCheck %s -check-prefix=KW_EXPR < %t.expr6
 // RUN: FileCheck %s -check-prefix=KW_EXPR_NEG < %t.expr6
 
+// KW_RETURN: Keyword[return]/None: return{{; name=.+$}}
+// KW_NO_RETURN-NOT: Keyword[return]
+
 // KW_DECL: Begin completions
 // KW_DECL-DAG: Keyword[class]/None: class{{; name=.+$}}
 // KW_DECL-DAG: Keyword/None: convenience{{; name=.+$}}
@@ -121,7 +161,6 @@
 // KW_DECL_STMT-DAG: Keyword[for]/None: for{{; name=.+$}}
 // KW_DECL_STMT-DAG: Keyword[in]/None: in{{; name=.+$}}
 // KW_DECL_STMT-DAG: Keyword[while]/None: while{{; name=.+$}}
-// KW_DECL_STMT-DAG: Keyword[return]/None: return{{; name=.+$}}
 // KW_DECL_STMT-DAG: Keyword[break]/None: break{{; name=.+$}}
 // KW_DECL_STMT-DAG: Keyword[continue]/None: continue{{; name=.+$}}
 // KW_DECL_STMT-DAG: Keyword[fallthrough]/None: fallthrough{{; name=.+$}}
@@ -194,6 +233,10 @@
 
 #^TOP_LEVEL_1^#
 
+for _ in 1...10 {
+  #^TOP_LEVEL_2^#
+}
+
 func testInFuncBody1() {
   #^IN_FUNC_BODY_1^#
 }
@@ -224,6 +267,25 @@ class InClassFunc {
   }
 }
 
+func testInClosure1() {
+  { #^IN_CLOSURE_1^# }
+}
+func testInClosure2() {
+  { #^IN_CLOSURE_2^#
+}
+struct InVarClosureInit {
+  let x = { #^IN_CLOSURE_3^# }()
+}
+
+{ #^IN_CLOSURE_4^# }
+
+struct InSubscript {
+  subscript(x: Int) -> Int { #^IN_SUBSCRIPT_1^# }
+}
+
+struct InInit {
+  init?() { #^IN_INIT_1^# }
+}
 
 struct InStruct {
   #^IN_NOMINAL_DECL_1^#

test/SourceKit/CodeComplete/complete_requestlimit.swift

@@ -23,6 +23,7 @@ func test001() {
 // TOP_LEVEL_0_ALL-NEXT: if
 // TOP_LEVEL_0_ALL-NEXT: for
 // TOP_LEVEL_0_ALL-NEXT: while
+// TOP_LEVEL_0_ALL-NEXT: return
 // TOP_LEVEL_0_ALL-NEXT: func
 // TOP_LEVEL_0_ALL-NEXT: x
 // TOP_LEVEL_0_ALL-NEXT: y

test/SourceKit/CodeComplete/complete_sort_order.swift

@@ -52,6 +52,7 @@ func test1() {
 // STMT: if
 // STMT: for
 // STMT: while
+// STMT: return
 // STMT: func
 // STMT: foo(a: Int)
 
@@ -61,14 +62,14 @@ func test5() {
   #^STMT_1,r,ret,retur,return^#
 }
 // STMT_1-LABEL: Results for filterText: r [
+// STMT_1-NEXT:    return
 // STMT_1-NEXT:    retLocal
 // STMT_1-NEXT:    repeat
-// STMT_1-NEXT:    return
 // STMT_1-NEXT:    required
 // STMT_1: ]
 // STMT_1-LABEL: Results for filterText: ret [
-// STMT_1-NEXT:    retLocal
 // STMT_1-NEXT:    return
+// STMT_1-NEXT:    retLocal
 // STMT_1-NEXT:    repeat
 // STMT_1: ]
 // STMT_1-LABEL: Results for filterText: retur [

tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp

@@ -439,6 +439,7 @@ static bool isHighPriorityKeyword(CodeCompletionKeywordKind kind) {
   case CodeCompletionKeywordKind::kw_if:
   case CodeCompletionKeywordKind::kw_for:
   case CodeCompletionKeywordKind::kw_while:
+  case CodeCompletionKeywordKind::kw_return:
   case CodeCompletionKeywordKind::kw_func:
     return true;
   default:
@@ -739,6 +740,7 @@ static int compareHighPriorityKeywords(Item &a_, Item &b_) {
     CodeCompletionKeywordKind::kw_if,
     CodeCompletionKeywordKind::kw_for,
     CodeCompletionKeywordKind::kw_while,
+    CodeCompletionKeywordKind::kw_return,
     CodeCompletionKeywordKind::kw_func,
   };
   auto size = sizeof(order) / sizeof(order[0]);