[CSApply] Delay solution application to multi-statement closure bodies

Let's delay solution application to multi-statement closure bodies,
because declarations found in the body of such closures may depend
on its `ExtInfo` flags e.g. no-escape is set based on parameter if
closure is used in a call.

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -5149,6 +5149,23 @@ class ConstraintSystem {
         Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>
           rewriteTarget);
 
+  /// Apply the given solution to the given closure body.
+  ///
+  ///
+  /// \param solution The solution to apply.
+  /// \param closure The closure to which the solution is being applied.
+  /// \param currentDC The declaration context in which transformations
+  /// will be applied.
+  /// \param rewriteTarget Function that performs a rewrite of any
+  /// solution application target within the context.
+  ///
+  /// \returns true if solution cannot be applied.
+  bool applySolutionToBody(Solution &solution, ClosureExpr *closure,
+                           DeclContext *&currentDC,
+                           std::function<Optional<SolutionApplicationTarget>(
+                               SolutionApplicationTarget)>
+                               rewriteTarget);
+
   /// Reorder the disjunctive clauses for a given expression to
   /// increase the likelihood that a favored constraint will be successfully
   /// resolved before any others.

lib/Sema/CSApply.cpp

@@ -8011,18 +8011,66 @@ namespace {
   public:
     ExprWalker(ExprRewriter &Rewriter) : Rewriter(Rewriter) { }
 
+    ~ExprWalker() {
+      assert(ClosuresToTypeCheck.empty());
+      assert(TapsToTypeCheck.empty());
+    }
+
     bool shouldWalkIntoPropertyWrapperPlaceholderValue() override {
       // Property wrapper placeholder underlying values are filled in
       // with already-type-checked expressions. Don't walk into them.
       return false;
     }
 
-    const SmallVectorImpl<ClosureExpr *> &getClosuresToTypeCheck() const {
-      return ClosuresToTypeCheck;
-    }
+    /// Process delayed closure bodies and `Tap` expressions.
+    ///
+    /// \returns true if any part of the processing fails.
+    bool processDelayed() {
+      bool hadError = false;
+
+      while (!ClosuresToTypeCheck.empty()) {
+        auto *closure = ClosuresToTypeCheck.pop_back_val();
+        // If experimental multi-statement closure support
+        // is enabled, solution should have all of required
+        // information.
+        //
+        // Note that in this mode `ClosuresToTypeCheck` acts
+        // as a stack because multi-statement closures could
+        // have other multi-statement closures in the body.
+        auto &ctx = closure->getASTContext();
+        if (ctx.TypeCheckerOpts.EnableMultiStatementClosureInference) {
+          auto &solution = Rewriter.solution;
+          auto &cs = solution.getConstraintSystem();
+
+          hadError |= cs.applySolutionToBody(
+              solution, closure, Rewriter.dc,
+              [&](SolutionApplicationTarget target) {
+                auto resultTarget = rewriteTarget(target);
+                if (resultTarget) {
+                  if (auto expr = resultTarget->getAsExpr())
+                    solution.setExprTypes(expr);
+                }
+
+                return resultTarget;
+              });
+          continue;
+        }
 
-    const SmallVectorImpl<std::pair<TapExpr *, DeclContext *>> &getTapsToTypeCheck() const {
-      return TapsToTypeCheck;
+        hadError |= TypeChecker::typeCheckClosureBody(closure);
+      }
+
+      // Tap expressions too; they should or should not be
+      // type-checked under the same conditions as closure bodies.
+      {
+        for (const auto &tuple : TapsToTypeCheck) {
+          auto tap = std::get<0>(tuple);
+          auto tapDC = std::get<1>(tuple);
+          hadError |= TypeChecker::typeCheckTapBody(tap, tapDC);
+        }
+        TapsToTypeCheck.clear();
+      }
+
+      return hadError;
     }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
@@ -8819,19 +8867,10 @@ Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
   if (!resultTarget)
     return None;
 
-  // Visit closures that have non-single expression bodies.
-  bool hadError = false;
-
-  for (auto *closure : walker.getClosuresToTypeCheck())
-    hadError |= TypeChecker::typeCheckClosureBody(closure);
-
-  // Tap expressions too; they should or should not be
-  // type-checked under the same conditions as closure bodies.
-  for (auto tuple : walker.getTapsToTypeCheck()) {
-    auto tap = std::get<0>(tuple);
-    auto tapDC = std::get<1>(tuple);
-    hadError |= TypeChecker::typeCheckTapBody(tap, tapDC);
-  }
+  // Visit closures that have non-single expression bodies, tap expressions,
+  // and possibly other types of AST nodes which could only be processed
+  // after contextual expression.
+  bool hadError = walker.processDelayed();
 
   // If any of them failed to type check, bail.
   if (hadError)

lib/Sema/CSClosure.cpp

@@ -1268,17 +1268,41 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
 
   // If this closure is checked as part of the enclosing expression, handle
   // that now.
-  if (shouldTypeCheckInEnclosingExpression(closure)) {
-    ClosureConstraintApplication application(
-        solution, closure, closureFnType->getResult(), rewriteTarget);
-    application.visit(fn.getBody());
-    closure->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
-
-    return SolutionApplicationToFunctionResult::Success;
+  //
+  // Multi-statement closures are handled separately because they need to
+  // wait until all of the `ExtInfo` flags are propagated from the context
+  // e.g. parameter could be no-escape if closure is applied to a call.
+  if (closure->hasSingleExpressionBody()) {
+    bool hadError =
+        applySolutionToBody(solution, closure, currentDC, rewriteTarget);
+    return hadError ? SolutionApplicationToFunctionResult::Failure
+                    : SolutionApplicationToFunctionResult::Success;
   }
 
   // Otherwise, we need to delay type checking of the closure until later.
   solution.setExprTypes(closure);
   closure->setBodyState(ClosureExpr::BodyState::ReadyForTypeChecking);
   return SolutionApplicationToFunctionResult::Delay;
 }
+
+bool ConstraintSystem::applySolutionToBody(Solution &solution,
+                                           ClosureExpr *closure,
+                                           DeclContext *&currentDC,
+                                           RewriteTargetFn rewriteTarget) {
+  auto &cs = solution.getConstraintSystem();
+  // Enter the context of the function before performing any additional
+  // transformations.
+  llvm::SaveAndRestore<DeclContext *> savedDC(currentDC, closure);
+
+  auto closureType = cs.getType(closure)->castTo<FunctionType>();
+  ClosureConstraintApplication application(
+      solution, closure, closureType->getResult(), rewriteTarget);
+  application.visit(closure->getBody());
+
+  if (application.hadError)
+    return true;
+
+  TypeChecker::checkClosureAttributes(closure);
+  closure->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
+  return false;
+}