[SR-5869][Parse/Sema] Expand recovery support for trailing closure in statement conditions

lib/Parse/ParseExpr.cpp

@@ -983,22 +983,44 @@ static bool isValidTrailingClosure(bool isExprBasic, Parser &P){
   if (P.peekToken().isAtStartOfLine())
     return false;
 
-
   // Determine if the {} goes with the expression by eating it, and looking
-  // to see if it is immediately followed by '{', 'where', or comma.  If so,
-  // we consider it to be part of the proceeding expression.
+  // to see if it is immediately followed by '{', 'where', ',', or any possible
+  // postfix or binary expression on the same line. If so, we consider it to be
+  // part of the proceeding expression at this point. Sema will emit a diagnostic
+  // later. See: Sema/MiscDiagnostics.cpp#checkStmtConditionTrailingClosure
   Parser::BacktrackingScope backtrack(P);
   P.consumeToken(tok::l_brace);
   P.skipUntil(tok::r_brace);
-  SourceLoc endLoc;
-  if (!P.consumeIf(tok::r_brace, endLoc) ||
-      P.Tok.isNot(tok::l_brace, tok::kw_where, tok::comma)) {
+  if (!P.consumeIf(tok::r_brace))
     return false;
-  }
 
-  // Recoverable case. Just return true here and Sema will emit a diagnostic
-  // later. see: Sema/MiscDiagnostics.cpp#checkStmtConditionTrailingClosure
-  return true;
+  switch (P.Tok.getKind()) {
+  case tok::l_brace:
+  case tok::kw_where:
+  case tok::comma:
+    // e.g.:
+    //   for x in ary.map { $0 + 1 }
+    //     where ... {
+    return true;
+  case tok::period:
+  case tok::period_prefix:
+  case tok::l_paren:
+  case tok::l_square:
+  case tok::question_postfix:
+  case tok::exclaim_postfix:
+  case tok::oper_binary_spaced:
+  case tok::oper_binary_unspaced:
+  case tok::question_infix:
+  case tok::equal:
+  case tok::kw_is:
+  case tok::kw_as:
+  case tok::colon:
+    // e.g.:
+    //   if ary.map { $0 + 1 }[0] == 1 {}
+    return !P.Tok.isAtStartOfLine();
+  default:
+    return false;
+  }
 }
 
 

lib/Sema/MiscDiagnostics.cpp

@@ -2835,7 +2835,7 @@ void swift::fixItEncloseTrailingClosure(TypeChecker &TC,
 static void checkStmtConditionTrailingClosure(TypeChecker &TC, const Expr *E) {
   if (E == nullptr || isa<ErrorExpr>(E)) return;
 
-  // Shallow walker. just dig into implicit expression.
+  // Find call expression with trailing closure.
   class DiagnoseWalker : public ASTWalker {
     TypeChecker &TC;
 
@@ -2871,13 +2871,31 @@ static void checkStmtConditionTrailingClosure(TypeChecker &TC, const Expr *E) {
     DiagnoseWalker(TypeChecker &tc) : TC(tc) { }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      // Dig into implicit expression.
-      if (E->isImplicit()) return { true, E };
-      // Diagnose call expression.
-      if (auto CE = dyn_cast<CallExpr>(E))
+      if (E->isImplicit())
+        return { true, E };
+
+      switch (E->getKind()) {
+      case ExprKind::Closure:
+      case ExprKind::Paren:
+      case ExprKind::Tuple:
+      case ExprKind::TupleShuffle:
+      case ExprKind::ObjectLiteral:
+      case ExprKind::Array:
+      case ExprKind::Dictionary:
+        // Don't dig into expression which syntactically enclose other
+        // expressinos. e.g. [...], (...), and { ... }
+        return { false, E };
+      case ExprKind::Call: {
+        auto *CE = cast<CallExpr>(E);
+        // Diagnose call expression.
         diagnoseIt(CE);
-      // Don't dig any further.
-      return { false, E };
+        // Only dig into the callee.
+        CE->getFn()->walk(*this);
+        return { false, E };
+      }
+      default:
+        return { true, E };
+      }
     }
   };
 

test/expr/closure/trailing.swift

@@ -138,6 +138,7 @@ limitXY(someInt, toGamut: intArray) {}  // expected-error {{extra argument 'toGa
 // <rdar://problem/23036383> QoI: Invalid trailing closures in stmt-conditions produce lowsy diagnostics
 func retBool(x: () -> Int) -> Bool {}
 func maybeInt(_: () -> Int) -> Int? {}
+func retAny(x: () -> Int) -> Any {}
 class Foo23036383 {
   init() {}
   func map(_: (Int) -> Int) -> Int? {}
@@ -205,6 +206,24 @@ func r23036383(foo: Foo23036383?, obj: Foo23036383) {
   // expected-error@-2 {{trailing closure requires parentheses for disambiguation in this context}} {{37-38=(x: }} {{43-43=)}} 
 }
 
+// SR-5869 - trailing closure not parsed correctly in if-statement conditional.
+func SR5869() {
+  let ary: [Int] = []
+  for _ in ary.map { $0 } .filter { $0 % 2 == 0 } { }
+  // expected-error@-1 {{trailing closure requires parentheses for disambiguation in this context}} {{34-35=(}} {{50-50=)}}
+  // expected-error@-2 {{trailing closure requires parentheses for disambiguation in this context}} {{19-20=(}} {{26-26=)}}
+  if ary.map { $0 + 1 }[0] == 1 { } // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{13-14=(}} {{24-24=)}}
+  if ary.map { $0 + 1 } == [1,2,3] { } // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{13-14=(}} {{24-24=)}}
+  if retAny { 1 } is Int { } // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{12-13=(x: }} {{18-18=)}}
+  if let _ = retAny { 1 } as? Int { } // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{20-21=(x: }} {{26-26=)}}
+  if retAny { 1 } is Int { } // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{12-13=(x: }} {{18-18=)}}
+  if retBool { 1 } ? retBool { 2 } : retBool { 3 } {}
+  // expected-error@-1 {{trailing closure requires parentheses for disambiguation in this context}} {{13-14=(x: }} {{19-19=)}}
+  // expected-error@-2 {{trailing closure requires parentheses for disambiguation in this context}} {{29-30=(x: }} {{35-35=)}}
+  // expected-error@-3 {{trailing closure requires parentheses for disambiguation in this context}} {{45-46=(x: }} {{51-51=)}}
+  if let _ = (retAny { 1 } as? Int) { } // OK.
+}
+
 func overloadOnLabel(a: () -> Void) {}
 func overloadOnLabel(b: () -> Void) {}
 func overloadOnLabel(c: () -> Void) {}