Merge pull request #60207 from simanerush/code-completion-foreach

Author: ahoppen

include/swift/IDE/CodeCompletionResult.h

@@ -214,6 +214,9 @@ enum class CompletionKind : uint8_t {
   ReturnStmtExpr,
   YieldStmtExpr,
   ForEachSequence,
+
+  /// The \c in keyword in a for-each loop.
+  ForEachInKw,
   AfterPoundExpr,
   AfterPoundDirective,
   PlatformConditon,

include/swift/Parse/CodeCompletionCallbacks.h

@@ -137,6 +137,9 @@ class CodeCompletionCallbacks {
   /// -- no tokens provided by user.
   virtual void completeForEachSequenceBeginning(CodeCompletionExpr *E) {};
 
+  /// Complete the \c in keyword in a for-each loop.
+  virtual void completeForEachInKeyword(){};
+
   /// Complete a given expr-postfix.
   virtual void completePostfixExpr(Expr *E, bool hasSpace) {};
 

lib/IDE/CodeCompletion.cpp

@@ -242,6 +242,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   void completeStmtOrExpr(CodeCompletionExpr *E) override;
   void completePostfixExprBeginning(CodeCompletionExpr *E) override;
   void completeForEachSequenceBeginning(CodeCompletionExpr *E) override;
+  void completeForEachInKeyword() override;
   void completePostfixExpr(Expr *E, bool hasSpace) override;
   void completePostfixExprParen(Expr *E, Expr *CodeCompletionE) override;
   void completeExprKeyPath(KeyPathExpr *KPE, SourceLoc DotLoc) override;
@@ -365,6 +366,12 @@ void CodeCompletionCallbacksImpl::completeForEachSequenceBeginning(
   CodeCompleteTokenExpr = E;
 }
 
+void CodeCompletionCallbacksImpl::completeForEachInKeyword() {
+  assert(P.Tok.is(tok::code_complete));
+  Kind = CompletionKind::ForEachInKw;
+  CurDeclContext = P.CurDeclContext;
+}
+
 void CodeCompletionCallbacksImpl::completePostfixExpr(Expr *E, bool hasSpace) {
   assert(P.Tok.is(tok::code_complete));
 
@@ -1069,6 +1076,9 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
       addKeyword(Sink, "await", CodeCompletionKeywordKind::None);
     addKeyword(Sink, "var", CodeCompletionKeywordKind::kw_var);
     addKeyword(Sink, "case", CodeCompletionKeywordKind::kw_case);
+    break;
+  case CompletionKind::ForEachInKw:
+    addKeyword(Sink, "in", CodeCompletionKeywordKind::kw_in);
   }
 }
 
@@ -1960,6 +1970,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::CaseStmtKeyword:
   case CompletionKind::EffectsSpecifier:
   case CompletionKind::ForEachPatternBeginning:
+  case CompletionKind::ForEachInKw:
     // Handled earlier by keyword completions.
     break;
   }

lib/Parse/ParseStmt.cpp

@@ -2292,14 +2292,23 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
     diagnose(LBraceLoc, diag::expected_foreach_container);
     Container = makeParserErrorResult(new (Context) ErrorExpr(LBraceLoc));
   } else if (Tok.is(tok::code_complete)) {
-    Container =
-        makeParserResult(new (Context) CodeCompletionExpr(Tok.getLoc()));
-    Container.setHasCodeCompletionAndIsError();
-    Status |= Container;
-    if (CodeCompletion)
-      CodeCompletion->completeForEachSequenceBeginning(
-          cast<CodeCompletionExpr>(Container.get()));
-    consumeToken(tok::code_complete);
+    // If there is no "in" keyword, suggest it. Otherwise, complete the
+    // sequence.
+    if (InLoc.isInvalid()) {
+      if (CodeCompletion)
+        CodeCompletion->completeForEachInKeyword();
+      consumeToken(tok::code_complete);
+      return makeParserCodeCompletionStatus();
+    } else {
+      Container =
+          makeParserResult(new (Context) CodeCompletionExpr(Tok.getLoc()));
+      Container.setHasCodeCompletionAndIsError();
+      Status |= Container;
+      if (CodeCompletion)
+        CodeCompletion->completeForEachSequenceBeginning(
+            cast<CodeCompletionExpr>(Container.get()));
+      consumeToken(tok::code_complete);
+    }
   } else {
     Container = parseExprBasic(diag::expected_foreach_container);
     Status |= Container;

test/IDE/complete_loop.swift

@@ -5,6 +5,8 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LOOP_3 | %FileCheck %s -check-prefix=LOOP_3
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LOOP_4 | %FileCheck %s -check-prefix=LOOP_4
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LOOP_5 | %FileCheck %s -check-prefix=LOOP_5
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LOOP_6 | %FileCheck %s -check-prefix=LOOP_6
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LOOP_7 | %FileCheck %s -check-prefix=LOOP_6
 
 class Gen {
 	func IntGen() -> Int { return 0 }
@@ -67,3 +69,14 @@ class C {
   }
 // LOOP_5: Begin completions
 }
+
+// https://github.com/apple/swift/issues/58633
+do {
+  for value #^LOOP_6^#
+}
+do {
+  for value #^LOOP_7^# 1..<7 {}
+}
+// LOOP_6: Begin completions, 1 items
+// LOOP_6-CHECK-NEXT: Keyword[in]/None:                   in; name=in
+// LOOP_6-CHECK-NEXT: End completions