[Sema] Fix two assertion failures/crashes in solver-based key-path completion

include/swift/Sema/ConstraintSystem.h

@@ -1293,6 +1293,10 @@ class Solution {
 
   bool hasType(ASTNode node) const;
 
+  /// Returns \c true if the \p ComponentIndex-th component in \p KP has a type
+  /// associated with it.
+  bool hasType(const KeyPathExpr *KP, unsigned ComponentIndex) const;
+
   /// Retrieve the type of the given node, as recorded in this solution.
   Type getType(ASTNode node) const;
 

lib/Sema/CSApply.cpp

@@ -8899,6 +8899,11 @@ bool Solution::hasType(ASTNode node) const {
   return cs.hasType(node);
 }
 
+bool Solution::hasType(const KeyPathExpr *KP, unsigned ComponentIndex) const {
+  auto &cs = getConstraintSystem();
+  return cs.hasType(KP, ComponentIndex);
+}
+
 Type Solution::getType(ASTNode node) const {
   auto result = nodeTypes.find(node);
   if (result != nodeTypes.end())

lib/Sema/CSGen.cpp

@@ -3233,7 +3233,7 @@ namespace {
           break;
         }
         case KeyPathExpr::Component::Kind::Identity:
-          continue;
+          break;
         case KeyPathExpr::Component::Kind::DictionaryKey:
           llvm_unreachable("DictionaryKey only valid in #keyPath");
           break;

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -618,10 +618,10 @@ class CompletionContextFinder : public ASTWalker {
   /// If we are completing inside an expression, the \c CodeCompletionExpr that
   /// represents the code completion token.
 
-  /// The AST node that represents the code completion token, either as an
-  /// expression or a KeyPath component.
-  llvm::PointerUnion<CodeCompletionExpr *, const KeyPathExpr::Component *>
-      CompletionNode;
+  /// The AST node that represents the code completion token, either as a
+  /// \c CodeCompletionExpr or a \c KeyPathExpr which contains a code completion
+  /// component.
+  llvm::PointerUnion<CodeCompletionExpr *, const KeyPathExpr *> CompletionNode;
 
   Expr *InitialExpr = nullptr;
   DeclContext *InitialDC;
@@ -675,10 +675,13 @@ class CompletionContextFinder : public ASTWalker {
       for (auto &component : KeyPath->getComponents()) {
         if (component.getKind() ==
             KeyPathExpr::Component::Kind::CodeCompletion) {
-          CompletionNode = &component;
+          CompletionNode = KeyPath;
           return std::make_pair(false, nullptr);
         }
       }
+      // Code completion in key paths is modelled by a code completion component
+      // Don't walk the key path's parsed expressions.
+      return std::make_pair(false, E);
     }
 
     return std::make_pair(true, E);
@@ -713,12 +716,33 @@ class CompletionContextFinder : public ASTWalker {
   }
 
   bool hasCompletionKeyPathComponent() const {
-    return CompletionNode.dyn_cast<const KeyPathExpr::Component *>() != nullptr;
+    return CompletionNode.dyn_cast<const KeyPathExpr *>() != nullptr;
   }
 
-  const KeyPathExpr::Component *getCompletionKeyPathComponent() const {
+  /// If we are completing in a key path, returns the \c KeyPath that contains
+  /// the code completion component.
+  const KeyPathExpr *getKeyPathContainingCompletionComponent() const {
     assert(hasCompletionKeyPathComponent());
-    return CompletionNode.get<const KeyPathExpr::Component *>();
+    return CompletionNode.get<const KeyPathExpr *>();
+  }
+
+  /// If we are completing in a key path, returns the index at which the key
+  /// path has the code completion component.
+  size_t getKeyPathCompletionComponentIndex() const {
+    assert(hasCompletionKeyPathComponent());
+    size_t ComponentIndex = 0;
+    auto Components =
+        getKeyPathContainingCompletionComponent()->getComponents();
+    for (auto &Component : Components) {
+      if (Component.getKind() == KeyPathExpr::Component::Kind::CodeCompletion) {
+        break;
+      } else {
+        ComponentIndex++;
+      }
+    }
+    assert(ComponentIndex < Components.size() &&
+           "No completion component in the key path?");
+    return ComponentIndex;
   }
 
   struct Fallback {
@@ -949,20 +973,27 @@ bool TypeChecker::typeCheckForCodeCompletion(
     if (contextAnalyzer.locatedInMultiStmtClosure()) {
       auto &solution = solutions.front();
 
-      if (solution.hasType(contextAnalyzer.getCompletionExpr())) {
-        llvm::for_each(solutions, callback);
-        return CompletionResult::Ok;
+      bool HasTypeForCompletionNode = false;
+      if (completionExpr) {
+        HasTypeForCompletionNode = solution.hasType(completionExpr);
+      } else {
+        assert(contextAnalyzer.hasCompletionKeyPathComponent());
+        HasTypeForCompletionNode = solution.hasType(
+            contextAnalyzer.getKeyPathContainingCompletionComponent(),
+            contextAnalyzer.getKeyPathCompletionComponentIndex());
       }
 
-      // At this point we know the code completion expression wasn't checked
-      // with the closure's surrounding context, so can defer to regular type-
-      // checking for the current call to typeCheckExpression. If that succeeds
-      // we will get a second call to typeCheckExpression for the body of the
-      // closure later and can gather completions then. If it doesn't we rely
-      // on the fallback typechecking in the subclasses of
-      // TypeCheckCompletionCallback that considers in isolation a
-      // sub-expression of the closure that contains the completion location.
-      return CompletionResult::NotApplicable;
+      if (!HasTypeForCompletionNode) {
+        // At this point we know the code completion node wasn't checked with
+        // the closure's surrounding context, so can defer to regular
+        // type-checking for the current call to typeCheckExpression. If that
+        // succeeds we will get a second call to typeCheckExpression for the
+        // body of the closure later and can gather completions then. If it
+        // doesn't we rely on the fallback typechecking in the subclasses of
+        // TypeCheckCompletionCallback that considers in isolation a
+        // sub-expression of the closure that contains the completion location.
+        return CompletionResult::NotApplicable;
+      }
     }
 
     llvm::for_each(solutions, callback);

test/IDE/complete_swift_key_path.swift

@@ -35,6 +35,10 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=GENERIC_KEY_PATH_BASE | %FileCheck %s -check-prefix=PERSONTYPE-DOT
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=GENERIC_KEY_PATH_RESULT | %FileCheck %s -check-prefix=PERSONTYPE-DOT
 
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=COMPLETE_AFTER_SELF | %FileCheck %s -check-prefix=OBJ-NODOT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_RESULT_BUILDER | %FileCheck %s -check-prefix=PERSONTYPE-DOT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=IN_MULTI_STMT_CLOSURE | %FileCheck %s -check-prefix=PERSONTYPE-DOT
+
 class Person {
     var name: String
     var friends: [Person] = []
@@ -191,3 +195,36 @@ func genericKeyPathResult<KeyPathResult>(id: KeyPath<Person, KeyPathResult>) {
   genericKeyPathResult(\.#^GENERIC_KEY_PATH_RESULT^#)
   // Same as TYPE_DOT.
 }
+
+func completeAfterSelf(people: [Person]) {
+  people.map(\.self#^COMPLETE_AFTER_SELF^#)
+}
+
+func inResultBuilder() {
+  protocol View2 {}
+
+  @resultBuilder public struct ViewBuilder2 {
+    public static func buildBlock<Content>(_ content: Content) -> Content where Content : View2 { fatalError() }
+    public static func buildIf<Content>(_ content: Content?) -> Content? where Content : View2 { fatalError() }
+  }
+
+  struct VStack2<Content>: View2 {
+    init(@ViewBuilder2 view: () -> Content) {}
+  }
+
+  @ViewBuilder2 var body: some View2 {
+    VStack2 {
+      if true {
+        var people: [Person] = []
+        people.map(\.#^IN_RESULT_BUILDER^#)
+      }
+    }
+  }
+}
+
+func inMultiStmtClosure(closure: () -> Void) {
+  inMultiStmtClosure {
+    var people: [Person] = []
+    people.map(\.#^IN_MULTI_STMT_CLOSURE^#)
+  }
+}