Merge pull request swiftlang#33278 from DougGregor/statement-checker-cleanups

Author: swift-ci

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,

lib/Sema/BuilderTransform.cpp

@@ -1348,8 +1348,7 @@ class BuilderClosureRewriter
       // Check restrictions on '@unknown'.
       if (caseStmt->hasUnknownAttr()) {
         checkUnknownAttrRestrictions(
-            cs.getASTContext(), caseStmt, /*fallthroughDest=*/nullptr,
-            limitExhaustivityChecks);
+            cs.getASTContext(), caseStmt, limitExhaustivityChecks);
       }
 
       ++caseIndex;

lib/Sema/TypeCheckStmt.cpp

@@ -352,6 +352,50 @@ emitUnresolvedLabelDiagnostics(DiagnosticEngine &DE,
   }
 }
 
+/// Find the target of a break or continue statement without a label.
+///
+/// \returns the target, if one was found, or \c nullptr if no such target
+/// exists.
+static LabeledStmt *findUnlabeledBreakOrContinueStmtTarget(
+    ASTContext &ctx, SourceFile *sourceFile, SourceLoc loc,
+    bool isContinue, DeclContext *dc,
+    ArrayRef<LabeledStmt *> activeLabeledStmts) {
+  for (auto labeledStmt : activeLabeledStmts) {
+    // 'break' with no label looks through non-loop structures
+    // except 'switch'.
+    // 'continue' ignores non-loop structures.
+    if (!labeledStmt->requiresLabelOnJump() &&
+        (!isContinue || labeledStmt->isPossibleContinueTarget())) {
+      return labeledStmt;
+    }
+  }
+
+  // If we're in a defer, produce a tailored diagnostic.
+  if (isDefer(dc)) {
+    ctx.Diags.diagnose(
+        loc, diag::jump_out_of_defer, isContinue ? "continue": "break");
+    return nullptr;
+  }
+
+  // If we're dealing with an unlabeled break inside of an 'if' or 'do'
+  // statement, produce a more specific error.
+  if (!isContinue &&
+      llvm::any_of(activeLabeledStmts,
+                   [&](Stmt *S) -> bool {
+                     return isa<IfStmt>(S) || isa<DoStmt>(S);
+                   })) {
+    ctx.Diags.diagnose(
+        loc, diag::unlabeled_break_outside_loop);
+    return nullptr;
+  }
+
+  // Otherwise produce a generic error.
+  ctx.Diags.diagnose(
+      loc,
+      isContinue ? diag::continue_outside_loop : diag::break_outside_loop);
+  return nullptr;
+}
+
 /// Find the target of a break or continue statement.
 ///
 /// \returns the target, if one was found, or \c nullptr if no such target
@@ -361,7 +405,6 @@ static LabeledStmt *findBreakOrContinueStmtTarget(
     SourceLoc loc, Identifier targetName, SourceLoc targetLoc,
     bool isContinue, DeclContext *dc,
     ArrayRef<LabeledStmt *> oldActiveLabeledStmts) {
-  TopCollection<unsigned, LabeledStmt *> labelCorrections(3);
 
   // Retrieve the active set of labeled statements.
   // FIXME: Once everything uses ASTScope lookup, \c oldActiveLabeledStmts
@@ -375,43 +418,37 @@ static LabeledStmt *findBreakOrContinueStmtTarget(
         oldActiveLabeledStmts.rbegin(), oldActiveLabeledStmts.rend());
   }
 
-  // Pick the nearest break target that matches the specified name.
+  // Handle an unlabeled break separately; that's the easy case.
   if (targetName.empty()) {
-    for (auto labeledStmt : activeLabeledStmts) {
-      // 'break' with no label looks through non-loop structures
-      // except 'switch'.
-      // 'continue' ignores non-loop structures.
-      if (!labeledStmt->requiresLabelOnJump() &&
-          (!isContinue || labeledStmt->isPossibleContinueTarget())) {
-        return labeledStmt;
-      }
-    }
-  } else {
-    // Scan inside out until we find something with the right label.
-    for (auto labeledStmt : activeLabeledStmts) {
-      if (targetName == labeledStmt->getLabelInfo().Name) {
-        // Continue cannot be used to repeat switches, use fallthrough instead.
-        if (isContinue && !labeledStmt->isPossibleContinueTarget()) {
-          ctx.Diags.diagnose(
-              loc, diag::continue_not_in_this_stmt,
-              isa<SwitchStmt>(labeledStmt) ? "switch" : "if");
-          return nullptr;
-        }
+    return findUnlabeledBreakOrContinueStmtTarget(
+        ctx, sourceFile, loc, isContinue, dc, activeLabeledStmts);
+  }
 
-        return labeledStmt;
+  // Scan inside out until we find something with the right label.
+  TopCollection<unsigned, LabeledStmt *> labelCorrections(3);
+  for (auto labeledStmt : activeLabeledStmts) {
+    if (targetName == labeledStmt->getLabelInfo().Name) {
+      // Continue cannot be used to repeat switches, use fallthrough instead.
+      if (isContinue && !labeledStmt->isPossibleContinueTarget()) {
+        ctx.Diags.diagnose(
+            loc, diag::continue_not_in_this_stmt,
+            isa<SwitchStmt>(labeledStmt) ? "switch" : "if");
+        return nullptr;
       }
 
-      unsigned distance =
-        TypeChecker::getCallEditDistance(
-            DeclNameRef(targetName),
-            labeledStmt->getLabelInfo().Name,
-            TypeChecker::UnreasonableCallEditDistance);
-      if (distance < TypeChecker::UnreasonableCallEditDistance)
-        labelCorrections.insert(distance, std::move(labeledStmt));
+      return labeledStmt;
     }
-    labelCorrections.filterMaxScoreRange(
-      TypeChecker::MaxCallEditDistanceFromBestCandidate);
+
+    unsigned distance =
+      TypeChecker::getCallEditDistance(
+          DeclNameRef(targetName),
+          labeledStmt->getLabelInfo().Name,
+          TypeChecker::UnreasonableCallEditDistance);
+    if (distance < TypeChecker::UnreasonableCallEditDistance)
+      labelCorrections.insert(distance, std::move(labeledStmt));
   }
+  labelCorrections.filterMaxScoreRange(
+    TypeChecker::MaxCallEditDistanceFromBestCandidate);
 
   // If we're in a defer, produce a tailored diagnostic.
   if (isDefer(dc)) {
@@ -420,26 +457,6 @@ static LabeledStmt *findBreakOrContinueStmtTarget(
     return nullptr;
   }
 
-  if (targetName.empty()) {
-    // If we're dealing with an unlabeled break inside of an 'if' or 'do'
-    // statement, produce a more specific error.
-    if (!isContinue &&
-        llvm::any_of(activeLabeledStmts,
-                     [&](Stmt *S) -> bool {
-                       return isa<IfStmt>(S) || isa<DoStmt>(S);
-                     })) {
-      ctx.Diags.diagnose(
-          loc, diag::unlabeled_break_outside_loop);
-      return nullptr;
-    }
-
-    // Otherwise produce a generic error.
-    ctx.Diags.diagnose(
-        loc,
-        isContinue ? diag::continue_outside_loop : diag::break_outside_loop);
-    return nullptr;
-  }
-
   // Provide potential corrections for an incorrect label.
   emitUnresolvedLabelDiagnostics(
       ctx.Diags, targetLoc, targetName, labelCorrections);
@@ -1316,8 +1333,7 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
         assert(parentKind == CaseParentKind::Switch &&
                "'@unknown' can only appear on switch cases");
         checkUnknownAttrRestrictions(
-            getASTContext(), caseBlock, FallthroughDest,
-            limitExhaustivityChecks);
+            getASTContext(), caseBlock, limitExhaustivityChecks);
       }
 
       Stmt *body = caseBlock->getBody();
@@ -2237,8 +2253,9 @@ void TypeChecker::typeCheckTopLevelCodeDecl(TopLevelCodeDecl *TLCD) {
 }
 
 void swift::checkUnknownAttrRestrictions(
-    ASTContext &ctx, CaseStmt *caseBlock, CaseStmt *fallthroughDest,
+    ASTContext &ctx, CaseStmt *caseBlock,
     bool &limitExhaustivityChecks) {
+  CaseStmt *fallthroughDest = caseBlock->findNextCaseStmt();
   if (caseBlock->getCaseLabelItems().size() != 1) {
     assert(!caseBlock->getCaseLabelItems().empty() &&
            "parser should not produce case blocks with no items");

lib/Sema/TypeChecker.h

@@ -1368,8 +1368,7 @@ bool areGenericRequirementsSatisfied(const DeclContext *DC,
 /// Check for restrictions on the use of the @unknown attribute on a
 /// case statement.
 void checkUnknownAttrRestrictions(
-    ASTContext &ctx, CaseStmt *caseBlock, CaseStmt *fallthroughDest,
-    bool &limitExhaustivityChecks);
+    ASTContext &ctx, CaseStmt *caseBlock, bool &limitExhaustivityChecks);
 
 /// Bind all of the pattern variables that occur within a case statement and
 /// all of its case items to their "parent" pattern variables, forming chains

test/Constraints/function_builder_diags.swift

@@ -433,16 +433,27 @@ struct TestConstraintGenerationErrors {
 
 // Check @unknown
 func testUnknownInSwitchSwitch(e: E) {
-    tuplify(true) { c in
+  tuplify(true) { c in
     "testSwitch"
     switch e {
-    @unknown case .a: // expected-error{{'@unknown' is only supported for catch-all cases ("case _")}}
-      "a"
     case .b(let i, let s?):
       i * 2
       s + "!"
     default:
       "nothing"
+    @unknown case .a: // expected-error{{'@unknown' is only supported for catch-all cases ("case _")}}
+      // expected-error@-1{{'case' blocks cannot appear after the 'default' block of a 'switch'}}
+      "a"
+    }
+  }
+
+  tuplify(true) { c in
+    "testSwitch"
+    switch e {
+      @unknown case _: // expected-error{{'@unknown' can only be applied to the last case in a switch}}
+      "a"
+    default:
+      "default"
     }
   }
 }