[TypeChecker] Fix constraint solver to respect LeaveClosureBodyUnchecked flag

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -5014,7 +5014,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.
   ///
@@ -5245,6 +5246,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 *));
 
@@ -6069,18 +6080,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

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

@@ -466,7 +466,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;
       });
@@ -3594,7 +3595,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

@@ -2105,7 +2105,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));
       }();
@@ -2627,7 +2627,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

@@ -2040,7 +2040,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;
@@ -11662,7 +11662,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

@@ -5874,6 +5874,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

@@ -893,6 +893,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;
@@ -1080,9 +1082,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; }
 
@@ -1428,8 +1432,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.
@@ -2082,11 +2091,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

@@ -206,7 +206,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);
@@ -305,8 +307,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();
@@ -788,7 +793,8 @@ bool TypeChecker::typeCheckForCodeCompletion(
     // expression and folding sequence expressions.
     auto failedPreCheck = ConstraintSystem::preCheckExpression(
         expr, DC,
-        /*replaceInvalidRefsWithErrors=*/true);
+        /*replaceInvalidRefsWithErrors=*/true,
+        /*leaveClosureBodiesUnchecked=*/true);
 
     target.setExpr(expr);
 
@@ -892,7 +898,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) {

lib/Sema/TypeCheckConstraints.cpp

@@ -344,7 +344,8 @@ TypeChecker::typeCheckExpression(
   // First, pre-check the expression, validating any types that occur in the
   // expression and folding sequence expressions.
   if (ConstraintSystem::preCheckExpression(
-          expr, dc, /*replaceInvalidRefsWithErrors=*/true)) {
+        expr, dc, /*replaceInvalidRefsWithErrors=*/true,
+        options.contains(TypeCheckExprFlags::LeaveClosureBodyUnchecked))) {
     target.setExpr(expr);
     return None;
   }
@@ -588,14 +589,17 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
   // Precheck the sequence.
   Expr *sequence = stmt->getSequence();
   if (ConstraintSystem::preCheckExpression(
-          sequence, dc, /*replaceInvalidRefsWithErrors=*/true))
+          sequence, dc, /*replaceInvalidRefsWithErrors=*/true,
+          /*leaveClosureBodiesUnchecked=*/false))
     return failed();
   stmt->setSequence(sequence);
 
   // Precheck the filtering condition.
   if (Expr *whereExpr = stmt->getWhere()) {
     if (ConstraintSystem::preCheckExpression(
-            whereExpr, dc, /*replaceInvalidRefsWithErrors=*/true))
+            whereExpr, dc,
+            /*replaceInvalidRefsWithErrors=*/true,
+            /*leaveClosureBodiesUnchecked=*/false))
       return failed();
 
     stmt->setWhere(whereExpr);
@@ -2052,26 +2056,19 @@ HasDynamicCallableAttributeRequest::evaluate(Evaluator &evaluator,
   });
 }
 
-bool swift::shouldTypeCheckInEnclosingExpression(ClosureExpr *expr) {
-  if (expr->hasSingleExpressionBody())
-    return true;
-
-  auto &ctx = expr->getASTContext();
-  return !expr->hasEmptyBody() &&
-         ctx.TypeCheckerOpts.EnableMultiStatementClosureInference;
-}
-
-void swift::forEachExprInConstraintSystem(
+void ConstraintSystem::forEachExpr(
     Expr *expr, llvm::function_ref<Expr *(Expr *)> callback) {
   struct ChildWalker : ASTWalker {
+    ConstraintSystem &CS;
     llvm::function_ref<Expr *(Expr *)> callback;
 
-    ChildWalker(llvm::function_ref<Expr *(Expr *)> callback)
-    : callback(callback) {}
+    ChildWalker(ConstraintSystem &CS,
+                llvm::function_ref<Expr *(Expr *)> callback)
+        : CS(CS), callback(callback) {}
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       if (auto closure = dyn_cast<ClosureExpr>(E)) {
-        if (!shouldTypeCheckInEnclosingExpression(closure))
+        if (!CS.participatesInInference(closure))
           return { false, callback(E) };
       }
       return { true, callback(E) };
@@ -2084,5 +2081,5 @@ void swift::forEachExprInConstraintSystem(
     bool walkToTypeReprPre(TypeRepr *T) override { return false; }
   };
 
-  expr->walk(ChildWalker(callback));
+  expr->walk(ChildWalker(*this, callback));
 }