[IDE] Allow typeCheckASTNodeAtLoc to type check declarations

We need the ability to type check declarations when entering using `typeCheckASTNodeAtLoc` for solver-based cursor info.

This causes more context to be type checked for code completion and breaks some type completion after dot if the base type is not fully qualified because the type checker will mark `TypeRepr` as invalid. But since we don’t actually need to run the type checker for type member completion, we can just not call `typeCheckContextAt` for `TypeIdentifierWithDot` completions.

Author: ahoppen

lib/IDE/CodeCompletion.cpp

@@ -174,6 +174,25 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
 
   /// \returns true on success, false on failure.
   bool typecheckParsedType() {
+    // If the type appeared inside an extension, make sure that extension has
+    // been bound.
+    auto SF = CurDeclContext->getParentSourceFile();
+    auto visitTopLevelDecl = [&](Decl *D) {
+      if (auto ED = dyn_cast<ExtensionDecl>(D)) {
+        if (ED->getSourceRange().contains(ParsedTypeLoc.getLoc())) {
+          ED->computeExtendedNominal();
+        }
+      }
+    };
+    for (auto item : SF->getTopLevelItems()) {
+      if (auto D = item.dyn_cast<Decl *>()) {
+        visitTopLevelDecl(D);
+      }
+    }
+    for (auto *D : SF->getHoistedDecls()) {
+      visitTopLevelDecl(D);
+    }
+
     assert(ParsedTypeLoc.getTypeRepr() && "should have a TypeRepr");
     if (ParsedTypeLoc.wasValidated() && !ParsedTypeLoc.isError()) {
       return true;
@@ -1576,11 +1595,14 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
     return;
 
   undoSingleExpressionReturn(CurDeclContext);
-  typeCheckContextAt(
-      TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
-      ParsedExpr
-          ? ParsedExpr->getLoc()
-          : CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc());
+  if (Kind != CompletionKind::TypeIdentifierWithDot) {
+    // Type member completion does not need a type-checked AST.
+    typeCheckContextAt(
+        TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
+        ParsedExpr
+            ? ParsedExpr->getLoc()
+            : CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc());
+  }
 
   // Add keywords even if type checking fails completely.
   addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);

lib/Sema/TypeCheckStmt.cpp

@@ -1926,9 +1926,19 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
         if (!braceCharRange.contains(Loc))
           return Action::SkipChildren(S);
 
-        // Reset the node found in a parent context.
-        if (!brace->isImplicit())
-          FoundNode = nullptr;
+        // Reset the node found in a parent context if it's not part of this
+        // brace statement.
+        // We must not reset FoundNode if it's inside thei BraceStmt's source
+        // range because the found node could be inside a capture list, which is
+        // syntactically part of the brace stmt's range but won't be walked as
+        // a child of the brace stmt.
+        if (!brace->isImplicit() && FoundNode) {
+          auto foundNodeCharRange = Lexer::getCharSourceRangeFromSourceRange(
+              SM, FoundNode->getSourceRange());
+          if (!braceCharRange.contains(foundNodeCharRange)) {
+            FoundNode = nullptr;
+          }
+        }
 
         for (ASTNode &node : brace->getElements()) {
           if (SM.isBeforeInBuffer(Loc, node.getStartLoc()))
@@ -1985,6 +1995,18 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
     PreWalkAction walkToDeclPre(Decl *D) override {
       if (auto *newDC = dyn_cast<DeclContext>(D))
         DC = newDC;
+
+      if (!SM.isBeforeInBuffer(Loc, D->getStartLoc())) {
+        // NOTE: We need to check the character loc here because the target
+        // loc can be inside the last token of the node. i.e. interpolated
+        // string.
+        SourceLoc endLoc = Lexer::getLocForEndOfToken(SM, D->getEndLoc());
+        if (!(SM.isBeforeInBuffer(endLoc, Loc) || endLoc == Loc)) {
+          if (!isa<TopLevelCodeDecl>(D)) {
+            FoundNode = new ASTNode(D);
+          }
+        }
+      }
       return Action::Continue();
     }
 
@@ -1998,7 +2020,7 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
   }
 
   // Nothing found at the location, or the decl context does not own the 'Loc'.
-  if (finder.isNull())
+  if (finder.isNull() || !finder.getDeclContext())
     return true;
 
   DeclContext *DC = finder.getDeclContext();