[Sema] Prevent fallthrough jumping out of if expressions

Impose the same limit we impose on other
forms of control flow statements (e.g `break`,
`continue`, `return`), where it cannot transfer
control out of the expression. This fixes a crash
where we'd fail to find a fallthrough nested in an
`if` expression. It is technically source breaking,
as we would have allowed the case where the `if`
expression is a directly nested result of an outer
`switch` expression, but I would be very surprised
if anyone is relying on that.

rdar://133845101

Author: hamishknight

lib/Sema/TypeCheckStmt.cpp

@@ -3022,11 +3022,14 @@ namespace {
 /// An ASTWalker that searches for any break/continue/return statements that
 /// jump out of the context the walker starts at.
 class JumpOutOfContextFinder : public ASTWalker {
+  const Stmt *ParentStmt;
   TinyPtrVector<Stmt *> &Jumps;
   SmallPtrSet<Stmt *, 4> ParentLabeledStmts;
+  SmallPtrSet<CaseStmt *, 4> ParentCaseStmts;
 
 public:
-  JumpOutOfContextFinder(TinyPtrVector<Stmt *> &jumps) : Jumps(jumps) {}
+  JumpOutOfContextFinder(const Stmt *parentStmt, TinyPtrVector<Stmt *> &jumps)
+      : ParentStmt(parentStmt), Jumps(jumps) {}
 
   MacroWalking getMacroWalkingBehavior() const override {
     return MacroWalking::Expansion;
@@ -3035,9 +3038,11 @@ class JumpOutOfContextFinder : public ASTWalker {
   PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
     if (auto *LS = dyn_cast<LabeledStmt>(S))
       ParentLabeledStmts.insert(LS);
+    if (auto *CS = dyn_cast<CaseStmt>(S))
+      ParentCaseStmts.insert(CS);
 
-    // Cannot 'break', 'continue', or 'return' out of the statement. A jump to
-    // a statement within a branch however is fine.
+    // Cannot 'break', 'continue', 'fallthrough' or 'return' out of the
+    // statement. A jump to a statement within a branch however is fine.
     if (auto *BS = dyn_cast<BreakStmt>(S)) {
       if (!ParentLabeledStmts.contains(BS->getTarget()))
         Jumps.push_back(BS);
@@ -3046,6 +3051,17 @@ class JumpOutOfContextFinder : public ASTWalker {
       if (!ParentLabeledStmts.contains(CS->getTarget()))
         Jumps.push_back(CS);
     }
+    if (auto *FS = dyn_cast<FallthroughStmt>(S)) {
+      // The source must either be in the parent statement, or must be a
+      // nested CaseStmt we've seen. If there's no source, we will have
+      // already diagnosed.
+      if (auto *source = FS->getFallthroughSource()) {
+        if (source->getParentStmt() != ParentStmt &&
+            !ParentCaseStmts.contains(source)) {
+          Jumps.push_back(FS);
+        }
+      }
+    }
     if (isa<ReturnStmt>(S) || isa<FailStmt>(S))
       Jumps.push_back(S);
 
@@ -3057,6 +3073,11 @@ class JumpOutOfContextFinder : public ASTWalker {
       assert(removed);
       (void)removed;
     }
+    if (auto *CS = dyn_cast<CaseStmt>(S)) {
+      auto removed = ParentCaseStmts.erase(CS);
+      assert(removed);
+      (void)removed;
+    }
     return Action::Continue(S);
   }
 
@@ -3072,10 +3093,11 @@ class JumpOutOfContextFinder : public ASTWalker {
 } // end anonymous namespace
 
 IsSingleValueStmtResult
-areBranchesValidForSingleValueStmt(ASTContext &ctx, ArrayRef<Stmt *> branches) {
+areBranchesValidForSingleValueStmt(ASTContext &ctx, const Stmt *parentStmt,
+                                   ArrayRef<Stmt *> branches) {
   TinyPtrVector<Stmt *> invalidJumps;
   TinyPtrVector<Stmt *> unterminatedBranches;
-  JumpOutOfContextFinder jumpFinder(invalidJumps);
+  JumpOutOfContextFinder jumpFinder(parentStmt, invalidJumps);
 
   // Must have a single expression brace, and non-single-expression branches
   // must end with a throw.
@@ -3145,17 +3167,19 @@ IsSingleValueStmtRequest::evaluate(Evaluator &eval, const Stmt *S,
       return IsSingleValueStmtResult::nonExhaustiveIf();
 
     SmallVector<Stmt *, 4> scratch;
-    return areBranchesValidForSingleValueStmt(ctx, IS->getBranches(scratch));
+    return areBranchesValidForSingleValueStmt(ctx, IS,
+                                              IS->getBranches(scratch));
   }
   if (auto *SS = dyn_cast<SwitchStmt>(S)) {
     SmallVector<Stmt *, 4> scratch;
-    return areBranchesValidForSingleValueStmt(ctx, SS->getBranches(scratch));
+    return areBranchesValidForSingleValueStmt(ctx, SS,
+                                              SS->getBranches(scratch));
   }
   if (auto *DS = dyn_cast<DoStmt>(S)) {
     if (!ctx.LangOpts.hasFeature(Feature::DoExpressions))
       return IsSingleValueStmtResult::unhandledStmt();
 
-    return areBranchesValidForSingleValueStmt(ctx, DS->getBody());
+    return areBranchesValidForSingleValueStmt(ctx, DS, DS->getBody());
   }
   if (auto *DCS = dyn_cast<DoCatchStmt>(S)) {
     if (!ctx.LangOpts.hasFeature(Feature::DoExpressions))
@@ -3165,7 +3189,8 @@ IsSingleValueStmtRequest::evaluate(Evaluator &eval, const Stmt *S,
       return IsSingleValueStmtResult::nonExhaustiveDoCatch();
 
     SmallVector<Stmt *, 4> scratch;
-    return areBranchesValidForSingleValueStmt(ctx, DCS->getBranches(scratch));
+    return areBranchesValidForSingleValueStmt(ctx, DCS,
+                                              DCS->getBranches(scratch));
   }
   return IsSingleValueStmtResult::unhandledStmt();
 }

test/SILGen/if_expr.swift

@@ -659,3 +659,36 @@ struct LazyProp {
     0
   }
 }
+
+func testNestedFallthrough1() throws -> Int {
+  let x = if .random() {
+    switch Bool.random() {
+    case true:
+      fallthrough
+    case false:
+      break
+    }
+    throw Err()
+  } else {
+    0
+  }
+  return x
+}
+
+func testNestedFallthrough2() throws -> Int {
+  let x = if .random() {
+    switch Bool.random() {
+    case true:
+      if .random() {
+        fallthrough
+      }
+      break
+    case false:
+      break
+    }
+    throw Err()
+  } else {
+    0
+  }
+  return x
+}

test/SILGen/switch_expr.swift

@@ -197,6 +197,23 @@ func testFallthrough2() -> Int {
 // CHECK:       dealloc_stack [[RESULT]] : $*Int
 // CHECK:       return [[VAL_RESULT:[%0-9]+]] : $Int
 
+func testFallthrough3() throws -> Int {
+  switch Bool.random() {
+  case true:
+    switch Bool.random() {
+    case true:
+      if .random() {
+        fallthrough
+      }
+      throw Err()
+    case false:
+      1
+    }
+  case false:
+    0
+  }
+}
+
 func testClosure() throws -> Int {
   let fn = {
     switch Bool.random() {

test/expr/unary/if_expr.swift

@@ -987,6 +987,63 @@ func return4() throws -> Int {
   return i
 }
 
+// https://github.com/swiftlang/swift/issues/75880
+func fallthrough1() throws {
+  switch Bool.random() {
+  case true:
+    let _ = if .random() {
+      if .random () {
+        fallthrough // expected-error {{cannot use 'fallthrough' to transfer control out of 'if' expression}}
+      }
+      throw Err()
+    } else {
+      0
+    }
+  case false:
+    break
+  }
+}
+
+func fallthrough2() throws -> Int {
+  let x = switch Bool.random() {
+  case true:
+    if .random() {
+      if .random () {
+        fallthrough // expected-error {{cannot use 'fallthrough' to transfer control out of 'if' expression}}
+      }
+      throw Err()
+    } else {
+      0
+    }
+  case false:
+    1
+  }
+  return x
+}
+
+func fallthrough3() -> Int {
+  let x = switch Bool.random() {
+  case true:
+    if .random() {
+      fallthrough // expected-error {{cannot use 'fallthrough' to transfer control out of 'if' expression}}
+    } else {
+      0
+    }
+  case false:
+    1
+  }
+  return x
+}
+
+func fallthrough4() -> Int {
+  let x = if .random() {
+    fallthrough // expected-error {{'fallthrough' is only allowed inside a switch}}
+  } else {
+    0
+  }
+  return x
+}
+
 // MARK: Effect specifiers
 
 func tryIf1() -> Int {

test/expr/unary/switch_expr.swift

@@ -1212,6 +1212,16 @@ func fallthrough5() -> Int {
   return x
 }
 
+func fallthrough6() -> Int {
+  let x = switch true {
+  case true:
+    0
+  case false:
+    fallthrough // expected-error {{'fallthrough' without a following 'case' or 'default' block}}
+  }
+  return x
+}
+
 func breakAfterNeverExpr() -> String {
   // We avoid turning this into a switch expression because of the 'break'.
   switch Bool.random() {