[CodeCompletion] Add a new StmtOrExpr code completion kind

We want to distinguish keywords that are only valid at
statement/declaration context from those valid in any expression.

include/swift/IDE/CodeCompletion.h

@@ -421,6 +421,7 @@ enum class CompletionKind {
   Import,
   UnresolvedMember,
   DotExpr,
+  StmtOrExpr,
   PostfixExprBeginning,
   PostfixExpr,
   PostfixExprParen,

include/swift/Parse/CodeCompletionCallbacks.h

@@ -102,6 +102,9 @@ class CodeCompletionCallbacks {
   /// \brief Complete expr-dot after we have consumed the dot.
   virtual void completeDotExpr(Expr *E, SourceLoc DotLoc) = 0;
 
+  /// \brief Complete the beginning of a statement or expression.
+  virtual void completeStmtOrExpr() = 0;
+
   /// \brief Complete the beginning of expr-postfix -- no tokens provided
   /// by user.
   virtual void completePostfixExprBeginning(CodeCompletionExpr *E) = 0;

lib/IDE/CodeCompletion.cpp

@@ -808,7 +808,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   CodeCompletionContext &CompletionContext;
   std::vector<RequestedCachedModule> RequestedModules;
   CodeCompletionConsumer &Consumer;
-  CodeCompletionExpr *CodeCompleteTokenExpr;
+  CodeCompletionExpr *CodeCompleteTokenExpr = nullptr;
   AssignExpr *AssignmentExpr;
   CallExpr *FuncCallExpr;
   UnresolvedMemberExpr *UnresolvedExpr;
@@ -950,6 +950,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
 
   void completeExpr() override;
   void completeDotExpr(Expr *E, SourceLoc DotLoc) override;
+  void completeStmtOrExpr() override;
   void completePostfixExprBeginning(CodeCompletionExpr *E) override;
   void completePostfixExpr(Expr *E, bool hasSpace) override;
   void completePostfixExprParen(Expr *E, Expr *CodeCompletionE) override;
@@ -3456,6 +3457,14 @@ void CodeCompletionCallbacksImpl::completeDotExpr(Expr *E, SourceLoc DotLoc) {
   CurDeclContext = P.CurDeclContext;
 }
 
+void CodeCompletionCallbacksImpl::completeStmtOrExpr() {
+  assert(P.Tok.is(tok::code_complete));
+  Kind = CompletionKind::StmtOrExpr;
+  CurDeclContext = P.CurDeclContext;
+  CStyleForLoopIterationVariable =
+      CodeCompletionCallbacks::CStyleForLoopIterationVariable;
+}
+
 void CodeCompletionCallbacksImpl::completePostfixExprBeginning(CodeCompletionExpr *E) {
   assert(P.Tok.is(tok::code_complete));
 
@@ -3741,6 +3750,7 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink) {
   case CompletionKind::AfterPound:
     break;
 
+  case CompletionKind::StmtOrExpr:
   case CompletionKind::PostfixExprBeginning:
     addSuperKeyword(Sink);
     addDeclKeywords(Sink);
@@ -4060,6 +4070,10 @@ void CodeCompletionCallbacksImpl::doneParsing() {
     break;
   }
 
+  case CompletionKind::StmtOrExpr:
+    DoPostfixExprBeginning();
+    break;
+
   case CompletionKind::PostfixExprBeginning: {
     CodeCompletionTypeContextAnalyzer Analyzer(CurDeclContext,
                                                CodeCompleteTokenExpr);

lib/Parse/ParseStmt.cpp

@@ -97,6 +97,13 @@ ParserStatus Parser::parseExprOrStmt(ASTNode &Result) {
   if (CodeCompletion)
     CodeCompletion->setExprBeginning(getParserPosition());
 
+  if (Tok.is(tok::code_complete)) {
+    if (CodeCompletion)
+      CodeCompletion->completeStmtOrExpr();
+    consumeToken(tok::code_complete);
+    return makeParserCodeCompletionStatus();
+  }
+
   ParserResult<Expr> ResultExpr = parseExpr(diag::expected_expr);
   if (ResultExpr.isNonNull()) {
     Result = ResultExpr.get();