[ConstraintSystem] Switch applySolutionToBody to accept AnyFunctionRef

Result builder body could be either a closure or function e.g. accessor,
which means that solution application needs to handle both.

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -647,6 +647,10 @@ T *getAsPattern(ASTNode node) {
   return nullptr;
 }
 
+template <typename T = Stmt> T *castToStmt(ASTNode node) {
+  return cast<T>(node.get<Stmt *>());
+}
+
 SourceLoc getLoc(ASTNode node);
 SourceRange getSourceRange(ASTNode node);
 
@@ -5656,14 +5660,14 @@ class ConstraintSystem {
   ///
   ///
   /// \param solution The solution to apply.
-  /// \param closure The closure to which the solution is being applied.
+  /// \param fn The function or 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,
+  bool applySolutionToBody(Solution &solution, AnyFunctionRef fn,
                            DeclContext *&currentDC,
                            std::function<Optional<SolutionApplicationTarget>(
                                SolutionApplicationTarget)>

lib/Sema/CSClosure.cpp

@@ -1755,25 +1755,40 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
 }
 
 bool ConstraintSystem::applySolutionToBody(Solution &solution,
-                                           ClosureExpr *closure,
+                                           AnyFunctionRef fn,
                                            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);
+  llvm::SaveAndRestore<DeclContext *> savedDC(currentDC, fn.getAsDeclContext());
+
+  Type resultTy;
+
+  if (auto transform = solution.getAppliedBuilderTransform(fn)) {
+    resultTy = solution.simplifyType(transform->bodyResultType);
+  } else if (auto *closure =
+                 getAsExpr<ClosureExpr>(fn.getAbstractClosureExpr())) {
+    resultTy =
+        solution.getResolvedType(closure)->castTo<FunctionType>()->getResult();
+  } else {
+    resultTy = fn.getBodyResultType();
+  }
+
+  SyntacticElementSolutionApplication application(solution, fn, resultTy,
+                                                  rewriteTarget);
 
-  auto closureType = cs.getType(closure)->castTo<FunctionType>();
-  SyntacticElementSolutionApplication application(
-      solution, closure, closureType->getResult(), rewriteTarget);
   auto body = application.apply();
 
   if (!body || application.hadError)
     return true;
 
-  closure->setBody(cast<BraceStmt>(body.get<Stmt *>()),
-                   closure->hasSingleExpressionBody());
-  closure->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
+  fn.setTypecheckedBody(castToStmt<BraceStmt>(body),
+                        fn.hasSingleExpressionBody());
+
+  if (auto *closure = getAsExpr<ClosureExpr>(fn.getAbstractClosureExpr())) {
+    closure->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
+  }
+
   return false;
 }