[Diag] [Typechecker] Diagnose throw in defer properly

lib/Sema/TypeCheckError.cpp

@@ -892,6 +892,9 @@ class Context {
 
     /// The pattern of a catch.
     CatchGuard,
+
+    /// A defer body
+    DeferBody
   };
 
 private:
@@ -917,6 +920,7 @@ class Context {
 
   Kind TheKind;
   bool DiagnoseErrorOnTry = false;
+  bool isInDefer = false;
   DeclContext *RethrowsDC = nullptr;
   InterpolatedStringLiteralExpr *InterpolatedString = nullptr;
 
@@ -959,6 +963,12 @@ class Context {
         D->getInterfaceType(), D->hasImplicitSelfDecl() ? 2 : 1));
   }
 
+  static Context forDeferBody() {
+    Context result(Kind::DeferBody);
+    result.isInDefer = true;
+    return result;
+  }
+
   static Context forInitializer(Initializer *init) {
     if (isa<DefaultArgumentInitializer>(init)) {
       return Context(Kind::DefaultArgument);
@@ -1029,13 +1039,20 @@ class Context {
   }
 
   static void diagnoseThrowInIllegalContext(TypeChecker &TC, ASTNode node,
-                                            StringRef description) {
+                                            StringRef description,
+                                            bool throwInDefer = false) {
     if (auto *e = node.dyn_cast<Expr*>())
       if (isa<ApplyExpr>(e)) {
         TC.diagnose(e->getLoc(), diag::throwing_call_in_illegal_context,
                     description);
         return;
       }
+
+    if (throwInDefer) {
+      // Return because this would've already been diagnosed in TypeCheckStmt.
+      return;
+    }
+
     TC.diagnose(node.getStartLoc(), diag::throw_in_illegal_context,
                 description);
   }
@@ -1196,6 +1213,9 @@ class Context {
     case Kind::CatchGuard:
       diagnoseThrowInIllegalContext(TC, E, "a catch guard expression");
       return;
+    case Kind::DeferBody:
+      diagnoseThrowInIllegalContext(TC, E, "a defer body", isInDefer);
+      return;
     }
     llvm_unreachable("bad context kind");
   }
@@ -1223,6 +1243,7 @@ class Context {
     case Kind::DefaultArgument:
     case Kind::CatchPattern:
     case Kind::CatchGuard:
+    case Kind::DeferBody:
       assert(!DiagnoseErrorOnTry);
       // Diagnosed at the call sites.
       return;
@@ -1670,7 +1691,10 @@ void TypeChecker::checkFunctionErrorHandling(AbstractFunctionDecl *fn) {
   PrettyStackTraceDecl debugStack("checking error handling for", fn);
 #endif
 
-  CheckErrorCoverage checker(*this, Context::forFunction(fn));
+  auto isDeferBody = isa<FuncDecl>(fn) && cast<FuncDecl>(fn)->isDeferBody();
+  auto context =
+      isDeferBody ? Context::forDeferBody() : Context::forFunction(fn);
+  CheckErrorCoverage checker(*this, context);
 
   // If this is a debugger function, suppress 'try' marking at the top level.
   if (fn->getAttrs().hasAttribute<LLDBDebuggerFunctionAttr>())

lib/Sema/TypeCheckStmt.cpp

@@ -545,6 +545,12 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
   }
 
   Stmt *visitThrowStmt(ThrowStmt *TS) {
+    // If the throw is in a defer, then it isn't valid.
+    if (isInDefer()) {
+      TC.diagnose(TS->getThrowLoc(), diag::jump_out_of_defer, "throw");
+      return nullptr;
+    }
+
     // Coerce the operand to the exception type.
     auto E = TS->getSubExpr();
 

test/stmt/statements.swift

@@ -356,13 +356,46 @@ func test_defer(_ a : Int) {
 
 class SomeTestClass {
   var x = 42
-  
+
   func method() {
     defer { x = 97 }  // self. not required here!
     // expected-warning@-1 {{'defer' statement before end of scope always executes immediately}}{{5-10=do}}
   }
 }
 
+enum DeferThrowError: Error {
+  case someError
+}
+
+func throwInDefer() {
+  defer { throw DeferThrowError.someError } // expected-error {{'throw' cannot transfer control out of a defer statement}}
+  print("Foo")
+}
+
+func throwingFuncInDefer1() throws {
+  defer { try throwingFunctionCalledInDefer() } // expected-error {{errors cannot be thrown out of a defer body}}
+  print("Bar")
+}
+
+func throwingFuncInDefer2() throws {
+  defer { throwingFunctionCalledInDefer() } // expected-error {{errors cannot be thrown out of a defer body}}
+  print("Bar")
+}
+
+func throwingFuncInDefer3() {
+  defer { try throwingFunctionCalledInDefer() } // expected-error {{errors cannot be thrown out of a defer body}}
+  print("Bar")
+}
+
+func throwingFuncInDefer4() {
+  defer { throwingFunctionCalledInDefer() } // expected-error {{errors cannot be thrown out of a defer body}}
+  print("Bar")
+}
+
+func throwingFunctionCalledInDefer() throws {
+  throw DeferThrowError.someError
+}
+
 class SomeDerivedClass: SomeTestClass {
   override init() {
     defer {