Merge pull request #39989 from xedin/more-diag-improvements-for-multi-stmt-closures

[TypeChecker] SE-0326: Enable multi-statement closure inference by default

Author: xedin

include/swift/Basic/LangOptions.h

@@ -682,7 +682,7 @@ namespace swift {
 
     /// Enable experimental support for type inference through multi-statement
     /// closures.
-    bool EnableMultiStatementClosureInference = false;
+    bool EnableMultiStatementClosureInference = true;
 
     /// See \ref FrontendOptions.PrintFullConvention
     bool PrintFullConvention = false;

include/swift/Sema/CSBindings.h

@@ -338,6 +338,14 @@ class BindingSet {
 
   TypeVariableType *getTypeVariable() const { return Info.TypeVar; }
 
+  /// Check whether this binding set belongs to a type variable
+  /// that represents a result type of a closure.
+  bool forClosureResult() const;
+
+  /// Check whether this binding set belongs to a type variable
+  /// that represents a generic parameter.
+  bool forGenericParameter() const;
+
   bool canBeNil() const;
 
   /// If this type variable doesn't have any viable bindings, or

include/swift/Sema/CSFix.h

@@ -2282,6 +2282,10 @@ class SpecifyClosureParameterType final : public ConstraintFix {
 
   bool diagnose(const Solution &solution, bool asNote = false) const override;
 
+  bool diagnoseForAmbiguity(CommonFixesArray commonFixes) const override {
+    return diagnose(*commonFixes.front().first);
+  }
+
   static SpecifyClosureParameterType *create(ConstraintSystem &cs,
                                              ConstraintLocator *locator);
 

include/swift/Sema/Constraint.h

@@ -436,6 +436,8 @@ class Constraint final : public llvm::ilist_node<Constraint>,
       ASTNode Element;
       /// Contextual information associated with the element (if any).
       ContextualTypeInfo Context;
+      /// Identifies whether result of this node is unused.
+      bool IsDiscarded;
     } ClosureElement;
   };
 
@@ -489,7 +491,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a closure body element constraint.
-  Constraint(ASTNode node, ContextualTypeInfo context,
+  Constraint(ASTNode node, ContextualTypeInfo context, bool isDiscarded,
              ConstraintLocator *locator,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
@@ -579,12 +581,14 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
   static Constraint *createClosureBodyElement(ConstraintSystem &cs,
                                               ASTNode node,
-                                              ConstraintLocator *locator);
+                                              ConstraintLocator *locator,
+                                              bool isDiscarded = false);
 
   static Constraint *createClosureBodyElement(ConstraintSystem &cs,
                                               ASTNode node,
                                               ContextualTypeInfo context,
-                                              ConstraintLocator *locator);
+                                              ConstraintLocator *locator,
+                                              bool isDiscarded = false);
 
   /// Determine the kind of constraint.
   ConstraintKind getKind() const { return Kind; }
@@ -850,6 +854,11 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return ClosureElement.Context;
   }
 
+  bool isDiscardedElement() const {
+    assert(Kind == ConstraintKind::ClosureBodyElement);
+    return ClosureElement.IsDiscarded;
+  }
+
   /// For an applicable function constraint, retrieve the trailing closure
   /// matching rule.
   Optional<TrailingClosureMatching> getTrailingClosureMatching() const;

include/swift/Sema/ConstraintSystem.h

@@ -4845,8 +4845,8 @@ class ConstraintSystem {
   /// Simplify a closure body element constraint by generating required
   /// constraints to represent the given element in constraint system.
   SolutionKind simplifyClosureBodyElementConstraint(
-      ASTNode element, ContextualTypeInfo context, TypeMatchOptions flags,
-      ConstraintLocatorBuilder locator);
+      ASTNode element, ContextualTypeInfo context, bool isDiscarded,
+      TypeMatchOptions flags, ConstraintLocatorBuilder locator);
 
 public: // FIXME: Public for use by static functions.
   /// Simplify a conversion constraint with a fix applied to it.
@@ -5005,7 +5005,8 @@ class ConstraintSystem {
   /// \param replaceInvalidRefsWithErrors Indicates whether it's allowed
   /// to replace any discovered invalid member references with `ErrorExpr`.
   static bool preCheckExpression(Expr *&expr, DeclContext *dc,
-                                 bool replaceInvalidRefsWithErrors);
+                                 bool replaceInvalidRefsWithErrors,
+                                 bool leaveClosureBodiesUnchecked);
 
   /// Solve the system of constraints generated from provided target.
   ///
@@ -5236,6 +5237,16 @@ class ConstraintSystem {
   /// imported from C/ObjectiveC.
   bool isArgumentOfImportedDecl(ConstraintLocatorBuilder locator);
 
+  /// Check whether given closure should participate in inference e.g.
+  /// if it's a single-expression closure - it always does, but
+  /// multi-statement closures require special flags.
+  bool participatesInInference(ClosureExpr *closure) const;
+
+  /// Visit each subexpression that will be part of the constraint system
+  /// of the given expression, including those in closure bodies that will be
+  /// part of the constraint system.
+  void forEachExpr(Expr *expr, llvm::function_ref<Expr *(Expr *)> callback);
+
   SWIFT_DEBUG_DUMP;
   SWIFT_DEBUG_DUMPER(dump(Expr *));
 
@@ -6059,18 +6070,6 @@ BraceStmt *applyResultBuilderTransform(
           constraints::SolutionApplicationTarget)>
             rewriteTarget);
 
-/// Determine whether the given closure expression should be type-checked
-/// within the context of its enclosing expression. Otherwise, it will be
-/// separately type-checked once its enclosing expression has determined all
-/// of the parameter and result types without looking at the body.
-bool shouldTypeCheckInEnclosingExpression(ClosureExpr *expr);
-
-/// Visit each subexpression that will be part of the constraint system
-/// of the given expression, including those in closure bodies that will be
-/// part of the constraint system.
-void forEachExprInConstraintSystem(
-    Expr *expr, llvm::function_ref<Expr *(Expr *)> callback);
-
 /// Whether the given parameter requires an argument.
 bool parameterRequiresArgument(
     ArrayRef<AnyFunctionType::Param> params,

include/swift/Sema/IDETypeChecking.h

@@ -47,7 +47,8 @@ namespace swift {
   struct PrintOptions;
 
   /// Typecheck binding initializer at \p bindingIndex.
-  void typeCheckPatternBinding(PatternBindingDecl *PBD, unsigned bindingIndex);
+  void typeCheckPatternBinding(PatternBindingDecl *PBD, unsigned bindingIndex,
+                               bool leaveClosureBodiesUnchecked);
 
   /// Check if T1 is convertible to T2.
   ///

lib/IDE/ExprContextAnalysis.cpp

@@ -98,7 +98,8 @@ void swift::ide::typeCheckContextAt(DeclContext *DC, SourceLoc Loc) {
             [](VarDecl *VD) { (void)VD->getInterfaceType(); });
         if (PBD->getInit(i)) {
           if (!PBD->isInitializerChecked(i))
-            typeCheckPatternBinding(PBD, i);
+            typeCheckPatternBinding(PBD, i,
+                                    /*LeaveClosureBodyUnchecked=*/true);
         }
       }
     } else if (auto *defaultArg = dyn_cast<DefaultArgumentInitializer>(DC)) {

lib/Sema/BuilderTransform.cpp

@@ -1932,7 +1932,9 @@ class PreCheckResultBuilderApplication : public ASTWalker {
       DiagnosticTransaction transaction(diagEngine);
 
       HasError |= ConstraintSystem::preCheckExpression(
-          E, DC, /*replaceInvalidRefsWithErrors=*/true);
+          E, DC, /*replaceInvalidRefsWithErrors=*/true,
+          /*leaveClosureBodiesUnchecked=*/false);
+
       HasError |= transaction.hasErrors();
 
       if (!HasError)

lib/Sema/CSApply.cpp

@@ -4056,8 +4056,7 @@ namespace {
       if (expr->isLiteralInit()) {
         auto *literalInit = expr->getSubExpr();
         if (auto *call = dyn_cast<CallExpr>(literalInit)) {
-          forEachExprInConstraintSystem(call->getFn(),
-                                        [&](Expr *subExpr) -> Expr * {
+          cs.forEachExpr(call->getFn(), [&](Expr *subExpr) -> Expr * {
             auto *TE = dyn_cast<TypeExpr>(subExpr);
             if (!TE)
               return subExpr;

lib/Sema/CSBindings.cpp

@@ -25,6 +25,14 @@ using namespace swift;
 using namespace constraints;
 using namespace inference;
 
+bool BindingSet::forClosureResult() const {
+  return Info.TypeVar->getImpl().isClosureResultType();
+}
+
+bool BindingSet::forGenericParameter() const {
+  return bool(Info.TypeVar->getImpl().getGenericParameter());
+}
+
 bool BindingSet::canBeNil() const {
   auto &ctx = CS.getASTContext();
   return Literals.count(

lib/Sema/CSClosure.cpp

@@ -16,6 +16,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "MiscDiagnostics.h"
 #include "TypeChecker.h"
 #include "swift/Sema/ConstraintSystem.h"
 
@@ -162,8 +163,8 @@ static bool isViableElement(ASTNode element) {
   return true;
 }
 
-using ElementInfo =
-    std::tuple<ASTNode, ContextualTypeInfo, ConstraintLocator *>;
+using ElementInfo = std::tuple<ASTNode, ContextualTypeInfo,
+                               /*isDiscarded=*/bool, ConstraintLocator *>;
 
 static void createConjunction(ConstraintSystem &cs,
                               ArrayRef<ElementInfo> elements,
@@ -190,7 +191,8 @@ static void createConjunction(ConstraintSystem &cs,
   for (const auto &entry : elements) {
     ASTNode element = std::get<0>(entry);
     ContextualTypeInfo context = std::get<1>(entry);
-    ConstraintLocator *elementLoc = std::get<2>(entry);
+    bool isDiscarded = std::get<2>(entry);
+    ConstraintLocator *elementLoc = std::get<3>(entry);
 
     if (!isViableElement(element))
       continue;
@@ -201,8 +203,8 @@ static void createConjunction(ConstraintSystem &cs,
     if (isIsolated)
       element.walk(paramCollector);
 
-    constraints.push_back(
-        Constraint::createClosureBodyElement(cs, element, context, elementLoc));
+    constraints.push_back(Constraint::createClosureBodyElement(
+        cs, element, context, elementLoc, isDiscarded));
   }
 
   // It's possible that there are no viable elements in the body,
@@ -220,8 +222,9 @@ static void createConjunction(ConstraintSystem &cs,
 }
 
 ElementInfo makeElement(ASTNode node, ConstraintLocator *locator,
-                        ContextualTypeInfo context = ContextualTypeInfo()) {
-  return std::make_tuple(node, context, locator);
+                        ContextualTypeInfo context = ContextualTypeInfo(),
+                        bool isDiscarded = false) {
+  return std::make_tuple(node, context, isDiscarded, locator);
 }
 
 static ProtocolDecl *getSequenceProtocol(ASTContext &ctx, SourceLoc loc,
@@ -751,6 +754,8 @@ class ClosureConstraintGenerator
 
   void visitBraceStmt(BraceStmt *braceStmt) {
     if (isSupportedMultiStatementClosure()) {
+      auto &ctx = cs.getASTContext();
+
       if (isChildOf(StmtKind::Case)) {
         auto *caseStmt = cast<CaseStmt>(
             locator->castLastElementTo<LocatorPathElt::ClosureBodyElement>()
@@ -765,10 +770,15 @@ class ClosureConstraintGenerator
 
       SmallVector<ElementInfo, 4> elements;
       for (auto element : braceStmt->getElements()) {
+        bool isDiscarded =
+            element.is<Expr *>() &&
+            (!ctx.LangOpts.Playground && !ctx.LangOpts.DebuggerSupport);
+
         elements.push_back(makeElement(
             element,
             cs.getConstraintLocator(
-                locator, LocatorPathElt::ClosureBodyElement(element))));
+                locator, LocatorPathElt::ClosureBodyElement(element)),
+            /*contextualInfo=*/{}, isDiscarded));
       }
 
       createConjunction(cs, elements, locator);
@@ -845,7 +855,7 @@ class ClosureConstraintGenerator
 
   bool isSupportedMultiStatementClosure() const {
     return !closure->hasSingleExpressionBody() &&
-           shouldTypeCheckInEnclosingExpression(closure);
+           cs.participatesInInference(closure);
   }
 
 #define UNSUPPORTED_STMT(STMT) void visit##STMT##Stmt(STMT##Stmt *) { \
@@ -876,7 +886,7 @@ class ClosureConstraintGenerator
 bool ConstraintSystem::generateConstraints(ClosureExpr *closure) {
   auto &ctx = closure->getASTContext();
 
-  if (shouldTypeCheckInEnclosingExpression(closure)) {
+  if (participatesInInference(closure)) {
     ClosureConstraintGenerator generator(*this, closure,
                                          getConstraintLocator(closure));
     generator.visit(closure->getBody());
@@ -925,28 +935,22 @@ bool isConditionOfStmt(ConstraintLocatorBuilder locator) {
 
 ConstraintSystem::SolutionKind
 ConstraintSystem::simplifyClosureBodyElementConstraint(
-    ASTNode element, ContextualTypeInfo context, TypeMatchOptions flags,
-    ConstraintLocatorBuilder locator) {
+    ASTNode element, ContextualTypeInfo context, bool isDiscarded,
+    TypeMatchOptions flags, ConstraintLocatorBuilder locator) {
   auto *closure = castToExpr<ClosureExpr>(locator.getAnchor());
 
   ClosureConstraintGenerator generator(*this, closure,
                                        getConstraintLocator(locator));
 
   if (auto *expr = element.dyn_cast<Expr *>()) {
-    if (context.purpose != CTP_Unused) {
-      SolutionApplicationTarget target(expr, closure, context.purpose,
-                                       context.getType(),
-                                       /*isDiscarded=*/false);
-
-      if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
-        return SolutionKind::Error;
-
-      setSolutionApplicationTarget(expr, target);
-      return SolutionKind::Solved;
-    }
+    SolutionApplicationTarget target(expr, closure, context.purpose,
+                                     context.getType(), isDiscarded);
 
-    if (!generateConstraints(expr, closure, /*isInputExpression=*/false))
+    if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
       return SolutionKind::Error;
+
+    setSolutionApplicationTarget(expr, target);
+    return SolutionKind::Solved;
   } else if (auto *stmt = element.dyn_cast<Stmt *>()) {
     generator.visit(stmt);
   } else if (auto *cond = element.dyn_cast<StmtCondition *>()) {
@@ -1160,12 +1164,17 @@ class ClosureConstraintApplication
 
     auto forEachTarget =
         rewriteTarget(*cs.getSolutionApplicationTarget(forEachStmt));
+
     if (!forEachTarget)
       hadError = true;
 
     auto body = visit(forEachStmt->getBody()).get<Stmt *>();
     forEachStmt->setBody(cast<BraceStmt>(body));
 
+    // Check to see if the sequence expr is throwing (in async context),
+    // if so require the stmt to have a `try`.
+    hadError |= diagnoseUnhandledThrowsInAsyncContext(closure, forEachStmt);
+
     return forEachStmt;
   }
 
@@ -1241,13 +1250,32 @@ class ClosureConstraintApplication
   }
 
   ASTNode visitBraceStmt(BraceStmt *braceStmt) {
+    auto &cs = solution.getConstraintSystem();
+
+    // Diagnose defer statement being last one in block.
+    if (!braceStmt->empty()) {
+      if (auto stmt = braceStmt->getLastElement().dyn_cast<Stmt *>()) {
+        if (auto deferStmt = dyn_cast<DeferStmt>(stmt)) {
+          auto &diags = closure->getASTContext().Diags;
+          diags
+              .diagnose(deferStmt->getStartLoc(), diag::defer_stmt_at_block_end)
+              .fixItReplace(deferStmt->getStartLoc(), "do");
+        }
+      }
+    }
+
     for (auto &node : braceStmt->getElements()) {
       if (auto expr = node.dyn_cast<Expr *>()) {
         // Rewrite the expression.
-        if (auto rewrittenExpr = rewriteExpr(expr))
-          node = rewrittenExpr;
-        else
+        auto target = *cs.getSolutionApplicationTarget(expr);
+        if (auto rewrittenTarget = rewriteTarget(target)) {
+          node = rewrittenTarget->getAsExpr();
+
+          if (target.isDiscardedExpr())
+            TypeChecker::checkIgnoredExpr(castToExpr(node));
+        } else {
           hadError = true;
+        }
       } else if (auto stmt = node.dyn_cast<Stmt *>()) {
         node = visit(stmt);
       } else {