Sema: Fix monthly TypeNullifier bug

This time, the issue is that TypeNullifier skips bodies of
multi-statement closures. However, ExprRewriter will type
happily pass them on to typeCheckClosureBody(). This could
trigger assertions. Fix this by skipping type checking of
multi-statement closures when diagnosing.

There seems to be a minor QoI regression in some test cases
that already looked pretty dodgy and/or had FIXMEs. However
I think its worth fixing a crash.

Author: slavapestov

lib/Sema/CSApply.cpp

@@ -295,6 +295,7 @@ namespace {
     DeclContext *dc;
     const Solution &solution;
     bool SuppressDiagnostics;
+    bool SkipClosures;
 
   private:
     /// \brief Coerce the given tuple to another tuple type.
@@ -1468,9 +1469,10 @@ namespace {
 
   public:
     ExprRewriter(ConstraintSystem &cs, const Solution &solution,
-                 bool suppressDiagnostics)
+                 bool suppressDiagnostics, bool skipClosures)
       : cs(cs), dc(cs.DC), solution(solution), 
-        SuppressDiagnostics(suppressDiagnostics) { }
+        SuppressDiagnostics(suppressDiagnostics),
+        SkipClosures(skipClosures) { }
 
     ConstraintSystem &getConstraintSystem() const { return cs; }
 
@@ -5646,6 +5648,11 @@ namespace {
     ExprWalker(ExprRewriter &Rewriter) : Rewriter(Rewriter) { }
 
     ~ExprWalker() {
+      // If we're re-typechecking an expression for diagnostics, don't
+      // visit closures that have non-single expression bodies.
+      if (Rewriter.SkipClosures)
+        return;
+
       auto &cs = Rewriter.getConstraintSystem();
       auto &tc = cs.getTypeChecker();
       for (auto *closure : closuresToTypeCheck)
@@ -6126,7 +6133,8 @@ bool ConstraintSystem::applySolutionFix(Expr *expr,
 Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
                                       Type convertType,
                                       bool discardedExpr,
-                                      bool suppressDiagnostics) {
+                                      bool suppressDiagnostics,
+                                      bool skipClosures) {
   // If any fixes needed to be applied to arrive at this solution, resolve
   // them to specific expressions.
   if (!solution.Fixes.empty()) {
@@ -6141,7 +6149,7 @@ Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
     return nullptr;
   }
 
-  ExprRewriter rewriter(*this, solution, suppressDiagnostics);
+  ExprRewriter rewriter(*this, solution, suppressDiagnostics, skipClosures);
   ExprWalker walker(rewriter);
 
   // Apply the solution to the expression.
@@ -6171,7 +6179,8 @@ Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
 Expr *ConstraintSystem::applySolutionShallow(const Solution &solution,
                                              Expr *expr,
                                              bool suppressDiagnostics) {
-  ExprRewriter rewriter(*this, solution, suppressDiagnostics);
+  ExprRewriter rewriter(*this, solution, suppressDiagnostics,
+                        /*skipClosures=*/false);
   rewriter.walkToExprPre(expr);
   Expr *result = rewriter.walkToExprPost(expr);
   if (result)
@@ -6183,7 +6192,9 @@ Expr *Solution::coerceToType(Expr *expr, Type toType,
                              ConstraintLocator *locator,
                              bool ignoreTopLevelInjection) const {
   auto &cs = getConstraintSystem();
-  ExprRewriter rewriter(cs, *this, /*suppressDiagnostics=*/false);
+  ExprRewriter rewriter(cs, *this,
+                        /*suppressDiagnostics=*/false,
+                        /*skipClosures=*/false);
   Expr *result = rewriter.coerceToType(expr, toType, locator);
   if (!result)
     return nullptr;
@@ -6309,7 +6320,9 @@ Expr *TypeChecker::callWitness(Expr *base, DeclContext *dc,
   }
 
   Solution &solution = solutions.front();
-  ExprRewriter rewriter(cs, solution, /*suppressDiagnostics=*/false);
+  ExprRewriter rewriter(cs, solution,
+                        /*suppressDiagnostics=*/false,
+                        /*skipClosures=*/false);
 
   auto memberRef = rewriter.buildMemberRef(base, openedFullType,
                                            base->getStartLoc(),

lib/Sema/CSDiag.cpp

@@ -2897,6 +2897,10 @@ typeCheckChildIndependently(Expr *subExpr, Type convertType,
   // the context is missing).
   TypeCheckExprOptions TCEOptions = TypeCheckExprFlags::DisableStructuralChecks;
 
+  // Don't walk into non-single expression closure bodies, because
+  // ExprTypeSaver and TypeNullifier skip them too.
+  TCEOptions |= TypeCheckExprFlags::SkipMultiStmtClosures;
+
   // Claim that the result is discarded to preserve the lvalue type of
   // the expression.
   if (options.contains(TCC_AllowLValue))

lib/Sema/ConstraintSystem.h

@@ -2368,9 +2368,14 @@ class ConstraintSystem {
   ///
   /// \param convertType the contextual type to which the
   /// expression should be converted, if any.
+  /// \param discardedExpr if true, the result of the expression
+  /// is contextually ignored.
+  /// \param skipClosures if true, don't descend into bodies of
+  /// non-single expression closures.
   Expr *applySolution(Solution &solution, Expr *expr,
                       Type convertType, bool discardedExpr,
-                      bool suppressDiagnostics);
+                      bool suppressDiagnostics,
+                      bool skipClosures);
 
   /// \brief Apply a given solution to the expression to the top-level
   /// expression, producing a fully type-checked expression.

lib/Sema/TypeCheckConstraints.cpp

@@ -1188,8 +1188,9 @@ bool TypeChecker::typeCheckExpression(Expr *&expr, DeclContext *dc,
   // Apply the solution to the expression.
   auto &solution = viable[0];
   bool isDiscarded = options.contains(TypeCheckExprFlags::IsDiscarded);
+  bool skipClosures = options.contains(TypeCheckExprFlags::SkipMultiStmtClosures);
   auto result = cs.applySolution(solution, expr, convertType, isDiscarded,
-                                 suppressDiagnostics);
+                                 suppressDiagnostics, skipClosures);
   if (!result) {
     // Failure already diagnosed, above, as part of applying the solution.
     return true;

lib/Sema/TypeChecker.h

@@ -194,6 +194,10 @@ enum class TypeCheckExprFlags {
   /// statement. This should only be used for syntactic restrictions, and should
   /// not affect type checking itself.
   IsExprStmt = 0x20,
+  
+  /// If set, this expression is being re-type checked as part of diagnostics,
+  /// and so we should not visit bodies of non-single expression closures.
+  SkipMultiStmtClosures = 0x40,
 };
 
 typedef OptionSet<TypeCheckExprFlags> TypeCheckExprOptions;

test/Constraints/closures.swift

@@ -162,3 +162,14 @@ func r15998821() {
 var _: (Int,Int)-> Int = {$0+$1+$2}  // expected-error {{contextual closure type '(Int, Int) -> Int' expects 2 arguments, but 3 were used in closure body}}
 
 
+// Crash when re-typechecking bodies of non-single expression closures
+
+struct CC {}
+func callCC<U>(f: CC -> U) -> () {}
+
+func typeCheckMultiStmtClosureCrash() {
+  callCC { // expected-error {{cannot invoke 'callCC' with an argument list of type '((CC) -> _)'}}
+    _ = $0 // expected-note@-1 {{expected an argument list of type '(CC -> U)'}}
+    return 1
+  }
+}

test/Constraints/diagnostics.swift

@@ -117,7 +117,7 @@ func rdar20142523() {
   map(0..<10, { x in // expected-error{{cannot invoke 'map' with an argument list of type '(Range<Int>, (_) -> _)'}}
     // expected-note @-1 {{overloads for 'map' exist with these partially matching parameter lists: (C, (C.Generator.Element) -> T), (T?, @noescape (T) -> U)}}
     ()
-    return x  // expected-error {{type of expression is ambiguous without more context}}
+    return x
   })
 }
 
@@ -425,8 +425,7 @@ func f20371273() {
 // FIXME: Should complain about not having a return type annotation in the closure.
 [0].map { _ in let r =  (1,2).0;  return r }
 // expected-error @-1 {{cannot invoke 'map' with an argument list of type '(@noescape (Int) throws -> _)'}}
-// expected-error @-2 {{cannot convert return expression of type 'Int' to return type 'T'}}
-// expected-note @-3 {{expected an argument list of type '(@noescape Int throws -> T)'}}
+// expected-note @-2 {{expected an argument list of type '(@noescape Int throws -> T)'}}
 
 // <rdar://problem/21078316> Less than useful error message when using map on optional dictionary type
 func rdar21078316() {

test/Constraints/patterns.swift

@@ -212,7 +212,7 @@ func bad(a: A<EE>) {
   a.map { // expected-error {{cannot invoke 'map' with an argument list of type '((EE) -> _)'}}
   // expected-note@-1 {{expected an argument list of type '(EE -> T)'}}
     let _: EE = $0
-    return 1 // expected-error {{cannot convert return expression of type 'Int' to return type 'T'}}
+    return 1
   }
 }
 
@@ -221,9 +221,9 @@ func ugly(a: A<EE>) {
   // expected-note@-1 {{expected an argument list of type '(EE -> T)'}}
     switch $0 {
     case .A:
-      return 1 // expected-error {{cannot convert return expression of type 'Int' to return type 'T'}}
+      return 1
     default:
-      return 2 // expected-error {{cannot convert return expression of type 'Int' to return type 'T'}}
+      return 2
     }
   }
 }

test/Constraints/subscript.swift

@@ -74,7 +74,7 @@ let _ = 1["1"]  // expected-error {{ambiguous use of 'subscript'}}
 // rdar://17687826 - QoI: error message when reducing to an untyped dictionary isn't helpful
 let squares = [ 1, 2, 3 ].reduce([:]) { (dict, n) in // expected-error {{cannot invoke 'reduce' with an argument list of type '([_ : _], @noescape (_, Int) throws -> _)'}}
   // expected-note @-1 {{expected an argument list of type '(T, combine: @noescape (T, Int) throws -> T)'}}
-  var dict = dict // expected-error {{type of expression is ambiguous without more context}}
+  var dict = dict
 
   dict[n] = n * n
   return dict

test/stmt/statements.swift

@@ -403,7 +403,7 @@ func test_is_as_patterns() {
 func matching_pattern_recursion() {
   switch 42 {
   case {  // expected-error {{expression pattern of type '() -> ()' cannot match values of type 'Int'}}
-      for i in zs {  // expected-error {{use of unresolved identifier 'zs'}}
+      for i in zs {
       }
   }: break
   }