[IDE] Allow typeCheckASTNodeAtLoc to type check declarations

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();