[CodeCompletion] Completion for control trasfer target

Completion after 'break' and 'continue'

rdar://problem/57016218

Author: rintaro

include/swift/IDE/CodeCompletion.h

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

include/swift/Parse/CodeCompletionCallbacks.h

@@ -219,6 +219,8 @@ class CodeCompletionCallbacks {
 
   virtual void completeGenericRequirement() {};
 
+  virtual void completeStmtLabel(StmtKind ParentKind) {};
+
   /// 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

@@ -1580,6 +1580,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   void completePlatformCondition() override;
   void completeGenericRequirement() override;
   void completeAfterIfStmt(bool hasElse) override;
+  void completeStmtLabel(StmtKind ParentKind) override;
 
   void doneParsing() override;
 
@@ -4350,6 +4351,52 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     for (auto PGD: precedenceGroups)
       addPrecedenceGroupRef(PGD);
   }
+
+  void getStmtLabelCompletions(SourceLoc Loc, bool isContinue) {
+    class LabelFinder : public ASTWalker {
+      SourceManager &SM;
+      SourceLoc TargetLoc;
+      bool IsContinue;
+
+    public:
+      SmallVector<Identifier, 2> Result;
+
+      LabelFinder(SourceManager &SM, SourceLoc TargetLoc, bool IsContinue)
+          : SM(SM), TargetLoc(TargetLoc), IsContinue(IsContinue) {}
+
+      std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+        if (SM.isBeforeInBuffer(S->getEndLoc(), TargetLoc))
+          return {false, S};
+
+        if (LabeledStmt *LS = dyn_cast<LabeledStmt>(S)) {
+          if (LS->getLabelInfo()) {
+            if (!IsContinue || LS->isPossibleContinueTarget()) {
+              auto label = LS->getLabelInfo().Name;
+              if (!llvm::is_contained(Result, label))
+                Result.push_back(label);
+            }
+          }
+        }
+
+        return {true, S};
+      }
+
+      Stmt *walkToStmtPost(Stmt *S) override { return nullptr; }
+
+      std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+        if (SM.isBeforeInBuffer(E->getEndLoc(), TargetLoc))
+          return {false, E};
+        return {true, E};
+      }
+    } Finder(CurrDeclContext->getASTContext().SourceMgr, Loc, isContinue);
+    const_cast<DeclContext *>(CurrDeclContext)->walkContext(Finder);
+    for (auto name : Finder.Result) {
+      CodeCompletionResultBuilder Builder(
+          Sink, CodeCompletionResult::ResultKind::Pattern,
+          SemanticContextKind::Local, {});
+      Builder.addTextChunk(name.str());
+    }
+  }
 };
 
 class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
@@ -5133,6 +5180,12 @@ void CodeCompletionCallbacksImpl::completeAccessorBeginning(
   CodeCompleteTokenExpr = E;
 }
 
+void CodeCompletionCallbacksImpl::completeStmtLabel(StmtKind ParentKind) {
+  CurDeclContext = P.CurDeclContext;
+  Kind = CompletionKind::StmtLabel;
+  ParentStmtKind = ParentKind;
+}
+
 static bool isDynamicLookup(Type T) {
   return T->getRValueType()->isAnyObject();
 }
@@ -5252,6 +5305,7 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
   case CompletionKind::KeyPathExprObjC:
   case CompletionKind::KeyPathExprSwift:
   case CompletionKind::PrecedenceGroup:
+  case CompletionKind::StmtLabel:
     break;
 
   case CompletionKind::AccessorBeginning: {
@@ -5850,6 +5904,11 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::PrecedenceGroup:
     Lookup.getPrecedenceGroupCompletions(SyntxKind);
     break;
+  case CompletionKind::StmtLabel: {
+    SourceLoc Loc = P.Context.SourceMgr.getCodeCompletionLoc();
+    Lookup.getStmtLabelCompletions(Loc, ParentStmtKind == StmtKind::Continue);
+    break;
+  }
   case CompletionKind::AfterIfStmtElse:
   case CompletionKind::CaseStmtKeyword:
     // Handled earlier by keyword completions.

lib/Parse/ParseStmt.cpp

@@ -679,6 +679,30 @@ ParserResult<BraceStmt> Parser::parseBraceItemList(Diag<> ID) {
                           BraceStmt::create(Context, LBLoc, Entries, RBLoc));
 }
 
+static ParserStatus parseOptionalControlTransferTarget(Parser &P,
+                                                       Identifier &Target,
+                                                       SourceLoc &TargetLoc,
+                                                       StmtKind Kind) {
+  // If we have an identifier after 'break' or 'continue', which is not the
+  // start of another stmt or decl, we assume it is the label to break to,
+  // unless there is a line break.  There is ambiguity with expressions (e.g.
+  // "break x+y") but since the expression after the them is dead, we don't feel
+  // bad eagerly parsing this.
+  if (!P.Tok.isAtStartOfLine()) {
+    if (P.Tok.is(tok::identifier) && !P.isStartOfStmt() &&
+        !P.isStartOfSwiftDecl()) {
+      TargetLoc = P.consumeIdentifier(&Target);
+      return makeParserSuccess();
+    } else if (P.Tok.is(tok::code_complete)) {
+      if (P.CodeCompletion)
+        P.CodeCompletion->completeStmtLabel(Kind);
+      TargetLoc = P.consumeToken(tok::code_complete);
+      return makeParserCodeCompletionStatus();
+    }
+  }
+  return makeParserSuccess();
+}
+
 /// parseStmtBreak
 ///
 ///   stmt-break:
@@ -689,17 +713,12 @@ ParserResult<Stmt> Parser::parseStmtBreak() {
   SourceLoc Loc = consumeToken(tok::kw_break);
   SourceLoc TargetLoc;
   Identifier Target;
+  ParserStatus Status;
+  Status |= parseOptionalControlTransferTarget(*this, Target, TargetLoc,
+                                               StmtKind::Break);
 
-  // If we have an identifier after this, which is not the start of another
-  // stmt or decl, we assume it is the label to break to, unless there is a
-  // line break.  There is ambiguity with expressions (e.g. "break x+y") but
-  // since the expression after the break is dead, we don't feel bad eagerly
-  // parsing this.
-  if (Tok.is(tok::identifier) && !Tok.isAtStartOfLine() &&
-      !isStartOfStmt() && !isStartOfSwiftDecl())
-    TargetLoc = consumeIdentifier(&Target);
-
-  return makeParserResult(new (Context) BreakStmt(Loc, Target, TargetLoc));
+  return makeParserResult(Status,
+                          new (Context) BreakStmt(Loc, Target, TargetLoc));
 }
 
 /// parseStmtContinue
@@ -712,17 +731,12 @@ ParserResult<Stmt> Parser::parseStmtContinue() {
   SourceLoc Loc = consumeToken(tok::kw_continue);
   SourceLoc TargetLoc;
   Identifier Target;
+  ParserStatus Status;
+  Status |= parseOptionalControlTransferTarget(*this, Target, TargetLoc,
+                                               StmtKind::Continue);
 
-  // If we have an identifier after this, which is not the start of another
-  // stmt or decl, we assume it is the label to continue to, unless there is a
-  // line break.  There is ambiguity with expressions (e.g. "continue x+y") but
-  // since the expression after the continue is dead, we don't feel bad eagerly
-  // parsing this.
-  if (Tok.is(tok::identifier) && !Tok.isAtStartOfLine() &&
-      !isStartOfStmt() && !isStartOfSwiftDecl())
-    TargetLoc = consumeIdentifier(&Target);
-
-  return makeParserResult(new (Context) ContinueStmt(Loc, Target, TargetLoc));
+  return makeParserResult(Status,
+                          new (Context) ContinueStmt(Loc, Target, TargetLoc));
 }
 
 

test/IDE/complete_controltransfer.swift

@@ -0,0 +1,80 @@
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LABEL_1 | %FileCheck %s -check-prefix=LABEL_1
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LABEL_2 | %FileCheck %s -check-prefix=LABEL_2
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LABEL_3 | %FileCheck %s -check-prefix=LABEL_3
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LABEL_4 | %FileCheck %s -check-prefix=LABEL_4
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LABEL_5 | %FileCheck %s -check-prefix=LABEL_5
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LABEL_6 | %FileCheck %s -check-prefix=LABEL_6
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=LABEL_7 | %FileCheck %s -check-prefix=LABEL_7
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TOPLEVEL_1 | %FileCheck %s -check-prefix=TOPLEVEL_1
+
+func test(subject: Int) {
+  OUTER_IF_1:
+  if subject == 1 {
+    break #^LABEL_1^#
+// LABEL_1: Begin completions, 1 items
+// LABEL_1-DAG: Pattern/Local:                      OUTER_IF_1;
+// LABEL_1: End completions
+  }
+
+  OUTER_SWITCH_1:
+  switch subject {
+  case var x where x < 2:
+    break #^LABEL_2^#
+// LABEL_2: Begin completions, 1 items
+// LABEL_2-DAG: Pattern/Local:                      OUTER_SWITCH_1;
+// LABEL_2: End completions
+
+    INNER_IF_1: if subject == 1 {
+      INNER_FOR_1: for 0 ..< 1 {
+        break #^LABEL_3^#
+// LABEL_3: Begin completions, 3 items
+// LABEL_3-DAG: Pattern/Local:                      INNER_FOR_1;
+// LABEL_3-DAG: Pattern/Local:                      INNER_IF_1;
+// LABEL_3-DAG: Pattern/Local:                      OUTER_SWITCH_1;
+// LABEL_3: End completions
+      }
+      break #^LABEL_4^#
+// LABEL_4: Begin completions, 2 items
+// LABEL_4-DAG: Pattern/Local:                      INNER_IF_1;
+// LABEL_4-DAG: Pattern/Local:                      OUTER_SWITCH_1;
+// LABEL_4: End completions
+    }
+
+    INNER_IF_2: if subject == 1 {
+      INNER_WHILE_1: while i == 1 {
+        break #^LABEL_5^#
+// LABEL_5: Begin completions, 3 items
+// LABEL_5-DAG: Pattern/Local:                      INNER_WHILE_1;
+// LABEL_5-DAG: Pattern/Local:                      INNER_IF_2;
+// LABEL_5-DAG: Pattern/Local:                      OUTER_SWITCH_1;
+// LABEL_5: End completions
+      }
+    }
+
+    INNER_GUARD_1: guard subject == 1 else {
+      INNER_DOCATCH_1: do {
+      }
+      catch let err {
+        continue #^LABEL_6^#
+// LABEL_6: Begin completions, 1 items
+// LABEL_6-DAG: Pattern/Local:                      INNER_DOCATCH_1;
+// LABEL_6: End completions
+      }
+    }
+
+  }
+
+  OUTER_FOR_1: for subject == 1 {
+    break #^LABEL_7^#
+// LABEL_7: Begin completions, 1 items
+// LABEL_7-DAG: Pattern/Local:                      OUTER_FOR_1;
+// LABEL_7: End completions
+  }
+}
+
+TOP_IF_1: if true {}
+TOP_IF_2: if true { break #^TOPLEVEL_1^# }
+TOP_IF_3: if true {}
+// TOPLEVEL_1: Begin completions, 1 items
+// TOPLEVEL_1-DAG: Pattern/Local:                      TOP_IF_2;
+// TOPLEVEL_1: End completions