[CodeCompletion] Teach TypeCheckASTNodeAtLocRequest to check initializers

This simplifies the logic in `ide::typeCheckContextAt` and fixes an issue that prevent code completion from working inside variables that are initialized by calling a closure.

Fixes rdar://90455154 [SR-16012]
Fixes rdar://85600167 [SR-15495]

Author: ahoppen

lib/AST/DeclContext.cpp

@@ -85,10 +85,11 @@ ProtocolDecl *DeclContext::getExtendedProtocolDecl() const {
 
 VarDecl *DeclContext::getNonLocalVarDecl() const {
   if (auto *init = dyn_cast<PatternBindingInitializer>(this)) {
-   if (auto *var =
-         init->getBinding()->getAnchoringVarDecl(init->getBindingIndex())) {
-      return var;
-     }
+    if (auto binding = init->getBinding()) {
+      if (auto *var = binding->getAnchoringVarDecl(init->getBindingIndex())) {
+        return var;
+      }
+    }
   }
   return nullptr;
 }

lib/IDE/ExprContextAnalysis.cpp

@@ -43,9 +43,6 @@ using namespace ide;
 //===----------------------------------------------------------------------===//
 
 void swift::ide::typeCheckContextAt(DeclContext *DC, SourceLoc Loc) {
-  while (isa<AbstractClosureExpr>(DC))
-    DC = DC->getParent();
-
   // Make sure the extension has been bound.
   {
     // Even if the extension is invalid (e.g. nested in a function or another
@@ -77,60 +74,7 @@ void swift::ide::typeCheckContextAt(DeclContext *DC, SourceLoc Loc) {
     }
   }
 
-  // Type-check this context.
-  switch (DC->getContextKind()) {
-  case DeclContextKind::AbstractClosureExpr:
-  case DeclContextKind::Module:
-  case DeclContextKind::FileUnit:
-  case DeclContextKind::SerializedLocal:
-  case DeclContextKind::EnumElementDecl:
-  case DeclContextKind::GenericTypeDecl:
-  case DeclContextKind::SubscriptDecl:
-  case DeclContextKind::ExtensionDecl:
-    // Nothing to do for these.
-    break;
-
-  case DeclContextKind::Initializer:
-    if (auto *patternInit = dyn_cast<PatternBindingInitializer>(DC)) {
-      if (auto *PBD = patternInit->getBinding()) {
-        auto i = patternInit->getBindingIndex();
-        PBD->getPattern(i)->forEachVariable(
-            [](VarDecl *VD) { (void)VD->getInterfaceType(); });
-        if (PBD->getInit(i)) {
-          if (!PBD->isInitializerChecked(i))
-            typeCheckPatternBinding(PBD, i,
-                                    /*LeaveClosureBodyUnchecked=*/true);
-        }
-      }
-    } else if (auto *defaultArg = dyn_cast<DefaultArgumentInitializer>(DC)) {
-      if (auto *AFD = dyn_cast<AbstractFunctionDecl>(defaultArg->getParent())) {
-        auto *Param = AFD->getParameters()->get(defaultArg->getIndex());
-        (void)Param->getTypeCheckedDefaultExpr();
-      }
-      if (auto *SD = dyn_cast<SubscriptDecl>(defaultArg->getParent())) {
-        auto *Param = SD->getIndices()->get(defaultArg->getIndex());
-        (void)Param->getTypeCheckedDefaultExpr();
-      }
-    }
-    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::typeCheckASTNodeAtLoc(DC, Loc);
-    } else {
-      assert(bodyRange.isInvalid() && "The body should not be parsed if the "
-                                      "completion happens in the signature");
-    }
-    break;
-  }
-  }
+  swift::typeCheckASTNodeAtLoc(DC, Loc);
 }
 
 //===----------------------------------------------------------------------===//

lib/Sema/TypeCheckStmt.cpp

@@ -1800,6 +1800,36 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(Evaluator &evaluator,
   assert(DiagnosticSuppression::isEnabled(ctx.Diags) &&
          "Diagnosing and Single ASTNode type checknig don't mix");
 
+  // Initializers aren't walked by ASTWalker and thus we don't find the context
+  // to type check using ASTNodeFinder. Also, initializers aren't representable
+  // by ASTNodes that can be type checked using typeCheckASTNode.
+  // Handle them specifically here.
+  if (auto *patternInit = dyn_cast<PatternBindingInitializer>(DC)) {
+    if (auto *PBD = patternInit->getBinding()) {
+      auto i = patternInit->getBindingIndex();
+      PBD->getPattern(i)->forEachVariable(
+          [](VarDecl *VD) { (void)VD->getInterfaceType(); });
+      if (PBD->getInit(i)) {
+        if (!PBD->isInitializerChecked(i)) {
+          typeCheckPatternBinding(PBD, i,
+                                  /*LeaveClosureBodyUnchecked=*/true);
+          return false;
+        }
+      }
+    }
+  } else if (auto *defaultArg = dyn_cast<DefaultArgumentInitializer>(DC)) {
+    if (auto *AFD = dyn_cast<AbstractFunctionDecl>(defaultArg->getParent())) {
+      auto *Param = AFD->getParameters()->get(defaultArg->getIndex());
+      (void)Param->getTypeCheckedDefaultExpr();
+      return false;
+    }
+    if (auto *SD = dyn_cast<SubscriptDecl>(defaultArg->getParent())) {
+      auto *Param = SD->getIndices()->get(defaultArg->getIndex());
+      (void)Param->getTypeCheckedDefaultExpr();
+      return false;
+    }
+  }
+
   // Find innermost ASTNode at Loc from DC. Results the reference to the found
   // ASTNode and the decl context of it.
   class ASTNodeFinder : public ASTWalker {

test/IDE/complete_call_arg.swift

@@ -1371,3 +1371,20 @@ func testDynamicMemberSubscriptLookup() {
 // DYNAMIC_MEMBER_SUBSCRIPT_LOOKUP-DAG: Decl[LocalVar]/Local/TypeRelation[Identical]: index[#Int#]; name=index
 // DYNAMIC_MEMBER_SUBSCRIPT_LOOKUP-DAG: Pattern/CurrNominal/Flair[ArgLabels]: ['[']{#keyPath: KeyPath<Binding<MyStruct>, Value>#}[']'][#Value#]; name=keyPath:
 // DYNAMIC_MEMBER_SUBSCRIPT_LOOKUP: End completions
+
+func testVarInitializedByCallingClosure() {
+  struct MyBundle {
+    func vrl(forResource: String, withExtension: String?)
+  }
+
+  struct Foo {
+    private lazy var calculatorContext: Void = {
+      let Bundle_main = MyBundle()
+      Bundle_main.vrl(forResource: "turnips", #^VAR_INITIALIZED_BY_CALLING_CLOSURE^#withExtension: "js")
+    }()
+  }
+
+// VAR_INITIALIZED_BY_CALLING_CLOSURE: Begin completions, 1 items
+// VAR_INITIALIZED_BY_CALLING_CLOSURE-DAG: Pattern/Local/Flair[ArgLabels]:     {#withExtension: String?#}[#String?#];
+// VAR_INITIALIZED_BY_CALLING_CLOSURE: End completions
+}

test/IDE/complete_multiple_files.swift

@@ -9,6 +9,10 @@
 //
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=MODULE_SCOPED %S/Inputs/multiple-files-1.swift %S/Inputs/multiple-files-2.swift | %FileCheck %s -check-prefix=MODULE_SCOPED
 
+// RUN: %empty-directory(%t)
+// RUN: echo "" > %t/empty.swift
+// RUN: %swift-ide-test --code-completion --code-completion-token VAR_INITIALIZED_BY_CALLING_CLOSURE --source-filename %s --second-source-filename %t/empty.swift | %FileCheck %s -check-prefix=VAR_INITIALIZED_BY_CALLING_CLOSURE
+
 func testObjectExpr() {
   fooObject.#^T1^#
 }
@@ -45,3 +49,16 @@ func moduleScoped() {
 // MODULE_SCOPED: Decl[Struct]/CurrModule: FooStruct[#FooStruct#]{{; name=.+$}}
 // MODULE_SCOPED-NOT: ERROR
 // MODULE_SCOPED: End completions
+
+enum Foo {
+    case bar
+}
+
+var sr15495: Void = {
+    let foo: Foo = .#^VAR_INITIALIZED_BY_CALLING_CLOSURE^#
+}()
+
+// VAR_INITIALIZED_BY_CALLING_CLOSURE:     Begin completions, 2 items
+// VAR_INITIALIZED_BY_CALLING_CLOSURE-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Identical]: bar[#Foo#];
+// VAR_INITIALIZED_BY_CALLING_CLOSURE-DAG: Decl[InstanceMethod]/CurrNominal/TypeRelation[Invalid]: hash({#(self): Foo#})[#(into: inout Hasher) -> Void#];
+// VAR_INITIALIZED_BY_CALLING_CLOSURE:     End completions