Merge pull request #32696 from rintaro/ide-completion-typecheckatloc

[CodeCompletion] Only type check target statement and related decls in function body

Author: rintaro

include/swift/AST/Expr.h

@@ -3759,7 +3759,25 @@ class SerializedAbstractClosureExpr : public SerializedLocalDeclContext {
 ///     { [weak c] (a : Int) -> Int in a + c!.getFoo() }
 /// \endcode
 class ClosureExpr : public AbstractClosureExpr {
+public:
+  enum class BodyState {
+    /// The body was parsed, but not ready for type checking because
+    /// the closure parameters haven't been type checked.
+    Parsed,
+
+    /// The type of the closure itself was type checked. But the body has not
+    /// been type checked yet.
+    ReadyForTypeChecking,
+
+    /// The body was typechecked with the enclosing closure.
+    /// i.e. single expression closure or function builder closure.
+    TypeCheckedWithSignature,
+
+    /// The body was type checked separately from the enclosing closure.
+    SeparatelyTypeChecked,
+  };
 
+private:
   /// The range of the brackets of the capture list, if present.
   SourceRange BracketRange;
 
@@ -3783,8 +3801,7 @@ class ClosureExpr : public AbstractClosureExpr {
   SourceLoc InLoc;
 
   /// The explicitly-specified result type.
-  llvm::PointerIntPair<TypeExpr *, 1, bool>
-    ExplicitResultTypeAndSeparatelyChecked;
+  llvm::PointerIntPair<TypeExpr *, 2, BodyState> ExplicitResultTypeAndBodyState;
 
   /// The body of the closure, along with a bit indicating whether it
   /// was originally just a single expression.
@@ -3799,7 +3816,7 @@ class ClosureExpr : public AbstractClosureExpr {
       BracketRange(bracketRange),
       CapturedSelfDecl(capturedSelfDecl),
       ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
-      ExplicitResultTypeAndSeparatelyChecked(explicitResultType, false),
+      ExplicitResultTypeAndBodyState(explicitResultType, BodyState::Parsed),
       Body(nullptr) {
     setParameterList(params);
     Bits.ClosureExpr.HasAnonymousClosureVars = false;
@@ -3854,15 +3871,13 @@ class ClosureExpr : public AbstractClosureExpr {
 
   Type getExplicitResultType() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultTypeAndSeparatelyChecked.getPointer()
-        ->getInstanceType();
+    return ExplicitResultTypeAndBodyState.getPointer()->getInstanceType();
   }
   void setExplicitResultType(Type ty);
 
   TypeRepr *getExplicitResultTypeRepr() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultTypeAndSeparatelyChecked.getPointer()
-        ->getTypeRepr();
+    return ExplicitResultTypeAndBodyState.getPointer()->getTypeRepr();
   }
 
   /// Determine whether the closure has a single expression for its
@@ -3904,14 +3919,20 @@ class ClosureExpr : public AbstractClosureExpr {
   /// captured non-weakly).
   bool capturesSelfEnablingImplictSelf() const;
 
-  /// Whether this closure's body was type checked separately from its
-  /// enclosing expression.
-  bool wasSeparatelyTypeChecked() const {
-    return ExplicitResultTypeAndSeparatelyChecked.getInt();
+
+  /// Get the type checking state of this closure's body.
+  BodyState getBodyState() const {
+    return ExplicitResultTypeAndBodyState.getInt();
+  }
+  void setBodyState(BodyState v) {
+    ExplicitResultTypeAndBodyState.setInt(v);
   }
 
-  void setSeparatelyTypeChecked(bool flag = true) {
-    ExplicitResultTypeAndSeparatelyChecked.setInt(flag);
+  /// Whether this closure's body is/was type checked separately from its
+  /// enclosing expression.
+  bool isSeparatelyTypeChecked() const {
+    return getBodyState() == BodyState::SeparatelyTypeChecked ||
+           getBodyState() == BodyState::ReadyForTypeChecking;
   }
 
   static bool classof(const Expr *E) {

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

@@ -514,7 +514,7 @@ namespace swift {
     /// Indicate that the type checker should skip type-checking non-inlinable
     /// function bodies.
     bool SkipNonInlinableFunctionBodies = 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/ASTWalker.cpp

@@ -833,7 +833,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
     // If the closure was separately type checked and we don't want to
     // visit separately-checked closure bodies, bail out now.
-    if (expr->wasSeparatelyTypeChecked() &&
+    if (expr->isSeparatelyTypeChecked() &&
         !Walker.shouldWalkIntoSeparatelyCheckedClosure(expr))
       return expr;
 

lib/AST/Expr.cpp

@@ -1996,7 +1996,7 @@ bool ClosureExpr::capturesSelfEnablingImplictSelf() const {
 
 void ClosureExpr::setExplicitResultType(Type ty) {
   assert(ty && !ty->hasTypeVariable());
-  ExplicitResultTypeAndSeparatelyChecked.getPointer()
+  ExplicitResultTypeAndBodyState.getPointer()
       ->setType(MetatypeType::get(ty));
 }
 

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

@@ -4129,6 +4129,10 @@ namespace {
       // If we end up here, we should have diagnosed somewhere else
       // already.
       Expr *simplified = simplifyExprType(expr);
+      // Invalidate 'VarDecl's inside the pattern.
+      expr->getSubPattern()->forEachVariable([](VarDecl *VD) {
+        VD->setInvalid();
+      });
       if (!SuppressDiagnostics
           && !cs.getType(simplified)->is<UnresolvedType>()) {
         auto &de = cs.getASTContext().Diags;
@@ -7834,8 +7838,11 @@ namespace {
         return true;
 
       case SolutionApplicationToFunctionResult::Delay: {
-        auto closure = cast<ClosureExpr>(fn.getAbstractClosureExpr());
-        ClosuresToTypeCheck.push_back(closure);
+        if (!Rewriter.cs.Options
+                .contains(ConstraintSystemFlags::LeaveClosureBodyUnchecked)) {
+          auto closure = cast<ClosureExpr>(fn.getAbstractClosureExpr());
+          ClosuresToTypeCheck.push_back(closure);
+        }
         return false;
       }
       }
@@ -8414,6 +8421,7 @@ Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
 
   // Visit closures that have non-single expression bodies.
   bool hadError = false;
+
   for (auto *closure : walker.getClosuresToTypeCheck())
     hadError |= TypeChecker::typeCheckClosureBody(closure);
 

lib/Sema/CSClosure.cpp

@@ -345,11 +345,13 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
     ClosureConstraintApplication application(
         solution, closure, closureFnType->getResult(), rewriteTarget);
     application.visit(fn.getBody());
+    closure->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
 
     return SolutionApplicationToFunctionResult::Success;
   }
 
   // Otherwise, we need to delay type checking of the closure until later.
   solution.setExprTypes(closure);
+  closure->setBodyState(ClosureExpr::BodyState::ReadyForTypeChecking);
   return SolutionApplicationToFunctionResult::Delay;
 }

lib/Sema/ConstraintSystem.h

@@ -1250,6 +1250,10 @@ enum class ConstraintSystemFlags {
   /// for a pre-configured set of expressions on line numbers by setting
   /// \c DebugConstraintSolverOnLines.
   DebugConstraints = 0x10,
+
+  /// Don't try to type check closure bodies, and leave them unchecked. This is
+  /// used for source tooling functionalities.
+  LeaveClosureBodyUnchecked = 0x20,
 };
 
 /// Options that affect the constraint system as a whole.

lib/Sema/MiscDiagnostics.cpp

@@ -86,7 +86,7 @@ static void diagSyntacticUseRestrictions(const Expr *E, const DeclContext *DC,
 
     bool walkToDeclPre(Decl *D) override {
       if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
-        return !closure->wasSeparatelyTypeChecked();
+        return !closure->isSeparatelyTypeChecked();
       return false;
     }
 
@@ -1501,7 +1501,7 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
     // Don't walk into nested decls.
     bool walkToDeclPre(Decl *D) override {
       if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
-        return !closure->wasSeparatelyTypeChecked();
+        return !closure->isSeparatelyTypeChecked();
       return false;
     }
 

lib/Sema/TypeCheckAvailability.cpp

@@ -1136,7 +1136,7 @@ static void findAvailabilityFixItNodes(SourceRange ReferenceRange,
         if (Expr *ParentExpr = Parent.getAsExpr()) {
           auto *ParentClosure = dyn_cast<ClosureExpr>(ParentExpr);
           if (!ParentClosure ||
-              ParentClosure->wasSeparatelyTypeChecked()) {
+              ParentClosure->isSeparatelyTypeChecked()) {
             return false;
           }
         } else if (auto *ParentStmt = Parent.getAsStmt()) {