[Sema] Introduce PreCheckFunctionBodyRequest

Split out the part of TypeCheckFunctionBodyRequest
that inserts and removes implicit returns.

Author: hamishknight

include/swift/AST/TypeCheckRequests.h

@@ -3868,6 +3868,30 @@ class PreCheckReturnStmtRequest
   bool isCached() const { return true; }
 };
 
+/// Performs some pre-checking of a function body, including inserting and
+/// removing implict returns where needed. This request is currently
+/// side-effectful, as it will mutate the body in-place.
+///
+/// Note this request is currently uncached as:
+/// - The result will ultimately be cached by TypeCheckFunctionBodyRequest.
+/// - When re-parsing function bodies in SourceKit, we can set a previously
+///   type-checked function body back to being parsed, so caching the result of
+///   this request would result in incorrectly attempting to restore the
+///   previous body. We either need to eliminate the mutation of the AST, or
+///   implement request cache invalidation through request dependencies.
+class PreCheckFunctionBodyRequest
+    : public SimpleRequest<PreCheckFunctionBodyRequest,
+                           BraceStmt *(AbstractFunctionDecl *),
+                           RequestFlags::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  BraceStmt *evaluate(Evaluator &evaluator, AbstractFunctionDecl *AFD) const;
+};
+
 /// The result of the query for whether a statement can produce a single value.
 class IsSingleValueStmtResult {
 public:

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -325,6 +325,9 @@ SWIFT_REQUEST(TypeChecker, SynthesizeAccessorRequest,
               SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TangentStoredPropertyRequest,
               llvm::Expected<VarDecl *>(VarDecl *, CanType), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PreCheckFunctionBodyRequest,
+              BraceStmt *(AbstractFunctionDecl *),
+              Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyRequest,
               BraceStmt *(AbstractFunctionDecl *), SeparatelyCached,
               NoLocationInfo)

lib/Sema/TypeCheckStmt.cpp

@@ -2641,7 +2641,60 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
 }
 
 BraceStmt *
-TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
+PreCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
+                                      AbstractFunctionDecl *AFD) const {
+  auto &ctx = AFD->getASTContext();
+  auto *body = AFD->getBody();
+  assert(body && "Expected body");
+  assert(!AFD->isBodyTypeChecked() && "Body already type-checked?");
+
+  if (auto *func = dyn_cast<FuncDecl>(AFD)) {
+    // Don't apply this pre-checking to functions with result builders.
+    if (getResultBuilderType(func))
+      return body;
+
+    if (func->hasSingleExpressionBody() &&
+        func->getResultInterfaceType()->isVoid()) {
+      // The function returns void.  We don't need an explicit return, no
+      // matter what the type of the expression is. Take the inserted return
+      // back out.
+      body->setLastElement(func->getSingleExpressionBody());
+    }
+    // If there is a single statement in the body that can be turned into a
+    // single expression return, do so now.
+    if (!func->getResultInterfaceType()->isVoid()) {
+      if (auto *S = body->getSingleActiveStatement()) {
+        if (S->mayProduceSingleValue(evaluator)) {
+          auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
+              ctx, S, /*DC*/ func, /*mustBeExpr*/ false);
+          auto *RS = new (ctx) ReturnStmt(SourceLoc(), SVE);
+          body->setLastElement(RS);
+          func->setHasSingleExpressionBody();
+          func->setSingleExpressionBody(SVE);
+        }
+      }
+    }
+  }
+
+  if (auto *ctor = dyn_cast<ConstructorDecl>(AFD)) {
+    if (body->empty() || !isKnownEndOfConstructor(body->getLastElement())) {
+      // For constructors, we make sure that the body ends with a "return"
+      // stmt, which we either implicitly synthesize, or the user can write.
+      // This simplifies SILGen.
+      SmallVector<ASTNode, 8> Elts(body->getElements().begin(),
+                                   body->getElements().end());
+      Elts.push_back(new (ctx) ReturnStmt(body->getRBraceLoc(),
+                                          /*value*/ nullptr,
+                                          /*implicit*/ true));
+      body = BraceStmt::create(ctx, body->getLBraceLoc(), Elts,
+                               body->getRBraceLoc(), body->isImplicit());
+    }
+  }
+  return body;
+}
+
+BraceStmt *
+TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
                                        AbstractFunctionDecl *AFD) const {
   ASTContext &ctx = AFD->getASTContext();
 
@@ -2672,6 +2725,12 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
                              range.End);
   };
 
+  // First do a pre-check of the body.
+  body = evaluateOrDefault(eval, PreCheckFunctionBodyRequest{AFD}, nullptr);
+  assert(body);
+
+  // Then apply a result builder if we have one, which if successful will
+  // produce a type-checked body.
   bool alreadyTypeChecked = false;
   if (auto *func = dyn_cast<FuncDecl>(AFD)) {
     if (Type builderType = getResultBuilderType(func)) {
@@ -2686,42 +2745,6 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
 
         body->walk(ContextualizeClosuresAndMacros(AFD));
       }
-    } else {
-      if (func->hasSingleExpressionBody() &&
-          func->getResultInterfaceType()->isVoid()) {
-        // The function returns void.  We don't need an explicit return, no
-        // matter what the type of the expression is. Take the inserted return
-        // back out.
-        body->setLastElement(func->getSingleExpressionBody());
-      }
-      // If there is a single statement in the body that can be turned into a
-      // single expression return, do so now.
-      if (!func->getResultInterfaceType()->isVoid()) {
-        if (auto *S = body->getSingleActiveStatement()) {
-          if (S->mayProduceSingleValue(evaluator)) {
-            auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
-                ctx, S, /*DC*/ func, /*mustBeExpr*/ false);
-            auto *RS = new (ctx) ReturnStmt(SourceLoc(), SVE);
-            body->setLastElement(RS);
-            func->setHasSingleExpressionBody();
-            func->setSingleExpressionBody(SVE);
-          }
-        }
-      }
-    }
-  } else if (auto *ctor = dyn_cast<ConstructorDecl>(AFD)) {
-    if (body->empty() ||
-        !isKnownEndOfConstructor(body->getLastElement())) {
-      // For constructors, we make sure that the body ends with a "return" stmt,
-      // which we either implicitly synthesize, or the user can write.  This
-      // simplifies SILGen.
-      SmallVector<ASTNode, 8> Elts(body->getElements().begin(),
-                                   body->getElements().end());
-      Elts.push_back(new (ctx) ReturnStmt(body->getRBraceLoc(),
-                                          /*value*/nullptr,
-                                          /*implicit*/true));
-      body = BraceStmt::create(ctx, body->getLBraceLoc(), Elts,
-                               body->getRBraceLoc(), body->isImplicit());
     }
   }