[ASTScope] Implement lookup for visible labeled statements.

Use the scope map to implement lookup for the labeled statements that
are visible from a given source location, which is a lexical property
that is currently handled with stateful tracking in the type checker.

For now, merely assert that the results of this approach are identical
to the state tracked in the statement type checker.

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -195,7 +195,7 @@ class ASTScopeImpl {
 
   /// Get ride of descendants and remove them from scopedNodes so the scopes
   /// can be recreated. Needed because typechecking inserts a return statment
-  /// into intiailizers.
+  /// into initializers.
   void disownDescendants(ScopeCreator &);
 
 public: // for addReusedBodyScopes
@@ -415,6 +415,10 @@ class ASTScopeImpl {
   unqualifiedLookup(SourceFile *, DeclNameRef, SourceLoc,
                     const DeclContext *startingContext, DeclConsumer);
 
+  /// Entry point into ASTScopeImpl-land for labeled statement lookups.
+  static llvm::SmallVector<LabeledStmt *, 4>
+  lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc);
+
   static Optional<bool>
   computeIsCascadingUse(ArrayRef<const ASTScopeImpl *> history,
                         Optional<bool> initialIsCascadingUse);
@@ -538,6 +542,12 @@ class ASTScopeImpl {
 
   NullablePtr<const ASTScopeImpl>
   ancestorWithDeclSatisfying(function_ref<bool(const Decl *)> predicate) const;
+
+  /// Whether this scope terminates lookup of labeled statements in the
+  /// children below it, because one cannot perform a "break" or a "continue"
+  /// in a child that goes outside of this scope.
+  virtual bool isLabeledStmtLookupTerminator() const;
+
 }; // end of ASTScopeImpl
 
 #pragma mark - specific scope classes
@@ -1357,6 +1367,9 @@ class ConditionalClauseScope final : public ASTScopeImpl {
 private:
   ArrayRef<StmtConditionElement> getCond() const;
   const StmtConditionElement &getStmtConditionElement() const;
+
+protected:
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 /// If, while, & guard statements all start with a conditional clause, then some
@@ -1380,6 +1393,7 @@ class ConditionalClausePatternUseScope final : public ASTScopeImpl {
   bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
                              DeclConsumer) const override;
   void printSpecifics(llvm::raw_ostream &out) const override;
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 
@@ -1713,6 +1727,9 @@ class AbstractStmtScope : public ASTScopeImpl {
   virtual Stmt *getStmt() const = 0;
   NullablePtr<Stmt> getStmtIfAny() const override { return getStmt(); }
   NullablePtr<const void> getReferrent() const override;
+
+protected:
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 class LabeledConditionalStmtScope : public AbstractStmtScope {
@@ -1910,6 +1927,7 @@ class ForEachPatternScope final : public ASTScopeImpl {
 protected:
   bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
                              DeclConsumer) const override;
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 class CaseStmtScope final : public AbstractStmtScope {

include/swift/AST/NameLookup.h

@@ -698,6 +698,19 @@ class ASTScope {
   computeIsCascadingUse(ArrayRef<const ast_scope::ASTScopeImpl *> history,
                         Optional<bool> initialIsCascadingUse);
 
+  /// Entry point to record the visible statement labels from the given
+  /// point.
+  ///
+  /// This lookup only considers labels that are visible within the current
+  /// function, so it will not return any labels from lexical scopes that
+  /// are not reachable via labeled control flow.
+  ///
+  /// \returns the set of labeled statements visible from the given source
+  /// location, with the innermost labeled statement first and proceeding
+  /// to the outermost labeled statement.
+  static llvm::SmallVector<LabeledStmt *, 4>
+  lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc);
+
   SWIFT_DEBUG_DUMP;
   void print(llvm::raw_ostream &) const;
   void dumpOneScopeMapLocation(std::pair<unsigned, unsigned>);

lib/AST/ASTScope.cpp

@@ -54,6 +54,11 @@ Optional<bool> ASTScope::computeIsCascadingUse(
   return ASTScopeImpl::computeIsCascadingUse(history, initialIsCascadingUse);
 }
 
+llvm::SmallVector<LabeledStmt *, 4> ASTScope::lookupLabeledStmts(
+    SourceFile *sourceFile, SourceLoc loc) {
+  return ASTScopeImpl::lookupLabeledStmts(sourceFile, loc);
+}
+
 #if SWIFT_COMPILER_IS_MSVC
 #pragma warning(push)
 #pragma warning(disable : 4996)

lib/AST/ASTScopeLookup.cpp

@@ -852,3 +852,47 @@ bool isLocWithinAnInactiveClause(const SourceLoc loc, SourceFile *SF) {
   SF->walk(tester);
   return tester.wasFoundWithinInactiveClause;
 }
+
+#pragma mark isLabeledStmtLookupTerminator implementations
+bool ASTScopeImpl::isLabeledStmtLookupTerminator() const {
+  return true;
+}
+
+bool ConditionalClauseScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+bool ConditionalClausePatternUseScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+bool AbstractStmtScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+bool ForEachPatternScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+llvm::SmallVector<LabeledStmt *, 4>
+ASTScopeImpl::lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc) {
+  // Find the innermost scope from which to start our search.
+  auto *const fileScope = sourceFile->getScope().impl;
+  const auto *innermost = fileScope->findInnermostEnclosingScope(loc, nullptr);
+  ASTScopeAssert(innermost->getWasExpanded(),
+                 "If looking in a scope, it must have been expanded.");
+
+  llvm::SmallVector<LabeledStmt *, 4> labeledStmts;
+  for (auto scope = innermost; scope && !scope->isLabeledStmtLookupTerminator();
+       scope = scope->getParent().getPtrOrNull()) {
+    // If we have a labeled statement, record it.
+    if (auto stmt = scope->getStmtIfAny()) {
+      if (auto labeledStmt = dyn_cast<LabeledStmt>(stmt.get())) {
+        labeledStmts.push_back(labeledStmt);
+        continue;
+      }
+    }
+  }
+
+  return labeledStmts;
+}

lib/Sema/TypeCheckStmt.cpp

@@ -347,6 +347,36 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
       // In any case, remember that we're in this labeled statement so that
       // break and continue are aware of it.
       SC.ActiveLabeledStmts.push_back(LS);
+
+      // Verify that the ASTScope-based query for active labeled statements
+      // is equivalent to what we have here.
+      if (LS->getStartLoc().isValid()) {
+        if (auto sourceFile = SC.DC->getParentSourceFile()) {
+          // The labeled statements from ASTScope lookup have the
+          // innermost labeled statement first, so reverse it to
+          // match the data structure maintained here.
+          auto activeFromASTScope = ASTScope::lookupLabeledStmts(
+              sourceFile, LS->getStartLoc());
+          assert(activeFromASTScope.front() == LS);
+          std::reverse(activeFromASTScope.begin(), activeFromASTScope.end());
+          if (activeFromASTScope != SC.ActiveLabeledStmts) {
+            llvm::errs() << "Old: ";
+            llvm::interleave(SC.ActiveLabeledStmts, [&](LabeledStmt *LS) {
+              llvm::errs() << LS;
+            }, [&] {
+              llvm::errs() << ' ';
+            });
+            llvm::errs() << "\nNew: ";
+            llvm::interleave(activeFromASTScope, [&](LabeledStmt *LS) {
+              llvm::errs() << LS;
+            }, [&] {
+              llvm::errs() << ' ';
+            });
+            llvm::errs() << "\n";
+          }
+          assert(activeFromASTScope == SC.ActiveLabeledStmts);
+        }
+      }
     }
     ~AddLabeledStmt() {
       SC.ActiveLabeledStmts.pop_back();