[TypeChecker] Fix constraint solver to respect LeaveClosureBodyUnchecked flag

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -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,

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=*/true);
+
       HasError |= transaction.hasErrors();
 
       if (!HasError)

lib/Sema/CSApply.cpp

@@ -4107,8 +4107,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/CSClosure.cpp

@@ -854,7 +854,7 @@ class ClosureConstraintGenerator
 
   bool isSupportedMultiStatementClosure() const {
     return !closure->hasSingleExpressionBody() &&
-           shouldTypeCheckInEnclosingExpression(closure);
+           cs.participatesInInference(closure);
   }
 
 #define UNSUPPORTED_STMT(STMT) void visit##STMT##Stmt(STMT##Stmt *) { \
@@ -885,7 +885,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());

lib/Sema/CSDiagnostics.cpp

@@ -463,7 +463,8 @@ bool MissingConformanceFailure::diagnoseAsError() {
       }
 
       bool hasFix = false;
-      forEachExprInConstraintSystem(caseExpr, [&](Expr *expr) -> Expr * {
+      auto &cs = getConstraintSystem();
+      cs.forEachExpr(caseExpr, [&](Expr *expr) -> Expr * {
         hasFix |= anchors.count(expr);
         return hasFix ? nullptr : expr;
       });
@@ -3591,7 +3592,8 @@ bool MissingMemberFailure::diagnoseAsError() {
 
     bool hasUnresolvedPattern = false;
     if (auto *E = getAsExpr(anchor)) {
-      forEachExprInConstraintSystem(const_cast<Expr *>(E), [&](Expr *expr) {
+      auto &cs = getConstraintSystem();
+      cs.forEachExpr(const_cast<Expr *>(E), [&](Expr *expr) {
         hasUnresolvedPattern |= isa<UnresolvedPatternExpr>(expr);
         return hasUnresolvedPattern ? nullptr : expr;
       });

lib/Sema/CSGen.cpp

@@ -2090,7 +2090,7 @@ namespace {
         // If this is a multi-statement closure, let's mark result
         // as potential hole right away.
         return Type(CS.createTypeVariable(
-            resultLocator, shouldTypeCheckInEnclosingExpression(closure)
+            resultLocator, CS.participatesInInference(closure)
                                ? 0
                                : TVO_CanBindToHole));
       }();
@@ -2612,7 +2612,7 @@ namespace {
       // genreation only if closure is going to participate
       // in the type-check. This allows us to delay validation of
       // multi-statement closures until body is opened.
-      if (shouldTypeCheckInEnclosingExpression(closure) &&
+      if (CS.participatesInInference(closure) &&
           collectVarRefs.hasErrorExprs) {
         return Type();
       }

lib/Sema/CSSimplify.cpp

@@ -2011,7 +2011,7 @@ static bool fixMissingArguments(ConstraintSystem &cs, ASTNode anchor,
 
       // Something like `foo { x in }` or `foo { $0 }`
       if (auto *closure = getAsExpr<ClosureExpr>(anchor)) {
-        forEachExprInConstraintSystem(closure, [&](Expr *expr) -> Expr * {
+        cs.forEachExpr(closure, [&](Expr *expr) -> Expr * {
           if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr)) {
             if (!isParam(UDE->getBase()))
               return expr;
@@ -11546,7 +11546,7 @@ bool ConstraintSystem::recordFix(ConstraintFix *fix, unsigned impact) {
 
   bool found = false;
   if (auto *expr = getAsExpr(anchor)) {
-    forEachExprInConstraintSystem(expr, [&](Expr *subExpr) -> Expr * {
+    forEachExpr(expr, [&](Expr *subExpr) -> Expr * {
       found |= anchors.count(subExpr);
       return subExpr;
     });

lib/Sema/ConstraintSystem.cpp

@@ -5845,6 +5845,18 @@ bool ConstraintSystem::isReadOnlyKeyPathComponent(
   return false;
 }
 
+bool ConstraintSystem::participatesInInference(ClosureExpr *closure) const {
+  if (closure->hasSingleExpressionBody())
+    return true;
+
+  if (Options.contains(ConstraintSystemFlags::LeaveClosureBodyUnchecked))
+    return false;
+
+  auto &ctx = closure->getASTContext();
+  return !closure->hasEmptyBody() &&
+         ctx.TypeCheckerOpts.EnableMultiStatementClosureInference;
+}
+
 TypeVarBindingProducer::TypeVarBindingProducer(BindingSet &bindings)
     : BindingProducer(bindings.getConstraintSystem(),
                       bindings.getTypeVariable()->getImpl().getLocator()),

lib/Sema/PreCheckExpr.cpp

@@ -894,6 +894,8 @@ namespace {
     /// implicit `ErrorExpr` in place of invalid references.
     bool UseErrorExprs;
 
+    bool LeaveClosureBodiesUnchecked;
+
     /// A stack of expressions being walked, used to determine where to
     /// insert RebindSelfInConstructorExpr nodes.
     llvm::SmallVector<Expr *, 8> ExprStack;
@@ -1081,9 +1083,11 @@ namespace {
 
   public:
     PreCheckExpression(DeclContext *dc, Expr *parent,
-                       bool replaceInvalidRefsWithErrors)
-        : Ctx(dc->getASTContext()), DC(dc), ParentExpr(parent),
-          UseErrorExprs(replaceInvalidRefsWithErrors) {}
+                       bool replaceInvalidRefsWithErrors,
+                       bool leaveClosureBodiesUnchecked)
+        : Ctx(dc->getASTContext()), DC(dc),
+          ParentExpr(parent), UseErrorExprs(replaceInvalidRefsWithErrors),
+          LeaveClosureBodiesUnchecked(leaveClosureBodiesUnchecked) {}
 
     ASTContext &getASTContext() const { return Ctx; }
 
@@ -1429,8 +1433,13 @@ namespace {
 /// true when we want the body to be considered part of this larger expression.
 bool PreCheckExpression::walkToClosureExprPre(ClosureExpr *closure) {
   // If we won't be checking the body of the closure, don't walk into it here.
-  if (!shouldTypeCheckInEnclosingExpression(closure))
-    return false;
+  if (!closure->hasSingleExpressionBody()) {
+    if (LeaveClosureBodiesUnchecked)
+      return false;
+
+    if (!Ctx.TypeCheckerOpts.EnableMultiStatementClosureInference)
+      return false;
+  }
 
   // Update the current DeclContext to be the closure we're about to
   // recurse into.
@@ -2086,11 +2095,15 @@ Expr *PreCheckExpression::simplifyTypeConstructionWithLiteralArg(Expr *E) {
 /// Pre-check the expression, validating any types that occur in the
 /// expression and folding sequence expressions.
 bool ConstraintSystem::preCheckExpression(Expr *&expr, DeclContext *dc,
-                                          bool replaceInvalidRefsWithErrors) {
+                                          bool replaceInvalidRefsWithErrors,
+                                          bool leaveClosureBodiesUnchecked) {
   auto &ctx = dc->getASTContext();
   FrontendStatsTracer StatsTracer(ctx.Stats, "precheck-expr", expr);
 
-  PreCheckExpression preCheck(dc, expr, replaceInvalidRefsWithErrors);
+  PreCheckExpression preCheck(dc, expr,
+                              replaceInvalidRefsWithErrors,
+                              leaveClosureBodiesUnchecked);
+
   // Perform the pre-check.
   if (auto result = expr->walk(preCheck)) {
     expr = result;

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -236,7 +236,9 @@ class SanitizeExpr : public ASTWalker {
       // If this is a closure, only walk into its children if they
       // are type-checked in the context of the enclosing expression.
       if (auto closure = dyn_cast<ClosureExpr>(expr)) {
-        if (!shouldTypeCheckInEnclosingExpression(closure))
+        // TODO: This has to be deleted once `EnableMultiStatementClosureInference`
+        //       is enabled by default.
+        if (!closure->hasSingleExpressionBody())
           return { false, expr };
         for (auto &Param : *closure->getParameters()) {
           Param->setSpecifier(swift::ParamSpecifier::Default);
@@ -335,8 +337,11 @@ getTypeOfExpressionWithoutApplying(Expr *&expr, DeclContext *dc,
   PrettyStackTraceExpr stackTrace(Context, "type-checking", expr);
   referencedDecl = nullptr;
 
+  ConstraintSystemOptions options;
+  options |= ConstraintSystemFlags::SuppressDiagnostics;
+
   // Construct a constraint system from this expression.
-  ConstraintSystem cs(dc, ConstraintSystemFlags::SuppressDiagnostics);
+  ConstraintSystem cs(dc, options);
 
   // Attempt to solve the constraint system.
   const Type originalType = expr->getType();
@@ -818,7 +823,8 @@ bool TypeChecker::typeCheckForCodeCompletion(
     // expression and folding sequence expressions.
     auto failedPreCheck = ConstraintSystem::preCheckExpression(
         expr, DC,
-        /*replaceInvalidRefsWithErrors=*/true);
+        /*replaceInvalidRefsWithErrors=*/true,
+        /*leaveClosureBodiesUnchecked=*/true);
 
     target.setExpr(expr);
 
@@ -922,7 +928,9 @@ static Optional<Type> getTypeOfCompletionContextExpr(
                         Expr *&parsedExpr,
                         ConcreteDeclRef &referencedDecl) {
   if (constraints::ConstraintSystem::preCheckExpression(
-          parsedExpr, DC, /*replaceInvalidRefsWithErrors=*/true))
+          parsedExpr, DC,
+          /*replaceInvalidRefsWithErrors=*/true,
+          /*leaveClosureBodiesUnchecked=*/true))
     return None;
 
   switch (kind) {