[CodeCompletion] Only type check related statements in function body

Introduce 'TypeCheckSingleASTNode' mode that only type checks single body
element and dependent necessities (i.e. referencing ValueDecls and their
dependencies).

Renamed swift::typeCheckAbstractFunctionBodyAtLoc() to
swift::typeCheckASTNodeAtLoc(DeclContext *, SourceLoc). That type checks
innermost 'ASTNode' at the location. Also, 'TypeCheckSingleASTNode' mode
skips type checking any "body" of the node (i.e. BraceStmt elements for
function body, if statement body, closure body, etc.)

Added on-demand type checking using it:
 - VarDecl in TapExpr
 - ParamDecl in ClosureExpr
 - Return type of ClosureExpr
 - Binding value in control statements
   (e.g. ForEachStmt, SwitchStmt, DoCatchStmt, etc.)

rdar://problem/63932852

Author: rintaro

include/swift/AST/Expr.h

@@ -950,11 +950,9 @@ class TapExpr : public Expr {
   BraceStmt * getBody() const { return Body; }
   void setBody(BraceStmt * b) { Body = b; }
 
-  SourceLoc getLoc() const { return SubExpr ? SubExpr->getLoc() : SourceLoc(); }
+  SourceLoc getLoc() const { return getStartLoc(); }
 
-  SourceLoc getStartLoc() const {
-    return SubExpr ? SubExpr->getStartLoc() : SourceLoc();
-  }
+  SourceLoc getStartLoc() const;
 
   SourceLoc getEndLoc() const;
 

include/swift/AST/TypeCheckRequests.h

@@ -892,9 +892,9 @@ class TypeCheckFunctionBodyRequest :
 /// 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),
+class TypeCheckASTNodeAtLocRequest
+    : public SimpleRequest<TypeCheckASTNodeAtLocRequest,
+                           bool(DeclContext *, SourceLoc),
                            RequestFlags::Uncached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -903,8 +903,7 @@ class TypeCheckFunctionBodyAtLocRequest
   friend SimpleRequest;
 
   // Evaluation.
-  bool evaluate(Evaluator &evaluator, AbstractFunctionDecl *func,
-                SourceLoc Loc) const;
+  bool evaluate(Evaluator &evaluator, DeclContext *DC, SourceLoc Loc) const;
 };
 
 /// Request to obtain a list of stored properties in a nominal type.

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -207,8 +207,8 @@ SWIFT_REQUEST(TypeChecker, TangentStoredPropertyRequest,
               llvm::Expected<VarDecl *>(VarDecl *, CanType), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyRequest,
               bool(AbstractFunctionDecl *), Cached, NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyAtLocRequest,
-              bool(AbstractFunctionDecl *, SourceLoc),
+SWIFT_REQUEST(TypeChecker, TypeCheckASTNodeAtLocRequest,
+              bool(DeclContext *, SourceLoc),
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, UnderlyingTypeRequest, Type(TypeAliasDecl *),
               SeparatelyCached, NoLocationInfo)

include/swift/Basic/LangOptions.h

@@ -507,7 +507,10 @@ namespace swift {
     /// Indicate that the type checker should skip type-checking non-inlinable
     /// function bodies.
     bool SkipNonInlinableFunctionBodies = false;
-    
+
+    /// Single statement (ASTNode) type checking mode for IDE tooling.
+    bool TypeCheckSingleASTNode = false;
+
     ///
     /// Flags for developers
     ///

include/swift/Sema/IDETypeChecking.h

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

lib/AST/Expr.cpp

@@ -2398,6 +2398,15 @@ VarDecl * TapExpr::getVar() const {
   return dyn_cast<VarDecl>(Body->getFirstElement().dyn_cast<Decl *>());
 }
 
+SourceLoc TapExpr::getStartLoc() const {
+  // Include the body in the range, assuming the body follows the SubExpr.
+  if (auto *const se = getSubExpr())
+    return se->getStartLoc();
+  if (auto *const b = getBody())
+    return b->getStartLoc();
+  return SourceLoc();
+}
+
 SourceLoc TapExpr::getEndLoc() const {
   // Include the body in the range, assuming the body follows the SubExpr.
   // Also, be (perhaps overly) defensive about null pointers & invalid

lib/IDE/CompletionInstance.cpp

@@ -520,6 +520,9 @@ bool CompletionInstance::performNewOperation(
 
     Invocation.setCodeCompletionPoint(completionBuffer, Offset);
 
+    // Enable single ASTNode type checking mode.
+    Invocation.getTypeCheckerOptions().TypeCheckSingleASTNode = true;
+
     if (CI.setup(Invocation)) {
       Error = "failed to setup compiler instance";
       return false;

lib/IDE/ExprContextAnalysis.cpp

@@ -48,17 +48,30 @@ void typeCheckContextImpl(DeclContext *DC, SourceLoc Loc) {
   if (DC->isModuleScopeContext())
     return;
 
-  typeCheckContextImpl(DC->getParent(), Loc);
+  // Make sure the extension has been bound, in case it is in an inactive #if
+  // or something weird like that.
+  {
+    SmallVector<ExtensionDecl *, 1> extensions;
+    for (auto typeCtx = DC->getInnermostTypeContext(); typeCtx != nullptr;
+         typeCtx = typeCtx->getParent()->getInnermostTypeContext()) {
+      if (auto *ext = dyn_cast<ExtensionDecl>(typeCtx))
+        extensions.push_back(ext);
+    }
+    while (!extensions.empty()) {
+      extensions.back()->computeExtendedNominal();
+      extensions.pop_back();
+    }
+  }
 
   // Type-check this context.
   switch (DC->getContextKind()) {
   case DeclContextKind::AbstractClosureExpr:
   case DeclContextKind::Module:
   case DeclContextKind::SerializedLocal:
-  case DeclContextKind::TopLevelCodeDecl:
   case DeclContextKind::EnumElementDecl:
   case DeclContextKind::GenericTypeDecl:
   case DeclContextKind::SubscriptDecl:
+  case DeclContextKind::ExtensionDecl:
     // Nothing to do for these.
     break;
 
@@ -76,25 +89,23 @@ void typeCheckContextImpl(DeclContext *DC, SourceLoc Loc) {
     }
     break;
 
+  case DeclContextKind::TopLevelCodeDecl:
+    swift::typeCheckASTNodeAtLoc(DC, Loc);
+    break;
+
   case DeclContextKind::AbstractFunctionDecl: {
     auto *AFD = cast<AbstractFunctionDecl>(DC);
     auto &SM = DC->getASTContext().SourceMgr;
     auto bodyRange = AFD->getBodySourceRange();
     if (SM.rangeContainsTokenLoc(bodyRange, Loc)) {
-      swift::typeCheckAbstractFunctionBodyAtLoc(AFD, Loc);
+      swift::typeCheckASTNodeAtLoc(DC, Loc);
     } else {
       assert(bodyRange.isInvalid() && "The body should not be parsed if the "
                                       "completion happens in the signature");
     }
     break;
   }
 
-  case DeclContextKind::ExtensionDecl:
-    // Make sure the extension has been bound, in case it is in an
-    // inactive #if or something weird like that.
-    cast<ExtensionDecl>(DC)->computeExtendedNominal();
-    break;
-
   case DeclContextKind::FileUnit:
     llvm_unreachable("module scope context handled above");
   }
@@ -105,11 +116,7 @@ void swift::ide::typeCheckContextAt(DeclContext *DC, SourceLoc Loc) {
   while (isa<AbstractClosureExpr>(DC))
     DC = DC->getParent();
 
-  if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(DC)) {
-    typeCheckTopLevelCodeDecl(TLCD);
-  } else {
-    typeCheckContextImpl(DC, Loc);
-  }
+  typeCheckContextImpl(DC, Loc);
 }
 
 //===----------------------------------------------------------------------===//
@@ -300,6 +307,10 @@ void swift::ide::collectPossibleReturnTypesFromContext(
   }
 
   if (auto ACE = dyn_cast<AbstractClosureExpr>(DC)) {
+    // Try type checking the closure signature if it hasn't.
+    if (!ACE->getType())
+      swift::typeCheckASTNodeAtLoc(ACE->getParent(), ACE->getLoc());
+
     // Use the type checked type if it has.
     if (ACE->getType() && !ACE->getType()->hasError() &&
         !ACE->getResultType()->hasUnresolvedType()) {

lib/Sema/CSApply.cpp

@@ -8271,22 +8271,25 @@ 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);
-  }
+  if (!getASTContext().TypeCheckerOpts.TypeCheckSingleASTNode) {
+    bool hadError = false;
 
-  // If any of them failed to type check, bail.
-  if (hadError)
-    return None;
+    // Visit closures that have non-single expression bodies.
+    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);
+    }
+
+    // If any of them failed to type check, bail.
+    if (hadError)
+      return None;
+  }
 
   rewriter.finalize();
 

lib/Sema/TypeCheckDecl.cpp

@@ -45,6 +45,7 @@
 #include "swift/Basic/Statistic.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/Parse/Parser.h"
+#include "swift/Sema/IDETypeChecking.h"
 #include "swift/Serialization/SerializedModuleLoader.h"
 #include "swift/Strings.h"
 #include "swift/AST/NameLookupRequests.h"
@@ -2233,16 +2234,35 @@ InterfaceTypeRequest::evaluate(Evaluator &eval, ValueDecl *D) const {
       }
     }
 
-    if (!PD->getTypeRepr())
-      return Type();
+    if (PD->getTypeRepr()) {
+      return validateParameterType(PD);
+    }
+
+    // Try type checking the closure if it hasn't.
+    // NOTE: Do this only in 'TypeCheckSingleASTNode' mode. Otherwise, the
+    // parent closure should have been type checked before.
+    if (Context.TypeCheckerOpts.TypeCheckSingleASTNode &&
+        isa<ClosureExpr>(PD->getDeclContext())) {
+      auto *closure = cast<ClosureExpr>(PD->getDeclContext());
+      if (!closure->getType()) {
+        swift::typeCheckASTNodeAtLoc(closure->getParent(), closure->getLoc());
+        if (PD->hasInterfaceType())
+          return PD->getInterfaceType();
+      }
+    }
 
-    return validateParameterType(PD);
+    // Failed.
+    return Type();
   }
 
   case DeclKind::Var: {
     auto *VD = cast<VarDecl>(D);
     Type interfaceType;
     if (auto *parentE = VD->getParentExpr()) {
+      if (Context.TypeCheckerOpts.TypeCheckSingleASTNode &&
+          !parentE->getType()) {
+        swift::typeCheckASTNodeAtLoc(VD->getDeclContext(), parentE->getLoc());
+      }
       interfaceType = parentE->getType();
     } else if (auto *namingPattern = VD->getNamingPattern()) {
       interfaceType = namingPattern->getType();
@@ -2405,13 +2425,17 @@ NamingPatternRequest::evaluate(Evaluator &evaluator, VarDecl *VD) const {
     namingPattern = canVD->NamingPattern;
   }
 
-  if (!namingPattern) {
-    // Try type checking parent conditional statement.
+  if (!namingPattern &&
+      VD->getASTContext().TypeCheckerOpts.TypeCheckSingleASTNode) {
+    // Try type checking parent control statement.
+    // NOTE: Do this only in 'TypeCheckSingleASTNode' mode. Otherwise, the
+    // parent control statement should have been type checked before.
     if (auto parentStmt = VD->getParentPatternStmt()) {
-      if (auto LCS = dyn_cast<LabeledConditionalStmt>(parentStmt)) {
-        TypeChecker::typeCheckConditionForStatement(LCS, VD->getDeclContext());
-        namingPattern = VD->NamingPattern;
-      }
+      if (auto CS = dyn_cast<CaseStmt>(parentStmt))
+        parentStmt = CS->getParentStmt();
+      ASTNode node(parentStmt);
+      TypeChecker::typeCheckASTNode(node, VD->getDeclContext());
+      namingPattern = VD->getCanonicalVarDecl()->NamingPattern;
     }
   }