[Parse/Sema][SR-1672] Improve diagnostics for trailing closures in st…

…mt-condition (swiftlang#3184)

Fix-it suggests normal argument expression, instead of of enclosing whole
expression with parens.

* Moved diagnostic logic to Sema, because we have to add correct argument
  label for the closure.

    if arr.starts(with: IDs) { $0.id == $2 } { ... }
                           ~~^
                           , isEquivalent:  )

* We now accept trailing closures for each expressions and right before `where`
  clause, as well as closures right before the body.

    if let _ = maybeInt { 1 }, someOtherCondition { ... }
                       ~^
                       (     )

    for let x in arr.map { $0 * 4 } where x != 0 { ... }
                        ~^
                        (          )

include/swift/AST/DiagnosticsParse.def

@@ -998,10 +998,6 @@ ERROR(unexpected_tokens_before_closure_in,none,
 ERROR(expected_closure_rbrace,none,
       "expected '}' at end of closure", ())
 
-ERROR(trailing_closure_requires_parens,none,
-      "trailing closure requires parentheses for disambiguation in this"
-      " context", ())
-
 WARNING(trailing_closure_excess_newlines,none,
         "trailing closure is separated from call site by multiple newlines", ())
 NOTE(trailing_closure_call_here,none,

include/swift/AST/DiagnosticsSema.def

@@ -2342,6 +2342,10 @@ NOTE(add_where_newline, none,
 NOTE(duplicate_where, none,
      "duplicate the 'where' on both patterns to check both patterns",())
 
+ERROR(trailing_closure_requires_parens,none,
+      "trailing closure requires parentheses for disambiguation in this"
+      " context", ())
+
 //------------------------------------------------------------------------------
 // Type Check Patterns
 //------------------------------------------------------------------------------

lib/Parse/ParseExpr.cpp

@@ -919,10 +919,10 @@ static bool isStartOfGetSetAccessor(Parser &P) {
          P.Tok.isContextualKeyword("willSet");
 }
 
-/// Disambiguate and diagnose invalid uses of trailing closures in a situation
+/// Recover invalid uses of trailing closures in a situation
 /// where the parser requires an expr-basic (which does not allow them).  We
-/// handle this by doing some lookahead in common situations and emitting a
-/// diagnostic with a fixit to add wrapping parens.
+/// handle this by doing some lookahead in common situations. And later, Sema
+/// will emit a diagnostic with a fixit to add wrapping parens.
 static bool isValidTrailingClosure(bool isExprBasic, Expr *baseExpr, Parser &P){
   assert(P.Tok.is(tok::l_brace) && "Couldn't be a trailing closure");
 
@@ -955,21 +955,19 @@ static bool isValidTrailingClosure(bool isExprBasic, Expr *baseExpr, Parser &P){
 
 
   // Determine if the {} goes with the expression by eating it, and looking
-  // to see if it is immediately followed by another {.  If so, we consider it
-  // to be part of the proceeding expression.
+  // to see if it is immediately followed by '{', 'where', or comma.  If so,
+  // we consider it to be part of the proceeding expression.
   Parser::BacktrackingScope backtrack(P);
-  auto startLoc = P.consumeToken(tok::l_brace);
+  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))
+      P.Tok.isNot(tok::l_brace, tok::kw_where, tok::comma)) {
     return false;
-
-  // Diagnose the bad case and return true so that the caller parses this as a
-  // trailing closure.
-  P.diagnose(startLoc, diag::trailing_closure_requires_parens)
-    .fixItInsert(baseExpr->getStartLoc(), "(")
-    .fixItInsertAfter(endLoc, ")");
+  }
+
+  // Recoverable case. Just return true here and Sema will emit a diagnostic
+  // later. see: Sema/MiscDiagnostics.cpp#checkStmtConditionTrailingClosure
   return true;
 }
 

lib/Sema/MiscDiagnostics.cpp

@@ -2784,6 +2784,120 @@ static void checkSwitch(TypeChecker &TC, const SwitchStmt *stmt) {
   }
 }
 
+// Perform checkStmtConditionTrailingClosure for single expression.
+static void checkStmtConditionTrailingClosure(TypeChecker &TC, const Expr *E) {
+  if (E == nullptr || isa<ErrorExpr>(E)) return;
+
+  // Shallow walker. just dig into implicit expression.
+  class DiagnoseWalker : public ASTWalker {
+    TypeChecker &TC;
+
+    void diagnoseIt(const CallExpr *E) {
+      auto argsExpr = E->getArg();
+      auto argsTy = argsExpr->getType();
+
+      // Ignore invalid argument type. Some diagnostics are already emitted.
+      if (!argsTy || argsTy->is<ErrorType>()) return;
+
+      if (auto TSE = dyn_cast<TupleShuffleExpr>(argsExpr))
+        argsExpr = TSE->getSubExpr();
+
+      SmallString<16> replacement;
+      SourceLoc lastLoc;
+      SourceRange closureRange;
+      if (auto PE = dyn_cast<ParenExpr>(argsExpr)) {
+        // Ignore non-trailing-closure.
+        if (!PE->hasTrailingClosure()) return;
+
+        closureRange = PE->getSubExpr()->getSourceRange();
+        lastLoc = PE->getLParenLoc();
+        if (lastLoc.isValid()) {
+          // Empty paren: e.g. if funcName() { 1 } { ... }
+          replacement = "";
+        } else {
+          // Bare trailing closure: e.g. if funcName { 1 } { ... }
+          replacement = "(";
+          lastLoc = E->getFn()->getEndLoc();
+        }
+      } else if (auto TE = dyn_cast<TupleExpr>(argsExpr)) {
+        // Ignore non-trailing-closure.
+        if (!TE->hasTrailingClosure()) return;
+
+        // Tuple + trailing closure: e.g. if funcName(x: 1) { 1 } { ... }
+        auto numElements = TE->getNumElements();
+        assert(numElements >= 2 && "Unexpected num of elements in TupleExpr");
+        closureRange = TE->getElement(numElements - 1)->getSourceRange();
+        lastLoc = TE->getElement(numElements - 2)->getEndLoc();
+        replacement = ", ";
+      } else {
+        // Can't be here.
+        return;
+      }
+
+      // Add argument label of the closure that is going to be enclosed in
+      // parens.
+      if (auto TT = argsTy->getAs<TupleType>()) {
+        assert(TT->getNumElements() != 0 && "Unexpected empty TupleType");
+        auto closureLabel = TT->getElement(TT->getNumElements() - 1).getName();
+        if(!closureLabel.empty()) {
+          replacement += closureLabel.str();
+          replacement += ": ";
+        }
+      }
+
+      // Emit diagnostics.
+      lastLoc = Lexer::getLocForEndOfToken(TC.Context.SourceMgr, lastLoc);
+      TC.diagnose(closureRange.Start, diag::trailing_closure_requires_parens)
+        .fixItReplaceChars(lastLoc, closureRange.Start, replacement)
+        .fixItInsertAfter(closureRange.End, ")");
+    }
+
+  public:
+    DiagnoseWalker(TypeChecker &tc) : TC(tc) { }
+
+    virtual std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+      // Dig into implict expression.
+      if (E->isImplicit()) return { true, E };
+      // Diagnose call expression.
+      if (auto CE = dyn_cast<CallExpr>(E))
+        diagnoseIt(CE);
+      // Don't dig any further.
+      return { false, E };
+    }
+  };
+
+  DiagnoseWalker Walker(TC);
+  const_cast<Expr *>(E)->walk(Walker);
+}
+
+/// \brief Diagnose trailing closure in statement-conditions.
+///
+/// Conditional statements, including 'for' or `switch` doesn't allow ambiguous
+/// trailing closures in these conditions part. Even if the parser can recover
+/// them, we force them to disambiguate.
+//
+/// E.g.:
+///   if let _ = arr?.map {$0+1} { ... }
+///   for _ in numbers.filter {$0 > 4} { ... }
+static void checkStmtConditionTrailingClosure(TypeChecker &TC, const Stmt *S) {
+  if (auto LCS = dyn_cast<LabeledConditionalStmt>(S)) {
+    for (auto elt : LCS->getCond()) {
+      if (elt.getKind() == StmtConditionElement::CK_PatternBinding)
+        checkStmtConditionTrailingClosure(TC, elt.getInitializer());
+      else if (elt.getKind() == StmtConditionElement::CK_Boolean)
+        checkStmtConditionTrailingClosure(TC, elt.getBoolean());
+      // No trailing closure for CK_Availability: e.g. `if #available() {}`.
+    }
+  } else if (auto SS = dyn_cast<SwitchStmt>(S)) {
+    checkStmtConditionTrailingClosure(TC, SS->getSubjectExpr());
+  } else if (auto FES = dyn_cast<ForEachStmt>(S)) {
+    checkStmtConditionTrailingClosure(TC, FES->getSequence());
+    checkStmtConditionTrailingClosure(TC, FES->getWhere());
+  } else if (auto DCS = dyn_cast<DoCatchStmt>(S)) {
+    for (auto CS : DCS->getCatches())
+      checkStmtConditionTrailingClosure(TC, CS->getGuardExpr());
+  }
+}
 
 static Optional<ObjCSelector>
 parseObjCSelector(ASTContext &ctx, StringRef string) {
@@ -3222,6 +3336,8 @@ void swift::performStmtDiagnostics(TypeChecker &TC, const Stmt *S) {
 
   if (auto switchStmt = dyn_cast<SwitchStmt>(S))
     checkSwitch(TC, switchStmt);
+
+  checkStmtConditionTrailingClosure(TC, S);
 }
 
 //===----------------------------------------------------------------------===//

test/Parse/recovery.swift

@@ -652,19 +652,7 @@ func postfixDot(a : String) {
     a.         // expected-error {{expected member name following '.'}}
 }
 
-// <rdar://problem/23036383> QoI: Invalid trailing closures in stmt-conditions produce lowsy diagnostics
-func r23036383(arr : [Int]?) {
-  if let _ = arr?.map {$0+1} {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{14-14=(}} {{29-29=)}}
-  }
-
-  let numbers = [1, 2]
-  for _ in numbers.filter {$0 > 4} {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{12-12=(}} {{35-35=)}}
-  }
-}
-
 // <rdar://problem/22290244> QoI: "UIColor." gives two issues, should only give one
 func f() {
   _ = ClassWithStaticDecls.  // expected-error {{expected member name following '.'}}
 }
-
-

test/expr/closure/trailing.swift

@@ -110,3 +110,73 @@ let someInt = 0
 let intArray = [someInt]
 limitXY(someInt, toGamut: intArray) {}  // expected-error {{extra argument 'toGamut' in call}}
 
+
+// <rdar://problem/23036383> QoI: Invalid trailing closures in stmt-conditions produce lowsy diagnostics
+func retBool(x: () -> Int) -> Bool {}
+func maybeInt(_: () -> Int) -> Int? {}
+class Foo23036383 {
+  init() {}
+  func map(_: (Int) -> Int) -> Int? {}
+  func meth1(x: Int, _: () -> Int) -> Bool {}
+  func meth2(_: Int, y: () -> Int) -> Bool {}
+  func filter(by: (Int) -> Bool) -> [Int] {}
+}
+enum MyErr : ErrorProtocol {
+  case A
+}
+
+func r23036383(foo: Foo23036383?, obj: Foo23036383) {
+
+  if retBool(x: { 1 }) { } // OK
+  if (retBool { 1 }) { } // OK
+
+  if retBool{ 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{13-13=(x: }} {{18-18=)}}
+  }
+  if retBool { 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{13-14=(x: }} {{19-19=)}}
+  }
+  if retBool() { 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{14-16=x: }} {{21-21=)}}
+  } else if retBool( ) { 0 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{21-24=x: }} {{29-29=)}}
+  }
+
+  if let _ = maybeInt { 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{22-23=(}} {{28-28=)}}
+  }
+  if let _ = maybeInt { 1 } , true {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{22-23=(}} {{28-28=)}}
+  }
+
+  if let _ = foo?.map {$0+1} {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{22-23=(}} {{29-29=)}}
+  }
+  if let _ = foo?.map() {$0+1} {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{23-25=}} {{31-31=)}}
+  }
+  if let _ = foo, retBool { 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{26-27=(x: }} {{32-32=)}}
+  }
+
+  if obj.meth1(x: 1) { 0 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{20-22=, }} {{27-27=)}}
+  }
+  if obj.meth2(1) { 0 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{17-19=, y: }} {{24-24=)}}
+  }
+
+  for _ in obj.filter {$0 > 4} {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{22-23=(by: }} {{31-31=)}}
+  }
+  for _ in obj.filter {$0 > 4} where true {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{22-23=(by: }} {{31-31=)}}
+  }
+  for _ in [1,2] where retBool { 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{31-32=(x: }} {{37-37=)}}
+  }
+
+  while retBool { 1 } { // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{16-17=(x: }} {{22-22=)}}
+  }
+  while let _ = foo, retBool { 1 } { // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{29-30=(x: }} {{35-35=)}}
+  }
+
+  switch retBool { return 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{17-18=(x: }} {{30-30=)}}
+    default: break
+  }
+
+  do {
+    throw MyErr.A;
+  } catch MyErr.A where retBool { 1 } {  // expected-error {{trailing closure requires parentheses for disambiguation in this context}} {{32-33=(x: }} {{38-38=)}}
+  } catch { }
+
+  if let _ = maybeInt { 1 }, retBool { 1 } { }
+  // expected-error@-1 {{trailing closure requires parentheses for disambiguation in this context}} {{22-23=(}} {{28-28=)}} 
+  // expected-error@-2 {{trailing closure requires parentheses for disambiguation in this context}} {{37-38=(x: }} {{43-43=)}} 
+}