[Completion] Type-check parent closures for local functions

Local functions can capture variables from parent
closures, so we need to make sure we type-check
parent closures when doing completion in a local
function. Ideally we ought to be able to be more
selective about the elements of the parent closure
that we type-check, but that's a more complex change
I'm leaving as future work for now.

Author: hamishknight

include/swift/Sema/CompletionContextFinder.h

@@ -52,6 +52,18 @@ class CompletionContextFinder : public ASTWalker {
   Expr *InitialExpr = nullptr;
   DeclContext *InitialDC;
 
+  /// Whether we're looking for any viable fallback expression.
+  bool ForFallback = false;
+
+  /// Finder for fallback completion contexts within the outermost non-closure
+  /// context of the code completion expression's direct context.
+  CompletionContextFinder(DeclContext *completionDC)
+    : InitialDC(completionDC), ForFallback(true) {
+    while (auto *ACE = dyn_cast<AbstractClosureExpr>(InitialDC))
+      InitialDC = ACE->getParent();
+    InitialDC->walkContext(*this);
+  }
+
 public:
   MacroWalking getMacroWalkingBehavior() const override {
     return MacroWalking::Arguments;
@@ -61,18 +73,16 @@ class CompletionContextFinder : public ASTWalker {
   CompletionContextFinder(constraints::SyntacticElementTarget target,
                           DeclContext *DC);
 
-  /// Finder for completion contexts within the outermost non-closure context of
-  /// the code completion expression's direct context.
-  CompletionContextFinder(DeclContext *completionDC) : InitialDC(completionDC) {
-    while (auto *ACE = dyn_cast<AbstractClosureExpr>(InitialDC))
-      InitialDC = ACE->getParent();
-    InitialDC->walkContext(*this);
+  static CompletionContextFinder forFallback(DeclContext *DC) {
+    return CompletionContextFinder(DC);
   }
 
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override;
 
   PostWalkResult<Expr *> walkToExprPost(Expr *E) override;
 
+  PreWalkAction walkToDeclPre(Decl *D) override;
+
   bool locatedInStringInterpolation() const {
     return hasContext(ContextKind::StringInterpolation);
   }

lib/IDE/PostfixCompletion.cpp

@@ -74,7 +74,7 @@ void PostfixCompletionCallback::fallbackTypeCheck(DeclContext *DC) {
   Expr *fallbackExpr = CompletionExpr;
   DeclContext *fallbackDC = DC;
 
-  CompletionContextFinder finder(DC);
+  auto finder = CompletionContextFinder::forFallback(DC);
   if (finder.hasCompletionExpr()) {
     if (auto fallback = finder.getFallbackCompletionExpr()) {
       fallbackExpr = fallback->E;

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -24,7 +24,7 @@ using namespace swift::constraints;
 void TypeCheckCompletionCallback::fallbackTypeCheck(DeclContext *DC) {
   assert(!GotCallback);
 
-  CompletionContextFinder finder(DC);
+  auto finder = CompletionContextFinder::forFallback(DC);
   if (!finder.hasCompletionExpr())
     return;
 

lib/Sema/CompletionContextFinder.cpp

@@ -86,6 +86,17 @@ CompletionContextFinder::walkToExprPost(Expr *E) {
   return Action::Continue(E);
 }
 
+ASTWalker::PreWalkAction CompletionContextFinder::walkToDeclPre(Decl *D) {
+  // Look through any decl if we're looking for any viable fallback expression.
+  if (ForFallback)
+    return Action::Continue();
+
+  // Otherwise, follow the same rule as the ConstraintSystem, where only
+  // nested PatternBindingDecls are solved as part of the system. Local decls
+  // are handled by TypeCheckASTNodeAtLocRequest.
+  return Action::VisitNodeIf(isa<PatternBindingDecl>(D));
+}
+
 size_t CompletionContextFinder::getKeyPathCompletionComponentIndex() const {
   size_t ComponentIndex = 0;
   auto Components = getKeyPathContainingCompletionComponent()->getComponents();

lib/Sema/TypeCheckStmt.cpp

@@ -2628,15 +2628,26 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
     }
   }
 
-  // If the context is a closure, type check the entire surrounding closure.
-  // Conjunction constraints ensure that statements unrelated to the one that
-  // contains the code completion token are not type checked.
-  if (auto CE = dyn_cast<ClosureExpr>(DC)) {
+  // If we're within a ClosureExpr that can propagate its captures to this
+  // DeclContext, we need to type-check the entire surrounding closure. If the
+  // completion token is contained within the closure itself, conjunction
+  // constraints ensure that statements unrelated to the one that contains the
+  // code completion token are not type checked. If it's in a nested local
+  // function, we unfortunately need to type-check everything since we need to
+  // apply the solution.
+  // FIXME: We ought to see if we can do better in that case.
+  if (auto *CE = DC->getInnermostClosureForCaptures()) {
     if (CE->getBodyState() == ClosureExpr::BodyState::Parsed) {
       swift::typeCheckASTNodeAtLoc(
           TypeCheckASTNodeAtLocContext::declContext(CE->getParent()),
           CE->getLoc());
-      return false;
+
+      // If the context itself is a ClosureExpr, we should have type-checked
+      // the completion expression now. If it's a nested local declaration,
+      // fall through to type-check the AST node now that we've type-checked
+      // the surrounding closure.
+      if (isa<ClosureExpr>(DC))
+        return false;
     }
   }
 

test/IDE/complete_if_switch_expr.swift

@@ -252,15 +252,22 @@ func testSkipTypeChecking9() -> E {
 }
 
 func testSkipTypeChecking10() -> E {
-  // We only need to type-check the inner-most function for this.
+  // Similar to the above case, we need to type-check everything for this since
+  // the type-checking of 'takesArgAndClosure' is required to correctly handle
+  // any potential captures in 'foo'.
   if Bool.random() {
-    NO.TYPECHECK
+    .e
   } else {
-    takesArgAndClosure(NO.TYPECHECK) {
+    takesArgAndClosure(0) {
       func foo() {
+        // We can however skip unrelated elements in the local function.
+        let x = NO.TYPECHECK
+        if NO.TYPECHECK {
+          takesE(NO.TYPECHECK)
+        }
         takesE(.#^DOT21?check=DOT^#)
       }
-      return NO.TYPECHECK
+      return .e
     }
   }
 }
@@ -320,7 +327,7 @@ func testSkipTypeChecking14() -> E {
   }
 }
 
-func testSkipTypeChecking14() -> E {
+func testSkipTypeChecking15() -> E {
   switch Bool.random() {
   case true:
     .#^DOT26?check=DOT^#
@@ -336,7 +343,7 @@ func testSkipTypeChecking14() -> E {
   }
 }
 
-func testSkipTypechecking15(_ x: inout Int) -> E {
+func testSkipTypechecking16(_ x: inout Int) -> E {
   switch Bool.random() {
   case true:
     .#^DOT27?check=DOT^#

test/IDE/complete_issue-77305.swift

@@ -0,0 +1,31 @@
+// RUN: %batch-code-completion
+
+// https://github.com/swiftlang/swift/issues/77305
+
+struct S {
+  var x: Int
+}
+
+func withFoo(_ x: (S) -> Void) {}
+
+withFoo { foo in
+  func bar() {
+    foo.#^FN_IN_CLOSURE^#
+    // FN_IN_CLOSURE: Decl[InstanceVar]/CurrNominal: x[#Int#]; name=x
+  }
+}
+
+withFoo { x in
+  _ = { y in
+    func bar() {
+      _ = { z in
+        func baz() {
+          func qux() {
+            z.#^VERY_NESTED_FN_IN_CLOSURE^#
+            // VERY_NESTED_FN_IN_CLOSURE: Decl[InstanceVar]/CurrNominal: x[#Int#]; name=x
+          }
+        }
+      }(y)
+    }
+  }(x)
+}