[Type checker] Fold typeCheckFunctionBuilderFuncBody() so it isn't so special

Instead of being a completely separate body type-checking pass, refactor
this function so it does the function builder transform on the body, then
continues with the rest of typeCheckFunctionBodyUntil() as normal.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -476,47 +476,25 @@ static bool hasReturnStmt(Stmt *stmt) {
   return finder.hasReturnStmt;
 }
 
-bool TypeChecker::typeCheckFunctionBuilderFuncBody(FuncDecl *FD,
-                                                   Type builderType) {
+BraceStmt *
+TypeChecker::applyFunctionBuilderBodyTransform(FuncDecl *FD,
+                                               BraceStmt *body,
+                                               Type builderType) {
   // Try to build a single result expression.
   BuilderClosureVisitor visitor(Context, nullptr,
                                 /*wantExpr=*/true, builderType);
-  Expr *returnExpr = visitor.visit(FD->getBody());
+  Expr *returnExpr = visitor.visit(body);
   if (!returnExpr)
-    return true;
+    return nullptr;
 
   // Make sure we have a usable result type for the body.
   Type returnType = AnyFunctionRef(FD).getBodyResultType();
   if (!returnType || returnType->hasError())
-    return true;
-
-  TypeCheckExprOptions options = {};
-  if (auto opaque = returnType->getAs<OpaqueTypeArchetypeType>()) {
-    if (opaque->getDecl()->isOpaqueReturnTypeOfFunction(FD))
-      options |= TypeCheckExprFlags::ConvertTypeIsOpaqueReturnType;
-  }
-
-  // If we are performing code-completion inside the functions body, supress
-  // diagnostics to workaround typechecking performance problems.
-  if (Context.SourceMgr.rangeContainsCodeCompletionLoc(
-          FD->getBody()->getSourceRange()))
-    options |= TypeCheckExprFlags::SuppressDiagnostics;
-
-  // Type-check the single result expression.
-  Type returnExprType = typeCheckExpression(returnExpr, FD,
-                                            TypeLoc::withoutLoc(returnType),
-                                            CTP_ReturnStmt, options);
-  if (!returnExprType)
-    return true;
-  assert(returnExprType->isEqual(returnType));
+    return nullptr;
 
   auto returnStmt = new (Context) ReturnStmt(SourceLoc(), returnExpr);
-  auto origBody = FD->getBody();
-  auto fakeBody = BraceStmt::create(Context, origBody->getLBraceLoc(),
-                                    { returnStmt },
-                                    origBody->getRBraceLoc());
-  FD->setBody(fakeBody);
-  return false;
+  return BraceStmt::create(Context, body->getLBraceLoc(), { returnStmt },
+                           body->getRBraceLoc());
 }
 
 ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(

lib/Sema/TypeCheckStmt.cpp

@@ -524,6 +524,14 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
       if (func->hasSingleExpressionBody()) {
         ctp = CTP_ReturnSingleExpr;
       }
+
+      // If we are performing code-completion inside a function builder body,
+      // suppress diagnostics to workaround typechecking performance problems.
+      if (func->getFunctionBuilderType() &&
+          TC.Context.SourceMgr.rangeContainsCodeCompletionLoc(
+              func->getBody()->getSourceRange())) {
+        options |= TypeCheckExprFlags::SuppressDiagnostics;
+      }
     }
 
     auto exprTy = TC.typeCheckExpression(E, DC, TypeLoc::withoutLoc(ResultTy),
@@ -2166,9 +2174,9 @@ bool TypeChecker::typeCheckAbstractFunctionBodyUntil(AbstractFunctionDecl *AFD,
 
   if (auto *func = dyn_cast<FuncDecl>(AFD)) {
     if (Type builderType = getFunctionBuilderType(func)) {
-      // FIXME: this should be a transform on the body, not a separate type
-      // checking operation.
-      return typeCheckFunctionBuilderFuncBody(func, builderType);
+      body = applyFunctionBuilderBodyTransform(func, body, builderType);
+      if (!body)
+        return true;
     } else if (func->hasSingleExpressionBody()) {
       auto resultTypeLoc = func->getBodyResultTypeLoc();
       auto expr = func->getSingleExpressionBody();

lib/Sema/TypeChecker.h

@@ -981,7 +981,9 @@ class TypeChecker final : public LazyResolver {
                                           SourceLoc EndTypeCheckLoc);
   bool typeCheckAbstractFunctionBody(AbstractFunctionDecl *AFD);
 
-  bool typeCheckFunctionBuilderFuncBody(FuncDecl *FD, Type builderType);
+  BraceStmt *applyFunctionBuilderBodyTransform(FuncDecl *FD,
+                                               BraceStmt *body,
+                                               Type builderType);
   bool typeCheckClosureBody(ClosureExpr *closure);
 
   bool typeCheckTapBody(TapExpr *expr, DeclContext *DC);