[CodeCompletion] Skip typechecking unrelated statements in func bodies

rdar://problem/58687608

Author: rintaro

lib/IDE/CodeCompletion.cpp

@@ -5736,7 +5736,9 @@ void CodeCompletionCallbacksImpl::doneParsing() {
 
   typeCheckContextUntil(
       CurDeclContext,
-      CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc());
+      ParsedExpr
+          ? ParsedExpr->getLoc()
+          : CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc());
 
   // Add keywords even if type checking fails completely.
   addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);

lib/IDE/ConformingMethodList.cpp

@@ -67,9 +67,7 @@ void ConformingMethodListCallbacks::doneParsing() {
   if (!ParsedExpr)
     return;
 
-  typeCheckContextUntil(
-      CurDeclContext,
-      CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc());
+  typeCheckContextUntil(CurDeclContext, ParsedExpr->getLoc());
 
   Type T = ParsedExpr->getType();
 

lib/IDE/TypeContextInfo.cpp

@@ -88,9 +88,7 @@ void ContextInfoCallbacks::doneParsing() {
   if (!ParsedExpr)
     return;
 
-  typeCheckContextUntil(
-      CurDeclContext,
-      CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc());
+  typeCheckContextUntil(CurDeclContext, ParsedExpr->getLoc());
 
   ExprContextInfo Info(CurDeclContext, ParsedExpr);
 

lib/Sema/TypeCheckStmt.cpp

@@ -1563,12 +1563,19 @@ Stmt *StmtChecker::visitBraceStmt(BraceStmt *BS) {
   }
 
   for (auto &elem : BS->getElements()) {
-    if (auto *SubExpr = elem.dyn_cast<Expr*>()) {
-      SourceLoc Loc = SubExpr->getStartLoc();
-      if (EndTypeCheckLoc.isValid() &&
-          (Loc == EndTypeCheckLoc || SM.isBeforeInBuffer(EndTypeCheckLoc, Loc)))
+    if (EndTypeCheckLoc.isValid()) {
+      if (SM.isBeforeInBuffer(EndTypeCheckLoc, elem.getStartLoc()))
         break;
 
+      // NOTE: We need to check the character loc here because the target loc
+      // can be inside the last token of the node. i.e. string interpolation.
+      SourceLoc endLoc = Lexer::getLocForEndOfToken(SM, elem.getEndLoc());
+      if (endLoc == EndTypeCheckLoc ||
+          SM.isBeforeInBuffer(endLoc, EndTypeCheckLoc))
+        continue;
+    }
+
+    if (auto *SubExpr = elem.dyn_cast<Expr*>()) {
       // Type check the expression.
       TypeCheckExprOptions options = TypeCheckExprFlags::IsExprStmt;
       bool isDiscarded = (!getASTContext().LangOpts.Playground &&
@@ -1607,22 +1614,12 @@ Stmt *StmtChecker::visitBraceStmt(BraceStmt *BS) {
     }
 
     if (auto *SubStmt = elem.dyn_cast<Stmt*>()) {
-      SourceLoc Loc = SubStmt->getStartLoc();
-      if (EndTypeCheckLoc.isValid() &&
-          (Loc == EndTypeCheckLoc || SM.isBeforeInBuffer(EndTypeCheckLoc, Loc)))
-        break;
-
       typeCheckStmt(SubStmt);
       elem = SubStmt;
       continue;
     }
 
     Decl *SubDecl = elem.get<Decl *>();
-    SourceLoc Loc = SubDecl->getStartLoc();
-    if (EndTypeCheckLoc.isValid() &&
-        (Loc == EndTypeCheckLoc || SM.isBeforeInBuffer(EndTypeCheckLoc, Loc)))
-      break;
-
     TypeChecker::typeCheckDecl(SubDecl);
   }
 

test/IDE/complete_skipbody.swift

@@ -0,0 +1,26 @@
+// RUN: %swift-ide-test -code-completion -source-filename %s -code-completion-token COMPLETE -debug-forbid-typecheck-prefix FORBIDDEN | %FileCheck %s
+struct FORBIDDEN_Struct {
+  func FORBIDDEN_method() -> Int? { 1 }
+}
+
+struct MyStruct {
+  var x: Int { 1 }
+  var y: Int { 1 }
+}
+
+func test(value: MyStruct) {
+
+  let FORBIDDEN_localVar = 1
+  let unrelated = FORBIDDEN_Struct()
+  guard let a = unrelated.FORBIDDEN_method() else {
+    return
+  }
+
+  _ = value.#^COMPLETE^#
+}
+
+// CHECK: Begin completions, 3 items
+// CHECK-DAG: Keyword[self]/CurrNominal:          self[#MyStruct#]; name=self
+// CHECK-DAG: Decl[InstanceVar]/CurrNominal:      x[#Int#]; name=x
+// CHECK-DAG: Decl[InstanceVar]/CurrNominal:      y[#Int#]; name=y
+// CHECK: End completions