[CodeCompletion] Separate TypeCheckFunctionBodyAtLocRequest

from TypeCheckFunctionBodyRequest, and only do necessary operations.

Author: rintaro

include/swift/AST/TypeCheckRequests.h

@@ -867,14 +867,13 @@ class LazyStoragePropertyRequest :
   bool isCached() const { return true; }
 };
 
-/// Request to type check the body of the given function up to the given
-/// source location.
+/// Request to type check the body of the given function.
 ///
 /// Produces true if an error occurred, false otherwise.
 /// FIXME: it would be far better to return the type-checked body.
 class TypeCheckFunctionBodyRequest :
     public SimpleRequest<TypeCheckFunctionBodyRequest,
-                         bool(AbstractFunctionDecl *, SourceLoc),
+                         bool(AbstractFunctionDecl *),
                          RequestFlags::Cached|RequestFlags::DependencySource> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -883,9 +882,7 @@ class TypeCheckFunctionBodyRequest :
   friend SimpleRequest;
 
   // Evaluation.
-  bool
-  evaluate(Evaluator &evaluator, AbstractFunctionDecl *func,
-           SourceLoc endTypeCheckLoc) const;
+  bool evaluate(Evaluator &evaluator, AbstractFunctionDecl *func) const;
 
 public:
   bool isCached() const { return true; }
@@ -896,6 +893,24 @@ class TypeCheckFunctionBodyRequest :
   readDependencySource(const evaluator::DependencyRecorder &) const;
 };
 
+/// Request to typecheck a function body element at the given source location.
+///
+/// Produces true if an error occurred, false otherwise.
+class TypeCheckFunctionBodyAtLocRequest
+    : public SimpleRequest<TypeCheckFunctionBodyAtLocRequest,
+                           bool(AbstractFunctionDecl *, SourceLoc),
+                           RequestFlags::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  bool evaluate(Evaluator &evaluator, AbstractFunctionDecl *func,
+                SourceLoc Loc) const;
+};
+
 /// Request to obtain a list of stored properties in a nominal type.
 ///
 /// This will include backing storage for lazy properties and

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -197,8 +197,10 @@ SWIFT_REQUEST(TypeChecker, SynthesizeAccessorRequest,
               AccessorDecl *(AbstractStorageDecl *, AccessorKind),
               SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyRequest,
+              bool(AbstractFunctionDecl *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyAtLocRequest,
               bool(AbstractFunctionDecl *, SourceLoc),
-              Cached, NoLocationInfo)
+              Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, UnderlyingTypeRequest, Type(TypeAliasDecl *),
               SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, USRGenerationRequest, std::string(const ValueDecl *),

include/swift/Sema/IDETypeChecking.h

@@ -133,9 +133,9 @@ namespace swift {
   /// Typecheck the given expression.
   bool typeCheckExpression(DeclContext *DC, Expr *&parsedExpr);
 
-  /// Partially typecheck the specified function body.
-  bool typeCheckAbstractFunctionBodyNodeAt(AbstractFunctionDecl *AFD,
-                                           SourceLoc TargetLoc);
+  /// Type check a function body element which is at \p TagetLoc .
+  bool typeCheckAbstractFunctionBodyAtLoc(AbstractFunctionDecl *AFD,
+                                          SourceLoc TargetLoc);
 
   /// Typecheck top-level code parsed during code completion.
   ///

lib/IDE/ExprContextAnalysis.cpp

@@ -79,7 +79,7 @@ void typeCheckContextImpl(DeclContext *DC, SourceLoc Loc) {
     auto &SM = DC->getASTContext().SourceMgr;
     auto bodyRange = AFD->getBodySourceRange();
     if (SM.rangeContainsTokenLoc(bodyRange, Loc)) {
-      swift::typeCheckAbstractFunctionBodyNodeAt(AFD, Loc);
+      swift::typeCheckAbstractFunctionBodyAtLoc(AFD, Loc);
     } else {
       assert(bodyRange.isInvalid() && "The body should not be parsed if the "
                                       "completion happens in the signature");
@@ -866,10 +866,10 @@ class ExprContextAnalyzer {
       break;
     }
     default:
-      if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-        assert(isSingleExpressionBodyForCodeCompletion(AFD->getBody()));
+      if (auto *FD = dyn_cast<FuncDecl>(D)) {
+        assert(isSingleExpressionBodyForCodeCompletion(FD->getBody()));
         singleExpressionBody = true;
-        recordPossibleType(getReturnTypeFromContext(AFD));
+        recordPossibleType(getReturnTypeFromContext(FD));
         break;
       }
       llvm_unreachable("Unhandled decl kind.");
@@ -1003,8 +1003,8 @@ class ExprContextAnalyzer {
         case DeclKind::PatternBinding:
           return true;
         default:
-          if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D))
-            if (auto *body = AFD->getBody())
+          if (auto *FD = dyn_cast<FuncDecl>(D))
+            if (auto *body = FD->getBody())
               return isSingleExpressionBodyForCodeCompletion(body);
           return false;
         }

lib/Sema/TypeCheckStmt.cpp

@@ -1635,9 +1635,8 @@ static Type getFunctionBuilderType(FuncDecl *FD) {
 }
 
 bool TypeChecker::typeCheckAbstractFunctionBody(AbstractFunctionDecl *AFD) {
-  auto res =
-      evaluateOrDefault(AFD->getASTContext().evaluator,
-                        TypeCheckFunctionBodyRequest{AFD, SourceLoc()}, true);
+  auto res = evaluateOrDefault(AFD->getASTContext().evaluator,
+                               TypeCheckFunctionBodyRequest{AFD}, true);
   TypeChecker::checkFunctionErrorHandling(AFD);
   TypeChecker::computeCaptures(AFD);
   return res;
@@ -1836,10 +1835,41 @@ static void checkClassConstructorBody(ClassDecl *classDecl,
   }
 }
 
+bool TypeCheckFunctionBodyAtLocRequest::evaluate(Evaluator &evaluator,
+                                                 AbstractFunctionDecl *AFD,
+                                                 SourceLoc Loc) const {
+  ASTContext &ctx = AFD->getASTContext();
+
+  BraceStmt *body = AFD->getBody();
+  if (!body || AFD->isBodyTypeChecked())
+    return false;
+
+  // Function builder doesn't support partial type checking.
+  if (auto *func = dyn_cast<FuncDecl>(AFD)) {
+    if (Type builderType = getFunctionBuilderType(func)) {
+      auto optBody =
+          TypeChecker::applyFunctionBuilderBodyTransform(func, builderType);
+      if (!optBody || !*optBody)
+        return true;
+      // Wire up the function body now.
+      AFD->setBody(*optBody, AbstractFunctionDecl::BodyKind::TypeChecked);
+      return false;
+    }
+  }
+
+  if (ctx.LangOpts.EnableASTScopeLookup)
+    ASTScope::expandFunctionBody(AFD);
+
+  StmtChecker SC(AFD);
+  SC.TargetTypeCheckLoc = Loc;
+  bool hadError = SC.typeCheckBody(body);
+  AFD->setBody(body, AbstractFunctionDecl::BodyKind::TypeChecked);
+  return hadError;
+}
+
 bool
 TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
-                                       AbstractFunctionDecl *AFD,
-                                       SourceLoc targetTypeCheckLoc) const {
+                                       AbstractFunctionDecl *AFD) const {
   ASTContext &ctx = AFD->getASTContext();
 
   Optional<FunctionBodyTimer> timer;
@@ -1896,7 +1926,6 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
   bool hadError = false;
   if (!alreadyTypeChecked) {
     StmtChecker SC(AFD);
-    SC.TargetTypeCheckLoc = targetTypeCheckLoc;
     hadError = SC.typeCheckBody(body);
   }
 
@@ -1927,7 +1956,7 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
   AFD->setBody(body, AbstractFunctionDecl::BodyKind::TypeChecked);
 
   // If nothing went wrong yet, perform extra checking.
-  if (!hadError && targetTypeCheckLoc.isInvalid())
+  if (!hadError)
     performAbstractFuncDeclDiagnostics(AFD);
 
   return hadError;

lib/Sema/TypeChecker.cpp

@@ -540,12 +540,13 @@ bool swift::typeCheckExpression(DeclContext *DC, Expr *&parsedExpr) {
   return !resultTy;
 }
 
-bool swift::typeCheckAbstractFunctionBodyNodeAt(AbstractFunctionDecl *AFD,
-                                                SourceLoc TargetLoc) {
+bool swift::typeCheckAbstractFunctionBodyAtLoc(AbstractFunctionDecl *AFD,
+                                               SourceLoc TargetLoc) {
   auto &Ctx = AFD->getASTContext();
   DiagnosticSuppression suppression(Ctx.Diags);
   return !evaluateOrDefault(Ctx.evaluator,
-                            TypeCheckFunctionBodyRequest{AFD, TargetLoc}, true);
+                            TypeCheckFunctionBodyAtLocRequest{AFD, TargetLoc},
+                            true);
 }
 
 bool swift::typeCheckTopLevelCodeDecl(TopLevelCodeDecl *TLCD) {