[CodeComplete] Use ASTScope lookup in getStmtLabelCompletions

Use the same lookup logic that we use for regular compilation.

Author: hamishknight

lib/IDE/CompletionLookup.cpp

@@ -3500,54 +3500,23 @@ void CompletionLookup::lookupExternalModuleDecls(
 }
 
 void CompletionLookup::getStmtLabelCompletions(SourceLoc Loc, bool isContinue) {
-  class LabelFinder : public ASTWalker {
-    SourceManager &SM;
-    SourceLoc TargetLoc;
-    bool IsContinue;
-
-  public:
-    SmallVector<Identifier, 2> Result;
-
-    /// Walk only the arguments of a macro.
-    MacroWalking getMacroWalkingBehavior() const override {
-      return MacroWalking::Arguments;
-    }
-
-    LabelFinder(SourceManager &SM, SourceLoc TargetLoc, bool IsContinue)
-        : SM(SM), TargetLoc(TargetLoc), IsContinue(IsContinue) {}
-
-    PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
-      if (SM.isBeforeInBuffer(S->getEndLoc(), TargetLoc))
-        return Action::SkipNode(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);
-          }
-        }
-      }
+  auto *SF = CurrDeclContext->getParentSourceFile();
+  llvm::SmallPtrSet<Identifier, 4> labels;
+  for (auto *LS : ASTScope::lookupLabeledStmts(SF, Loc)) {
+    if (isContinue && !LS->isPossibleContinueTarget())
+      continue;
 
-      return Action::Continue(S);
-    }
+    auto labelInfo = LS->getLabelInfo();
+    if (!labelInfo)
+      continue;
 
-    PostWalkResult<Stmt *> walkToStmtPost(Stmt *S) override {
-      return Action::Stop();
-    }
+    auto label = labelInfo.Name;
+    if (!labels.insert(label).second)
+      continue;
 
-    PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-      if (SM.isBeforeInBuffer(E->getEndLoc(), TargetLoc))
-        return Action::SkipNode(E);
-      return Action::Continue(E);
-    }
-  } Finder(CurrDeclContext->getASTContext().SourceMgr, Loc, isContinue);
-  const_cast<DeclContext *>(CurrDeclContext)->walkContext(Finder);
-  for (auto name : Finder.Result) {
     CodeCompletionResultBuilder Builder = makeResultBuilder(
         CodeCompletionResultKind::Pattern, SemanticContextKind::Local);
-    Builder.addTextChunk(name.str());
+    Builder.addTextChunk(label.str());
   }
 }
 

test/IDE/complete_controltransfer.swift

@@ -56,6 +56,15 @@ func test(subject: Int) {
 // LABEL_7: Begin completions, 1 items
 // LABEL_7-DAG: Pattern/Local:                      OUTER_FOR_1;
   }
+
+  // This is illegal, but make sure we don't duplicate in completion.
+  DUPLICATE_LABEL: do {
+    DUPLICATE_LABEL: do {
+      break #^LABEL_8^#
+      // LABEL_8:     Begin completions, 1 items
+      // LABEL_8-DAG: Pattern/Local:                      DUPLICATE_LABEL;
+    }
+  }
 }
 
 TOP_IF_1: if true {}