Merge pull request #77537 from hamishknight/complete-func

[Completion] Type-check parent closures for local functions

include/swift/AST/DeclContext.h

@@ -515,12 +515,12 @@ public:
         const_cast<DeclContext *>(this)->getInnermostSkippedFunctionContext();
   }
 
-  /// Returns the innermost context that is a ClosureExpr, which defines how
-  /// self behaves, unless within a type context that redefines self.
+  /// Returns the innermost ClosureExpr context that can propagate its captures
+  /// to this DeclContext.
   LLVM_READONLY
-  ClosureExpr *getInnermostClosureForSelfCapture();
-  const ClosureExpr *getInnermostClosureForSelfCapture() const {
-    return const_cast<DeclContext *>(this)->getInnermostClosureForSelfCapture();
+  ClosureExpr *getInnermostClosureForCaptures();
+  const ClosureExpr *getInnermostClosureForCaptures() const {
+    return const_cast<DeclContext *>(this)->getInnermostClosureForCaptures();
   }
 
   /// Returns the semantic parent of this context.  A context has a

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 @@ public:
   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/AST/DeclContext.cpp

@@ -280,21 +280,21 @@ DeclContext *DeclContext::getInnermostSkippedFunctionContext() {
   return nullptr;
 }
 
-ClosureExpr *DeclContext::getInnermostClosureForSelfCapture() {
-  auto dc = this;
-  if (auto closure = dyn_cast<ClosureExpr>(dc)) {
-    return closure;
-  }
+ClosureExpr *DeclContext::getInnermostClosureForCaptures() {
+  auto *DC = this;
+  do {
+    if (auto *CE = dyn_cast<ClosureExpr>(DC))
+      return CE;
 
-  // Stop searching if we find a type decl, since types always
-  // redefine what 'self' means, even when nested inside a closure.
-  if (dc->isTypeContext()) {
+    // Autoclosures and AbstractFunctionDecls can propagate captures.
+    switch (DC->getContextKind()) {
+    case DeclContextKind::AbstractClosureExpr:
+    case DeclContextKind::AbstractFunctionDecl:
+      continue;
+    default:
       return nullptr;
     }
-
-  if (auto parent = dc->getParent()) {
-    return parent->getInnermostClosureForSelfCapture();
-  }
+  } while ((DC = DC->getParent()));
 
   return nullptr;
 }

lib/AST/UnqualifiedLookup.cpp

@@ -362,7 +362,7 @@ ValueDecl *UnqualifiedLookupFactory::lookupBaseDecl(const DeclContext *baseDC) c
   // Perform an unqualified lookup for the base decl of this result. This
   // handles cases where self was rebound (e.g. `guard let self = self`)
   // earlier in this closure or some outer closure.
-  auto closureExpr = DC->getInnermostClosureForSelfCapture();
+  auto closureExpr = DC->getInnermostClosureForCaptures();
   if (!closureExpr) {
     return nullptr;
   }

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/CSApply.cpp

@@ -5620,6 +5620,9 @@ namespace {
             .fixItInsert(coercion->getStartLoc(), "consume ");
       }
 
+      // If we're doing code completion, avoid doing any further type-checking,
+      // that should instead be handled by TypeCheckASTNodeAtLocRequest.
+      if (!ctx.CompletionCallback) {
         // Type-check any local decls encountered.
         for (auto *D : LocalDeclsToTypeCheck)
           TypeChecker::typeCheckDecl(D);
@@ -5632,6 +5635,7 @@ namespace {
                                   std::nullopt);
         }
       }
+    }
 
     /// Diagnose an optional injection that is probably not what the
     /// user wanted, because it comes from a forced downcast, or from an

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/MiscDiagnostics.cpp

@@ -2221,7 +2221,7 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
         return nullptr;
       }
 
-      return parentContext->getInnermostClosureForSelfCapture();
+      return parentContext->getInnermostClosureForCaptures();
     }
 
     bool shouldRecordClosure(const AbstractClosureExpr *E) {

lib/Sema/TypeCheckStmt.cpp

@@ -2628,14 +2628,25 @@ 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());
+
+      // 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)
+}