[codecomplete] Add completion of if after else

At some point we stopped completing keywords after `else` in an if
statement. Bring back the completion of `if`, which is the only valid
continuation other than a brace.

rdar://37467474

Author: benlangmuir

include/swift/IDE/CodeCompletion.h

@@ -503,6 +503,7 @@ enum class CompletionKind {
   ReturnStmtExpr,
   ForEachSequence,
   AfterPound,
+  AfterIfStmtElse,
   GenericParams,
   SwiftKeyPath,
 };

include/swift/Parse/CodeCompletionCallbacks.h

@@ -202,6 +202,8 @@ class CodeCompletionCallbacks {
 
   virtual void completeAfterPound(CodeCompletionExpr *E, StmtKind ParentKind) = 0;
 
+  virtual void completeAfterIfStmt(bool hasElse) = 0;
+
   virtual void completeGenericParams(TypeLoc TL) = 0;
 
   /// \brief Signals that the AST for the all the delayed-parsed code was

lib/IDE/CodeCompletion.cpp

@@ -1467,6 +1467,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   void completeReturnStmt(CodeCompletionExpr *E) override;
   void completeAfterPound(CodeCompletionExpr *E, StmtKind ParentKind) override;
   void completeGenericParams(TypeLoc TL) override;
+  void completeAfterIfStmt(bool hasElse) override;
   void addKeywords(CodeCompletionResultSink &Sink, bool MaybeFuncBody);
 
   void doneParsing() override;
@@ -4745,6 +4746,15 @@ void CodeCompletionCallbacksImpl::completeAfterPound(CodeCompletionExpr *E,
   ParentStmtKind = ParentKind;
 }
 
+void CodeCompletionCallbacksImpl::completeAfterIfStmt(bool hasElse) {
+  CurDeclContext = P.CurDeclContext;
+  if (hasElse) {
+    Kind = CompletionKind::AfterIfStmtElse;
+  } else {
+    Kind = CompletionKind::StmtOrExpr;
+  }
+}
+
 void CodeCompletionCallbacksImpl::completeGenericParams(TypeLoc TL) {
   CurDeclContext = P.CurDeclContext;
   Kind = CompletionKind::GenericParams;
@@ -4898,6 +4908,10 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
     addDeclKeywords(Sink);
     addLetVarKeywords(Sink);
     break;
+
+  case CompletionKind::AfterIfStmtElse:
+    addKeyword(Sink, "if", CodeCompletionKeywordKind::kw_if);
+    break;
   }
 }
 
@@ -5474,6 +5488,9 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       }
     }
     break;
+  case CompletionKind::AfterIfStmtElse:
+    // Handled earlier by keyword completions.
+    break;
   }
 
   if (Lookup.RequestedCachedResults) {

lib/Parse/ParseStmt.cpp

@@ -1545,10 +1545,20 @@ ParserResult<Stmt> Parser::parseStmtIf(LabeledStmtInfo LabelInfo) {
     if (Tok.is(tok::kw_if)) {
       SyntaxParsingContext ElseIfCtxt(SyntaxContext, SyntaxKind::IfStmt);
       ElseBody = parseStmtIf(LabeledStmtInfo());
+    } else if (Tok.is(tok::code_complete)) {
+      if (CodeCompletion)
+        CodeCompletion->completeAfterIfStmt(/*hasElse*/true);
+      Status.setHasCodeCompletion();
+      consumeToken(tok::code_complete);
     } else {
       ElseBody = parseBraceItemList(diag::expected_lbrace_after_else);
     }
     Status |= ElseBody;
+  } else if (Tok.is(tok::code_complete)) {
+    if (CodeCompletion)
+      CodeCompletion->completeAfterIfStmt(/*hasElse*/false);
+    Status.setHasCodeCompletion();
+    consumeToken(tok::code_complete);
   }
 
   return makeParserResult(

test/IDE/complete_keywords.swift

@@ -6,6 +6,18 @@
 // 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=TOP_LEVEL_AFTER_IF_1 > %t.top3
+// RUN: %FileCheck %s -check-prefix=KW_DECL_STMT < %t.top3
+// RUN: %FileCheck %s -check-prefix=KW_NO_RETURN < %t.top3
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TOP_LEVEL_AFTER_IF_ELSE_1 | %FileCheck %s -check-prefix=AFTER_IF_ELSE
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=AFTER_IF_1 > %t.if1
+// RUN: %FileCheck %s -check-prefix=KW_DECL_STMT < %t.if1
+// RUN: %FileCheck %s -check-prefix=KW_RETURN < %t.if1
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=AFTER_IF_ELSE_1 | %FileCheck %s -check-prefix=AFTER_IF_ELSE
+
 // 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
@@ -237,6 +249,20 @@ for _ in 1...10 {
   #^TOP_LEVEL_2^#
 }
 
+if true {} #^TOP_LEVEL_AFTER_IF_1^#
+
+if true {} else #^TOP_LEVEL_AFTER_IF_ELSE_1^# {}
+
+// AFTER_IF_ELSE: Begin completions, 1 items
+// AFTER_IF_ELSE: Keyword[if]/None: if;
+
+func testAfterIf1() {
+  if true {} #^AFTER_IF_1^#
+}
+func testAfterIfElse1() {
+  if true {} else #^AFTER_IF_ELSE_1^# {}
+}
+
 func testInFuncBody1() {
   #^IN_FUNC_BODY_1^#
 }