[Sema] Fully allow fallthrough in switch expressions

Allow `fallthrough` to appear as the last statement
in the case of a `switch` expression. We already
allowed it in other positions, this was just an
oversight.

rdar://127670432

Author: hamishknight

include/swift/AST/DiagnosticsSema.def

@@ -1328,8 +1328,9 @@ ERROR(single_value_stmt_branch_empty,none,
       "expected expression in branch of '%0' expression",
       (StmtKind))
 ERROR(single_value_stmt_branch_must_end_in_result,none,
-      "non-expression branch of '%0' expression may only end with a 'throw'",
-      (StmtKind))
+      "non-expression branch of '%0' expression may only end with a 'throw'"
+      "%select{| or 'fallthrough'}1",
+      (StmtKind, bool))
 ERROR(cannot_jump_in_single_value_stmt,none,
       "cannot use '%0' to transfer control out of '%1' expression",
       (StmtKind, StmtKind))

lib/Sema/MiscDiagnostics.cpp

@@ -4478,7 +4478,7 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
           // default.
           Diags.diagnose(branch->getEndLoc(),
                          diag::single_value_stmt_branch_must_end_in_result,
-                         S->getKind());
+                         S->getKind(), isa<SwitchStmt>(S));
         }
         break;
       }

lib/Sema/TypeCheckStmt.cpp

@@ -3071,18 +3071,6 @@ class JumpOutOfContextFinder : public ASTWalker {
 };
 } // end anonymous namespace
 
-/// Whether the given brace statement ends with a 'throw'.
-static bool doesBraceEndWithThrow(BraceStmt *BS) {
-  if (BS->empty())
-    return false;
-
-  auto *S = BS->getLastElement().dyn_cast<Stmt *>();
-  if (!S)
-    return false;
-
-  return isa<ThrowStmt>(S);
-}
-
 IsSingleValueStmtResult
 areBranchesValidForSingleValueStmt(ASTContext &ctx, ArrayRef<Stmt *> branches) {
   TinyPtrVector<Stmt *> invalidJumps;
@@ -3108,8 +3096,20 @@ areBranchesValidForSingleValueStmt(ASTContext &ctx, ArrayRef<Stmt *> branches) {
       continue;
     }
 
-    // If there was no result, the branch must end in a 'throw'.
-    if (!doesBraceEndWithThrow(BS))
+    // If there was no result, the branch must end in a 'throw' or
+    // 'fallthrough'.
+    auto endsInJump = [&]() {
+      if (BS->empty())
+        return false;
+
+      auto *S = BS->getLastElement().dyn_cast<Stmt *>();
+      if (!S)
+        return false;
+
+      return isa<ThrowStmt>(S) || isa<FallthroughStmt>(S);
+    }();
+
+    if (!endsInJump)
       unterminatedBranches.push_back(BS);
   }
 

test/SILGen/switch_expr.swift

@@ -168,6 +168,35 @@ func testFallthrough() throws -> Int {
 // CHECK:       dealloc_stack [[RESULT]] : $*Int
 // CHECK:       return [[VAL]] : $Int
 
+func testFallthrough2() -> Int {
+  let x = switch Bool.random() {
+  case true:
+    fallthrough
+  case false:
+    1
+  }
+  return x
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s11switch_expr16testFallthrough2SiyF : $@convention(thin) () -> Int {
+// CHECK:       [[RESULT:%[0-9]+]] = alloc_stack $Int
+// CHECK:       switch_value {{%[0-9]+}} : $Builtin.Int1, case {{%[0-9]+}}: [[TRUEBB:bb[0-9]+]], case {{%[0-9]+}}: [[FALSEBB:bb[0-9]+]]
+//
+// CHECK:       [[TRUEBB]]:
+// CHECK-NEXT:  br [[ENDBB:bb[0-9]+]]
+//
+// CHECK:       [[FALSEBB]]:
+// CHECK-NEXT:  br [[ENDBB]]
+//
+// CHECK:       [[ENDBB]]:
+// CHECK:       [[ONELIT:%[0-9]+]] = integer_literal $Builtin.IntLiteral, 1
+// CHECK:       [[ONE:%[0-9]+]] = apply {{%[0-9]+}}([[ONELIT]], {{%[0-9]+}})
+// CHECK:       store [[ONE]] to [trivial] [[RESULT]] : $*Int
+// CHECK:       [[VAL:%[0-9]+]] = load [trivial] [[RESULT]] : $*Int
+// CHECK:       [[VAL_RESULT:%[0-9]+]] = move_value [var_decl] [[VAL]] : $Int
+// CHECK:       dealloc_stack [[RESULT]] : $*Int
+// CHECK:       return [[VAL_RESULT:[%0-9]+]] : $Int
+
 func testClosure() throws -> Int {
   let fn = {
     switch Bool.random() {

test/expr/unary/switch_expr.swift

@@ -820,7 +820,7 @@ func returnBranches6() -> Int {
   case true:
     print("hello")
     0 // expected-warning {{integer literal is unused}}
-    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     1
   }
@@ -835,7 +835,7 @@ func returnBranches6PoundIf() -> Int {
     print("hello")
     0 // expected-warning {{integer literal is unused}}
     #endif
-    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     1
   }
@@ -850,7 +850,7 @@ func returnBranches6PoundIf2() -> Int {
     print("hello")
     0
     #endif
-    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     1
   }
@@ -882,7 +882,7 @@ func returnBranches9() -> Int {
   let i = switch Bool.random() {
   case true:
     print("hello")
-    if .random() {} // expected-error {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+    if .random() {} // expected-error {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     1
   }
@@ -898,7 +898,7 @@ func returnBranches10() -> Int {
       0 // expected-warning {{integer literal is unused}}
     case false:
       2 // expected-warning {{integer literal is unused}}
-    } // expected-error {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+    } // expected-error {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     1
   }
@@ -914,7 +914,7 @@ func returnBranches11() -> Int {
       "" // expected-warning {{string literal is unused}}
     case false:
       2 // expected-warning {{integer literal is unused}}
-    } // expected-error {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+    } // expected-error {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     1
   }
@@ -1059,7 +1059,7 @@ func testPoundIfBranch3() -> Int {
     #if false
     0
     #endif
-  // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+  // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     0
   }
@@ -1100,7 +1100,7 @@ func testPoundIfBranch6() -> Int {
     0
     #endif
     0 // expected-warning {{integer literal is unused}}
-    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw'}}
+    // expected-error@-1 {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
   case false:
     1
   }
@@ -1180,6 +1180,38 @@ func fallthrough2() -> Int {
   return x
 }
 
+func fallthrough3() -> Int {
+  let x = switch true {
+  case true:
+    fallthrough
+  case false:
+    0
+  }
+  return x
+}
+
+func fallthrough4() -> Int {
+  let x = switch true {
+  case true:
+    fallthrough
+    return 0 // expected-error {{cannot use 'return' to transfer control out of 'switch' expression}}
+  case false:
+    0
+  }
+  return x
+}
+
+func fallthrough5() -> Int {
+  let x = switch true {
+  case true:
+    fallthrough
+    print(0) // expected-error {{non-expression branch of 'switch' expression may only end with a 'throw' or 'fallthrough'}}
+  case false:
+    0
+  }
+  return x
+}
+
 func breakAfterNeverExpr() -> String {
   // We avoid turning this into a switch expression because of the 'break'.
   switch Bool.random() {

test/stmt/then_stmt_exec.swift

@@ -8,9 +8,9 @@
 // Required for experimental features
 // REQUIRES: asserts
 
-func testDefer(_ x: Bool) -> Int {
+func testDeferIf(_ x: Bool) -> Int {
   defer {
-    print("defer fn")
+    print("defer testDeferIf")
   }
   let x = if x {
     defer {
@@ -24,9 +24,55 @@ func testDefer(_ x: Bool) -> Int {
   print("after if")
   return x
 }
-_ = testDefer(true)
+_ = testDeferIf(true)
 
 // CHECK:      enter if
 // CHECK-NEXT: defer if
 // CHECK-NEXT: after if
-// CHECK-NEXT: defer fn
+// CHECK-NEXT: defer testDeferIf
+
+func testDeferSwitch(_ x: Bool) -> Int {
+  defer {
+    print("defer testDeferSwitch")
+  }
+  let x = switch x {
+  case true:
+    defer {
+      print("defer case true")
+    }
+    print("enter case true")
+    fallthrough
+  case false:
+    defer {
+      print("defer case false")
+    }
+    print("enter case false")
+    then 1
+  }
+  print("after switch")
+  return x
+}
+_ = testDeferSwitch(true)
+// CHECK:      enter case true
+// CHECK-NEXT: defer case true
+// CHECK-NEXT: enter case false
+// CHECK-NEXT: defer case false
+// CHECK-NEXT: after switch
+// CHECK-NEXT: defer testDeferSwitch
+
+func testFallthrough(_ x: Bool) -> Int {
+  var z = 0
+  let y: Int
+  let x = switch x {
+  case true:
+    z = 1
+    fallthrough
+  case false:
+    y = 2
+    then 3
+  }
+  return x + y + z
+}
+print("fallthrough: \(testFallthrough(true))")
+
+// CHECK: fallthrough: 6