[CodeCompletion] Complete pattern introducer for 'for'

After 'for', suggest 'try', 'await', 'var' and 'case'.

rdar://76355581

Author: rintaro

include/swift/IDE/CodeCompletion.h

@@ -561,6 +561,7 @@ enum class CompletionKind {
   GenericRequirement,
   PrecedenceGroup,
   StmtLabel,
+  ForEachPatternBeginning,
 };
 
 /// A single code completion result.

include/swift/Parse/CodeCompletionCallbacks.h

@@ -231,6 +231,9 @@ class CodeCompletionCallbacks {
 
   virtual void completeStmtLabel(StmtKind ParentKind) {};
 
+  virtual
+  void completeForEachPatternBeginning(bool hasTry, bool hasAwait) {};
+
   /// Signals that the AST for the all the delayed-parsed code was
   /// constructed.  No \c complete*() callbacks will be done after this.
   virtual void doneParsing() = 0;

lib/IDE/CodeCompletion.cpp

@@ -1669,6 +1669,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   void completeGenericRequirement() override;
   void completeAfterIfStmt(bool hasElse) override;
   void completeStmtLabel(StmtKind ParentKind) override;
+  void completeForEachPatternBeginning(bool hasTry, bool hasAwait) override;
 
   void doneParsing() override;
 
@@ -5820,6 +5821,17 @@ void CodeCompletionCallbacksImpl::completeStmtLabel(StmtKind ParentKind) {
   ParentStmtKind = ParentKind;
 }
 
+void CodeCompletionCallbacksImpl::completeForEachPatternBeginning(
+    bool hasTry, bool hasAwait) {
+  CurDeclContext = P.CurDeclContext;
+  Kind = CompletionKind::ForEachPatternBeginning;
+  ParsedKeywords.clear();
+  if (hasTry)
+    ParsedKeywords.emplace_back("try");
+  if (hasAwait)
+    ParsedKeywords.emplace_back("await");
+}
+
 static bool isDynamicLookup(Type T) {
   return T->getRValueType()->isAnyObject();
 }
@@ -6055,6 +6067,13 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
   case CompletionKind::AfterIfStmtElse:
     addKeyword(Sink, "if", CodeCompletionKeywordKind::kw_if);
     break;
+  case CompletionKind::ForEachPatternBeginning:
+    if (!llvm::is_contained(ParsedKeywords, "try"))
+      addKeyword(Sink, "try", CodeCompletionKeywordKind::kw_try);
+    if (!llvm::is_contained(ParsedKeywords, "await"))
+      addKeyword(Sink, "await", CodeCompletionKeywordKind::None);
+    addKeyword(Sink, "var", CodeCompletionKeywordKind::kw_var);
+    addKeyword(Sink, "case", CodeCompletionKeywordKind::kw_case);
   }
 }
 
@@ -6945,6 +6964,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::AfterIfStmtElse:
   case CompletionKind::CaseStmtKeyword:
   case CompletionKind::EffectsSpecifier:
+  case CompletionKind::ForEachPatternBeginning:
     // Handled earlier by keyword completions.
     break;
   }

lib/Parse/ParseStmt.cpp

@@ -2136,6 +2136,17 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
     }
   }
 
+  if (Tok.is(tok::code_complete)) {
+    if (CodeCompletion) {
+      CodeCompletion->completeForEachPatternBeginning(TryLoc.isValid(),
+                                                      AwaitLoc.isValid());
+    }
+    consumeToken(tok::code_complete);
+    // Since 'completeForeachPatternBeginning' is a keyword only completion,
+    // we don't need to parse the rest of 'for' statement.
+    return makeParserCodeCompletionStatus();
+  }
+
   // Parse the pattern.  This is either 'case <refutable pattern>' or just a
   // normal pattern.
   if (consumeIf(tok::kw_case)) {

test/IDE/complete_stmt_controlling_expr.swift

@@ -142,22 +142,43 @@ func testRepeatWhile2(_ fooObject: FooStruct) {
   } while localFooObject.#^COND_DO_WHILE_2?check=COND-WITH-RELATION1^#
 }
 
-func testCStyleForInit1(_ fooObject: FooStruct) {
+func testForeachPattern1(_ fooObject: FooStruct) {
   var localInt = 42
   var localFooObject = FooStruct(localInt)
-  for #^C_STYLE_FOR_INIT_1?check=COND_NONE^#
+  for #^FOREACH_PATTERN_1^#
+// FOREACH_PATTERN_1: Begin completions, 4 items
+// FOREACH_PATTERN_1-DAG: Keyword[try]/None:                  try; name=try
+// FOREACH_PATTERN_1-DAG: Keyword/None:                       await; name=await
+// FOREACH_PATTERN_1-DAG: Keyword[var]/None:                  var; name=var
+// FOREACH_PATTERN_1-DAG: Keyword[case]/None:                 case; name=case
+// FOREACH_PATTERN_1: End completions
 }
 
-func testCStyleForInit2(_ fooObject: FooStruct) {
+func testForeachPattern2(_ fooObject: FooStruct) {
   var localInt = 42
   var localFooObject = FooStruct(localInt)
-  for #^C_STYLE_FOR_INIT_2?check=COND_COMMON^#;
+  for try #^FOREACH_PATTERN_2^#
+// FOREACH_PATTERN_2: Begin completions, 3 items
+// FOREACH_PATTERN_2-DAG: Keyword/None:                       await; name=await
+// FOREACH_PATTERN_2-DAG: Keyword[var]/None:                  var; name=var
+// FOREACH_PATTERN_2-DAG: Keyword[case]/None:                 case; name=case
+// FOREACH_PATTERN_2: End completions
 }
 
-func testCStyleForInit3(_ fooObject: FooStruct) {
+func testForeachPattern3(_ fooObject: FooStruct) {
   var localInt = 42
   var localFooObject = FooStruct(localInt)
-  for #^C_STYLE_FOR_INIT_3?check=COND_COMMON^# ;
+  for try await #^FOREACH_PATTERN_3^#
+// FOREACH_PATTERN_3: Begin completions, 2 items
+// FOREACH_PATTERN_3-DAG: Keyword[var]/None:                  var; name=var
+// FOREACH_PATTERN_3-DAG: Keyword[case]/None:                 case; name=case
+// FOREACH_PATTERN_3: End completions
+}
+
+func testForeachPattern4(_ fooObject: FooStruct) {
+  var localInt = 42
+  var localFooObject = FooStruct(localInt)
+  for var #^FOREACH_PATTERN_4?check=COND_NONE^#
 }
 
 func testCStyleForCond1(_ fooObject: FooStruct) {