Simplify ASTScopeImpl::lookupFallthroughSourceAndDest.

Author: DougGregor

include/swift/AST/Stmt.h

@@ -1073,6 +1073,10 @@ class CaseStmt final
     return *CaseBodyVariables;
   }
 
+  /// Find the next case statement within the same 'switch' or 'do-catch',
+  /// if there is one.
+  CaseStmt *findNextCaseStmt() const;
+
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Case; }
 
   size_t numTrailingObjects(OverloadToken<CaseLabelItem>) const {

lib/AST/ASTScopeLookup.cpp

@@ -914,9 +914,7 @@ std::pair<CaseStmt *, CaseStmt *> ASTScopeImpl::lookupFallthroughSourceAndDest(
   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;
+  // Look for the enclosing case statement of a 'switch'.
   for (auto scope = innermost; scope && !scope->isLabeledStmtLookupTerminator();
        scope = scope->getParent().getPtrOrNull()) {
     // If we have a case statement, record it.
@@ -926,34 +924,12 @@ std::pair<CaseStmt *, CaseStmt *> ASTScopeImpl::lookupFallthroughSourceAndDest(
     // 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;
+      if (caseStmt->getParentKind() == CaseParentKind::Switch)
+        return { caseStmt, caseStmt->findNextCaseStmt() };
 
       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 };
+  return { nullptr, nullptr };
 }

lib/AST/Stmt.cpp

@@ -458,6 +458,40 @@ CaseStmt *CaseStmt::create(ASTContext &ctx, CaseParentKind ParentKind,
                body, caseVarDecls, implicit, fallthroughStmt);
 }
 
+namespace {
+
+template<typename CaseIterator>
+CaseStmt *findNextCaseStmt(
+    CaseIterator first, CaseIterator last, const CaseStmt *caseStmt) {
+  for(auto caseIter = first; caseIter != last; ++caseIter) {
+    if (*caseIter == caseStmt) {
+      ++caseIter;
+      return caseIter == last ? nullptr : *caseIter;
+    }
+  }
+
+  return nullptr;
+}
+
+}
+
+CaseStmt *CaseStmt::findNextCaseStmt() const {
+  auto parent = getParentStmt();
+  if (!parent)
+    return nullptr;
+
+  if (auto switchParent = dyn_cast<SwitchStmt>(parent)) {
+    return ::findNextCaseStmt(
+        switchParent->getCases().begin(), switchParent->getCases().end(),
+        this);
+  }
+
+  auto doCatchParent = cast<DoCatchStmt>(parent);
+  return ::findNextCaseStmt(
+      doCatchParent->getCatches().begin(), doCatchParent->getCatches().end(),
+      this);
+}
+
 SwitchStmt *SwitchStmt::create(LabeledStmtInfo LabelInfo, SourceLoc SwitchLoc,
                                Expr *SubjectExpr,
                                SourceLoc LBraceLoc,