[ASTScope] Use ASTScope to find fallthrough source/dest cases.

Rather than depending on the tracking of state in switch cases to
remember the case statements that are the source and destination for a
`fallthrough` statement, compute them using ASTScope to find the
nearest enclosing case (the source) and then find the next case in the
same `switch` statement (the destination).

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -423,6 +423,9 @@ class ASTScopeImpl {
   computeIsCascadingUse(ArrayRef<const ASTScopeImpl *> history,
                         Optional<bool> initialIsCascadingUse);
 
+  static std::pair<CaseStmt *, CaseStmt *>
+  lookupFallthroughSourceAndDest(SourceFile *sourceFile, SourceLoc loc);
+
 #pragma mark - - lookup- starting point
 private:
   static const ASTScopeImpl *findStartingScopeForLookup(SourceFile *,

include/swift/AST/NameLookup.h

@@ -711,6 +711,18 @@ class ASTScope {
   static llvm::SmallVector<LabeledStmt *, 4>
   lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc);
 
+  /// Look for the directly enclosing case statement and the next case
+  /// statement, which together act as the source and destination for a
+  /// 'fallthrough' statement within a switch case.
+  ///
+  /// \returns a pair (fallthrough source, fallthrough dest). If the location
+  /// is not within the body of a case statement at all, the fallthrough
+  /// source will be \c nullptr. If there is a fallthrough source that case is
+  /// the last one, the fallthrough destination will be \c nullptr. A
+  /// well-formed 'fallthrough' statement has both a source and destination.
+  static std::pair<CaseStmt *, CaseStmt *>
+  lookupFallthroughSourceAndDest(SourceFile *sourceFile, SourceLoc loc);
+
   SWIFT_DEBUG_DUMP;
   void print(llvm::raw_ostream &) const;
   void dumpOneScopeMapLocation(std::pair<unsigned, unsigned>);

lib/AST/ASTScope.cpp

@@ -59,6 +59,11 @@ llvm::SmallVector<LabeledStmt *, 4> ASTScope::lookupLabeledStmts(
   return ASTScopeImpl::lookupLabeledStmts(sourceFile, loc);
 }
 
+std::pair<CaseStmt *, CaseStmt *> ASTScope::lookupFallthroughSourceAndDest(
+    SourceFile *sourceFile, SourceLoc loc) {
+  return ASTScopeImpl::lookupFallthroughSourceAndDest(sourceFile, loc);
+}
+
 #if SWIFT_COMPILER_IS_MSVC
 #pragma warning(push)
 #pragma warning(disable : 4996)

lib/AST/ASTScopeLookup.cpp

@@ -905,3 +905,55 @@ ASTScopeImpl::lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc) {
 
   return labeledStmts;
 }
+
+std::pair<CaseStmt *, CaseStmt *> ASTScopeImpl::lookupFallthroughSourceAndDest(
+    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.");
+
+  // Look for the enclosing case statement and its 'switch' statement.
+  CaseStmt *fallthroughSource = nullptr;
+  SwitchStmt *switchStmt = nullptr;
+  for (auto scope = innermost; scope && !scope->isLabeledStmtLookupTerminator();
+       scope = scope->getParent().getPtrOrNull()) {
+    // If we have a case statement, record it.
+    auto stmt = scope->getStmtIfAny();
+    if (!stmt) continue;
+
+    // If we've found the first case statement of a switch, record it as the
+    // fallthrough source. do-catch statements don't support fallthrough.
+    if (auto caseStmt = dyn_cast<CaseStmt>(stmt.get())) {
+      if (!fallthroughSource &&
+          caseStmt->getParentKind() == CaseParentKind::Switch)
+        fallthroughSource = caseStmt;
+
+      continue;
+    }
+
+    // If we've found the first switch statement, record it and we're done.
+    switchStmt = dyn_cast<SwitchStmt>(stmt.get());
+    if (switchStmt)
+      break;
+  }
+
+  // If we don't have both a fallthrough source and a switch statement
+  // enclosing it, the 'fallthrough' statement is ill-formed.
+  if (!fallthroughSource || !switchStmt)
+    return { nullptr, nullptr };
+
+  // Find this case in the list of cases for the switch. If we don't find it
+  // here, it means that the case isn't directly nested inside the switch, so
+  // the case and fallthrough are both ill-formed.
+  auto caseIter = llvm::find(switchStmt->getCases(), fallthroughSource);
+  if (caseIter == switchStmt->getCases().end())
+    return { nullptr, nullptr };
+
+  // Move along to the next case. This is the fallthrough destination.
+  ++caseIter;
+  auto fallthroughDest = caseIter == switchStmt->getCases().end() ? nullptr
+      : *caseIter;
+  return { fallthroughSource, fallthroughDest };
+}

lib/Sema/TypeCheckStmt.cpp

@@ -466,9 +466,6 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
   /// where this is greater than one.
   SmallVector<LabeledStmt*, 2> ActiveLabeledStmts;
 
-  /// The level of 'switch' nesting. 'fallthrough' is valid only in scopes where
-  /// this is greater than one.
-  unsigned SwitchLevel = 0;
   /// The destination block for a 'fallthrough' statement. Null if the switch
   /// scope depth is zero or if we are checking the final 'case' of the current
   /// switch.
@@ -536,11 +533,9 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
     CaseStmt *OuterFallthroughDest;
     AddSwitchNest(StmtChecker &SC) : SC(SC),
         OuterFallthroughDest(SC.FallthroughDest) {
-      ++SC.SwitchLevel;
     }
 
     ~AddSwitchNest() {
-      --SC.SwitchLevel;
       SC.FallthroughDest = OuterFallthroughDest;
     }
   };
@@ -896,33 +891,31 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
   }
 
   Stmt *visitFallthroughStmt(FallthroughStmt *S) {
-    auto sourceFile = DC->getParentSourceFile();
-    bool hasAnySwitches;
+    CaseStmt *fallthroughSource;
+    CaseStmt *fallthroughDest;
     if (getASTContext().LangOpts.EnableASTScopeLookup) {
-      auto activeLabeledStmts = ASTScope::lookupLabeledStmts(
-          sourceFile, S->getLoc());
-      auto numSwitches = llvm::count_if(activeLabeledStmts,
-                                        [](LabeledStmt *labeledStmt) {
-        return isa<SwitchStmt>(labeledStmt);
-      });
-      assert(numSwitches == SwitchLevel);
-      hasAnySwitches = numSwitches > 0;
+      auto sourceFile = DC->getParentSourceFile();
+      std::tie(fallthroughSource, fallthroughDest) =
+          ASTScope::lookupFallthroughSourceAndDest(sourceFile, S->getLoc());
+      assert(fallthroughSource == FallthroughSource);
+      assert(fallthroughDest == FallthroughDest);
     } else {
-      hasAnySwitches = SwitchLevel > 0;
+      fallthroughSource = FallthroughSource;
+      fallthroughDest = FallthroughDest;
     }
 
-    if (!hasAnySwitches) {
+    if (!fallthroughSource) {
       getASTContext().Diags.diagnose(S->getLoc(),
                                      diag::fallthrough_outside_switch);
       return nullptr;
     }
-    if (!FallthroughDest) {
+    if (!fallthroughDest) {
       getASTContext().Diags.diagnose(S->getLoc(),
                                      diag::fallthrough_from_last_case);
       return nullptr;
     }
-    S->setFallthroughSource(FallthroughSource);
-    S->setFallthroughDest(FallthroughDest);
+    S->setFallthroughSource(fallthroughSource);
+    S->setFallthroughDest(fallthroughDest);
     PreviousFallthrough = S;
     return S;
   }
@@ -1676,7 +1669,7 @@ Stmt *StmtChecker::visitBraceStmt(BraceStmt *BS) {
 void TypeChecker::typeCheckASTNode(ASTNode &node, DeclContext *DC,
                                    bool LeaveBodyUnchecked) {
   StmtChecker stmtChecker(DC);
-  // FIXME: 'ActiveLabeledStmts', 'SwitchLevel', etc. in StmtChecker are not
+  // FIXME: 'ActiveLabeledStmts', etc. in StmtChecker are not
   // populated. Since they don't affect "type checking", it's doesn't cause
   // any issue for now. But it should be populated nonetheless.
   stmtChecker.LeaveBraceStmtBodyUnchecked = LeaveBodyUnchecked;