[CodeCompletion] Fix issue when completing inside a switch in a closure

Author: ahoppen

lib/IDE/ArgumentCompletion.cpp

@@ -88,6 +88,15 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   while (ParentCall && ParentCall->getArgs() == nullptr) {
     ParentCall = CS.getParentExpr(ParentCall);
   }
+  if (auto TV = S.getType(CompletionExpr)->getAs<TypeVariableType>()) {
+    auto Locator = TV->getImpl().getLocator();
+    if (Locator->isLastElement<LocatorPathElt::PatternMatch>()) {
+      // The code completion token is inside a pattern, which got rewritten from
+      // a call by ResolvePattern. Thus, we aren't actually inside a call.
+      // Rest 'ParentCall' to nullptr to reflect that.
+      ParentCall = nullptr;
+    }
+  }
 
   if (!ParentCall || ParentCall == CompletionExpr) {
     // We might not have a call that contains the code completion expression if

lib/Parse/ParseExpr.cpp

@@ -1175,6 +1175,7 @@ Parser::parseExprPostfixSuffix(ParserResult<Expr> Result, bool isExprBasic,
       //   [.foo(), <HERE> .bar()]
       // '.bar()' is probably not a part of the inserting element. Moreover,
       // having suffixes doesn't help type inference in any way.
+
       return Result;
     }
 

lib/Parse/ParsePattern.cpp

@@ -1303,6 +1303,19 @@ ParserResult<Pattern> Parser::parseMatchingPattern(bool isExprBasic) {
   if (subExpr.isNull())
     return status;
 
+  if (isa<CodeCompletionExpr>(subExpr.get()) && Tok.isFollowingLParen()) {
+    // We are in the case like the following of parsing a pattern with the code
+    // completion token as base and associated value matches:
+    // #^COMPLETE^#(let a)
+    // We will have not consumed the `(let a)` in `parseExprPostfixSuffix`
+    // because usually suffixes don't influence the code completion's type and
+    // the suffix might be unrelated. But the trailing `(let a)` that is left
+    // prevents us from forming a valid pattern.
+    // Consume and discard the `(let a)`, which just leaves us with the base
+    // of the pattern.
+    (void)parseExprCallSuffix(subExpr, isExprBasic);
+  }
+
   // The most common case here is to parse something that was a lexically
   // obvious pattern, which will come back wrapped in an immediate
   // UnresolvedPatternExpr.  Transform this now to simplify later code.

lib/Sema/CSGen.cpp

@@ -2390,6 +2390,12 @@ namespace {
       // function, to set the type of the pattern.
       auto setType = [&](Type type) {
         CS.setType(pattern, type);
+        if (auto PE = dyn_cast<ExprPattern>(pattern)) {
+          // Set the type of the pattern's sub-expression as well, so code
+          // completion can retrieve the expression's type in case it is a code
+          // completion token.
+          CS.setType(PE->getSubExpr(), type);
+        }
         return type;
       };
 

lib/Sema/TypeCheckStmt.cpp

@@ -2393,6 +2393,8 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
     }
   }
 
+  bool LeaveBodyUnchecked = !ctx.CompletionCallback;
+
   // The enclosing closure might be a single expression closure or a function
   // builder closure. In such cases, the body elements are type checked with
   // the closure itself. So we need to try type checking the enclosing closure
@@ -2416,13 +2418,16 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
       auto ActorIsolation = determineClosureActorIsolation(
           CE, __Expr_getType, __AbstractClosureExpr_getActorIsolation);
       CE->setActorIsolation(ActorIsolation);
+      if (!LeaveBodyUnchecked) {
+        // Type checking the parent closure also type checked this node.
+        // Nothing to do anymore.
+        return false;
+      }
       if (CE->getBodyState() != ClosureExpr::BodyState::ReadyForTypeChecking)
         return false;
     }
   }
 
-  bool LeaveBodyUnchecked = !ctx.CompletionCallback;
-
   TypeChecker::typeCheckASTNode(finder.getRef(), DC, LeaveBodyUnchecked);
   return false;
 }

test/IDE/complete_in_closures.swift

@@ -416,3 +416,83 @@ func testClosureInPatternBindingInit() {
   // CLOSURE_IN_PATTERN_BINDING: End completion
 
 }
+
+func testSwitchInClosure() {
+  func executeClosure(closure: () -> Void) {}
+
+  struct Boredom {
+    static func doNothing() {}
+  }
+
+  enum MyEnum {
+    case first
+    case second
+  }
+
+  let item: MyEnum? = nil
+  executeClosure(closure: {
+    switch item {
+    case .#^SWITCH_IN_CLOSURE^#first:
+      break
+    case .second:
+      Boredom.doNothing()
+    }
+  })
+
+// SWITCH_IN_CLOSURE: Begin completions
+// SWITCH_IN_CLOSURE-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: first[#MyEnum#];
+// SWITCH_IN_CLOSURE-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: second[#MyEnum#];
+// SWITCH_IN_CLOSURE: End completions
+}
+
+func testSwitchWithAssociatedValueInClosure() {
+  func executeClosure(closure: () -> Void) {}
+
+  struct Boredom {
+    static func doNothing() {}
+  }
+
+  enum MyEnum {
+    case first(String)
+  }
+
+  let item: MyEnum? = nil
+  executeClosure(closure: {
+    switch item {
+    case .#^SWITCH_WITH_ASSOC_IN_CLOSURE^#first(_):
+      break
+    }
+  })
+
+// SWITCH_WITH_ASSOC_IN_CLOSURE: Begin completions
+// SWITCH_WITH_ASSOC_IN_CLOSURE-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: first({#String#})[#MyEnum#];
+// SWITCH_WITH_ASSOC_IN_CLOSURE: End completions
+}
+
+func testCompleteInMatchOfAssociatedValueInSwitchCase() {
+  func testSwitchWithAssociatedValueInClosure() {
+  func executeClosure(closure: () -> Void) {}
+
+  struct Boredom {
+    static func doNothing() {}
+  }
+
+  enum MyEnum {
+    case first(Bool, String)
+  }
+
+  let item: MyEnum? = nil
+  let str = "hi"
+  executeClosure(closure: {
+    switch item {
+    case .first(true, #^IN_ASSOC_OF_CASE_IN_CLOSURE^#str):
+      break
+    }
+  })
+
+// IN_ASSOC_OF_CASE_IN_CLOSURE: Begin completions
+// IN_ASSOC_OF_CASE_IN_CLOSURE-DAG: Decl[LocalVar]/Local:               str[#String#]; name=str
+// IN_ASSOC_OF_CASE_IN_CLOSURE: End completions
+}
+
+}

test/IDE/complete_in_result_builder.swift

@@ -329,3 +329,39 @@ func testInStringLiteralInResultBuilder() {
 // IN_STRING_LITERAL_IN_RESULT_BUILDER-DAG: Decl[InstanceVar]/CurrNominal:      bar[#Int#]; name=bar
 // IN_STRING_LITERAL_IN_RESULT_BUILDER: End completions
 }
+
+func testSwitchInResultBuilder() {
+  @resultBuilder
+  enum ReducerBuilder2<Action> {
+    static func buildBlock(_ r: Reduce2<Action>) -> Reduce2<Action> { r }
+    static func buildBlock(_ r0: Reduce2<Action>, _ r1: Reduce2<Action>) -> Reduce2<Action> { r0 }
+    static func buildExpression(_ r: Reduce2<Action>) -> Reduce2<Action> { r }
+  }
+
+  enum Action {
+    case alertDismissed
+  }
+
+  struct Reduce2<Action> {
+    init() {}
+
+    init(_ reduce: (Action) -> Int) {}
+  }
+
+  struct Login2 {
+    @ReducerBuilder2<Action>
+    var body: Reduce2<Action> {
+      Reduce2()
+      Reduce2 { action in
+        switch action {
+        case .#^SWITCH_IN_RESULT_BUILDER^# alertDismissed:
+          return 0
+        }
+      }
+    }
+  }
+// SWITCH_IN_RESULT_BUILDER: Begin completions, 2 items
+// SWITCH_IN_RESULT_BUILDER-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: alertDismissed[#Action#];
+// SWITCH_IN_RESULT_BUILDER-DAG: Decl[InstanceMethod]/CurrNominal/TypeRelation[Invalid]: hash({#(self): Action#})[#(into: inout Hasher) -> Void#];
+// SWITCH_IN_RESULT_BUILDER: End completions
+}

test/IDE/complete_unresolved_members.swift

@@ -732,3 +732,14 @@ func testSameType() {
 // SUGAR_TYPE: End completions
 }
 
+struct DispatchTime {
+  static func now() -> DispatchTime { .init() }
+}
+func +(_ x: DispatchTime, _ y: Double) -> DispatchTime { return x }
+
+let _: DispatchTime = .#^UNRESOLVED_FUNCTION_CALL^#now() + 0.2
+
+// UNRESOLVED_FUNCTION_CALL: Begin completions, 2 items
+// UNRESOLVED_FUNCTION_CALL-DAG: Decl[StaticMethod]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: now()[#DispatchTime#];
+// UNRESOLVED_FUNCTION_CALL-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Convertible]: init()[#DispatchTime#];
+// UNRESOLVED_FUNCTION_CALL: End completions