[code-completion] Add a new custom completion context for a for-each sequence

include/swift/IDE/CodeCompletion.h

@@ -501,6 +501,7 @@ enum class CompletionKind {
   AssignmentRHS,
   CallArg,
   ReturnStmtExpr,
+  ForEachSequence,
   AfterPound,
   GenericParams,
   SwiftKeyPath,

include/swift/Parse/CodeCompletionCallbacks.h

@@ -151,6 +151,10 @@ class CodeCompletionCallbacks {
   /// by user.
   virtual void completePostfixExprBeginning(CodeCompletionExpr *E) = 0;
 
+  /// \brief Complete the beginning of expr-postfix in a for-each loop sequqence
+  /// -- no tokens provided by user.
+  virtual void completeForEachSequenceBeginning(CodeCompletionExpr *E) = 0;
+
   /// \brief Complete a given expr-postfix.
   virtual void completePostfixExpr(Expr *E, bool hasSpace) = 0;
 

lib/IDE/CodeCompletion.cpp

@@ -1440,6 +1440,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   void completeDotExpr(Expr *E, SourceLoc DotLoc) override;
   void completeStmtOrExpr() override;
   void completePostfixExprBeginning(CodeCompletionExpr *E) override;
+  void completeForEachSequenceBeginning(CodeCompletionExpr *E) override;
   void completePostfixExpr(Expr *E, bool hasSpace) override;
   void completePostfixExprParen(Expr *E, Expr *CodeCompletionE) override;
   void completeExprSuper(SuperRefExpr *SRE) override;
@@ -4479,6 +4480,14 @@ void CodeCompletionCallbacksImpl::completePostfixExprBeginning(CodeCompletionExp
   CodeCompleteTokenExpr = E;
 }
 
+void CodeCompletionCallbacksImpl::completeForEachSequenceBeginning(
+    CodeCompletionExpr *E) {
+  assert(P.Tok.is(tok::code_complete));
+  Kind = CompletionKind::ForEachSequence;
+  CurDeclContext = P.CurDeclContext;
+  CodeCompleteTokenExpr = E;
+}
+
 void CodeCompletionCallbacksImpl::completePostfixExpr(Expr *E, bool hasSpace) {
   assert(P.Tok.is(tok::code_complete));
 
@@ -4840,6 +4849,7 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
   case CompletionKind::AssignmentRHS:
   case CompletionKind::ReturnStmtExpr:
   case CompletionKind::PostfixExprBeginning:
+  case CompletionKind::ForEachSequence:
     addSuperKeyword(Sink);
     addLetVarKeywords(Sink);
     addExprKeywords(Sink);
@@ -5191,6 +5201,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
     DoPostfixExprBeginning();
     break;
 
+  case CompletionKind::ForEachSequence:
   case CompletionKind::PostfixExprBeginning: {
     ::CodeCompletionTypeContextAnalyzer Analyzer(CurDeclContext,
                                                CodeCompleteTokenExpr);

lib/Parse/ParseStmt.cpp

@@ -2098,6 +2098,15 @@ ParserResult<Stmt> Parser::parseStmtForEach(SourceLoc ForLoc,
     SourceLoc LBraceLoc = Tok.getLoc();
     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.setHasCodeCompletion();
+    Status |= Container;
+    if (CodeCompletion)
+      CodeCompletion->completeForEachSequenceBeginning(
+          cast<CodeCompletionExpr>(Container.get()));
+    consumeToken(tok::code_complete);
   } else {
     Container = parseExprBasic(diag::expected_foreach_container);
     if (Container.isNull())

test/SourceKit/CodeComplete/Inputs/custom-completion/custom.json

@@ -24,6 +24,13 @@
         source.lang.swift.stmt,
         source.lang.swift.type
       ]
+    },
+    {
+      key.name: "customForEach",
+      key.kind: myuid.customForEach,
+      key.context: [
+        source.lang.swift.foreach.sequence,
+      ]
     }
   ]
 }

test/SourceKit/CodeComplete/complete_custom.swift

@@ -2,6 +2,7 @@ func test() {
   // stmt
   ()
   let foo: // type
+  for x in  { } // foreach.sequence
 }
 
 // ===--- Errors
@@ -62,6 +63,14 @@ func test() {
 // TYPE-NEXT: key.name: "customType"
 // TYPE-NOT: myuid
 
+// RUN: %sourcekitd-test -json-request-path %S/Inputs/custom-completion/custom.json == \
+// RUN:     -req=complete.open -pos=5:12 %s -- %s | %FileCheck %s -check-prefix=FOREACH
+
+// FOREACH-NOT: myuid
+// FOREACH: myuid.customForEach
+// FOREACH-NEXT: key.name: "customForEach"
+// FOREACH-NOT: myuid
+
 // ===--- Filtering.
 
 // RUN: %sourcekitd-test -json-request-path %S/Inputs/custom-completion/custom.json == \

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -141,6 +141,7 @@ struct CustomCompletionInfo {
     Stmt = 1 << 0,
     Expr = 1 << 1,
     Type = 1 << 2,
+    ForEachSequence = 1 << 3,
   };
   swift::OptionSet<Context> Contexts;
 };

tools/SourceKit/include/SourceKit/Core/ProtocolUIDs.def

@@ -276,6 +276,7 @@ KIND(ObjectLiteral, "source.lang.swift.syntaxtype.objectliteral")
 KIND(Expr, "source.lang.swift.expr")
 KIND(Stmt, "source.lang.swift.stmt")
 KIND(Type, "source.lang.swift.type")
+KIND(ForEachSequence, "source.lang.swift.foreach.sequence")
 
 KIND(DiagNote, "source.diagnostic.severity.note")
 KIND(DiagWarning, "source.diagnostic.severity.warning")

tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp

@@ -184,6 +184,12 @@ bool SourceKit::CodeCompletion::addCustomCompletions(
         addCompletion(custom);
       }
       break;
+    case CompletionKind::ForEachSequence:
+      if (custom.Contexts.contains(CustomCompletionInfo::ForEachSequence)) {
+        changed = true;
+        addCompletion(custom);
+      }
+      break;
     case CompletionKind::TypeSimpleBeginning:
       if (custom.Contexts.contains(CustomCompletionInfo::Type)) {
         changed = true;

tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp

@@ -593,6 +593,8 @@ void handleRequestImpl(sourcekitd_object_t ReqObj, ResponseReceiver Rec) {
           CCInfo.Contexts |= CustomCompletionInfo::Stmt;
         } else if (context == KindType) {
           CCInfo.Contexts |= CustomCompletionInfo::Type;
+        } else if (context == KindForEachSequence) {
+          CCInfo.Contexts |= CustomCompletionInfo::ForEachSequence;
         } else {
           err = createErrorRequestInvalid("invalid value for 'key.context'");
           return true;