swiftlang · ahoppen · Jul 1, 2021 · Jun 23, 2021
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -5409,6 +5409,7 @@ class KeyPathExpr : public Expr {
   /// - a resolved ValueDecl, referring to
   /// - a subscript index expression, which may or may not be resolved
   /// - an optional chaining, forcing, or wrapping component
+  /// - a code completion token
   class Component {
   public:
     enum class Kind: unsigned {
@@ -5423,6 +5424,7 @@ class KeyPathExpr : public Expr {
       Identity,
       TupleElement,
       DictionaryKey,
+      CodeCompletion,
     };
 
   private:
@@ -5598,8 +5600,12 @@ class KeyPathExpr : public Expr {
                                      SourceLoc loc) {
       return Component(fieldNumber, elementType, loc);
     }
-
-
+
+    static Component forCodeCompletion(SourceLoc Loc) {
+      return Component(nullptr, {}, nullptr, {}, {}, Kind::CodeCompletion,
+                       Type(), Loc);
+    }
+
     SourceLoc getLoc() const {
       return Loc;
     }
@@ -5637,6 +5643,7 @@ class KeyPathExpr : public Expr {
       case Kind::UnresolvedSubscript:
       case Kind::UnresolvedProperty:
       case Kind::Invalid:
+      case Kind::CodeCompletion:
         return false;
       }
       llvm_unreachable("unhandled kind");
@@ -5657,6 +5664,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         return nullptr;
       }
       llvm_unreachable("unhandled kind");
@@ -5677,6 +5685,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         llvm_unreachable("no subscript labels for this kind");
       }
       llvm_unreachable("unhandled kind");
@@ -5700,6 +5709,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         return {};
       }
       llvm_unreachable("unhandled kind");
@@ -5723,6 +5733,7 @@ class KeyPathExpr : public Expr {
       case Kind::Property:
       case Kind::Identity:
       case Kind::TupleElement:
+      case Kind::CodeCompletion:
         llvm_unreachable("no unresolved name for this kind");
       }
       llvm_unreachable("unhandled kind");
@@ -5743,6 +5754,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         return false;
       }
       llvm_unreachable("unhandled kind");
@@ -5763,6 +5775,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         llvm_unreachable("no decl ref for this kind");
       }
       llvm_unreachable("unhandled kind");
@@ -5783,6 +5796,7 @@ class KeyPathExpr : public Expr {
         case Kind::Property:
         case Kind::Subscript:
         case Kind::DictionaryKey:
+        case Kind::CodeCompletion:
           llvm_unreachable("no field number for this kind");
       }
       llvm_unreachable("unhandled kind");

diff --git a/include/swift/Sema/CodeCompletionTypeChecking.h b/include/swift/Sema/CodeCompletionTypeChecking.h
@@ -116,6 +116,29 @@ namespace swift {
     void sawSolution(const constraints::Solution &solution) override;
   };
 
+  class KeyPathTypeCheckCompletionCallback
+      : public TypeCheckCompletionCallback {
+  public:
+    struct Result {
+      /// The type on which completion should occur, i.e. a result type of the
+      /// previous component.
+      Type BaseType;
+      /// Whether code completion happens on the key path's root.
+      bool OnRoot;
+    };
+
+  private:
+    KeyPathExpr *KeyPath;
+    SmallVector<Result, 4> Results;
+
+  public:
+    KeyPathTypeCheckCompletionCallback(KeyPathExpr *KeyPath)
+        : KeyPath(KeyPath) {}
+
+    ArrayRef<Result> getResults() const { return Results; }
+
+    void sawSolution(const constraints::Solution &solution) override;
+  };
 }
 
 #endif
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -2880,6 +2880,9 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
         PrintWithColorRAII(OS, IdentifierColor)
           << "  key='" << component.getUnresolvedDeclName() << "'";
         break;
+      case KeyPathExpr::Component::Kind::CodeCompletion:
+        PrintWithColorRAII(OS, ASTNodeColor) << "completion";
+        break;
       }
       PrintWithColorRAII(OS, TypeColor)
         << " type='" << GetTypeOfKeyPathComponent(E, i) << "'";

diff --git a/lib/AST/ASTWalker.cpp b/lib/AST/ASTWalker.cpp
@@ -1125,6 +1125,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       case KeyPathExpr::Component::Kind::Identity:
       case KeyPathExpr::Component::Kind::TupleElement:
       case KeyPathExpr::Component::Kind::DictionaryKey:
+      case KeyPathExpr::Component::Kind::CodeCompletion:
         // No subexpr to visit.
         break;
       }

diff --git a/lib/AST/Expr.cpp b/lib/AST/Expr.cpp
@@ -2399,6 +2399,7 @@ void KeyPathExpr::Component::setSubscriptIndexHashableConformances(
   case Kind::Identity:
   case Kind::TupleElement:
   case Kind::DictionaryKey:
+  case Kind::CodeCompletion:
     llvm_unreachable("no hashable conformances for this kind");
   }
 }

diff --git a/lib/IDE/CodeCompletion.cpp b/lib/IDE/CodeCompletion.cpp
@@ -6727,6 +6727,28 @@ void deliverUnresolvedMemberResults(
   deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
 }
 
+void deliverKeyPathResults(
+    ArrayRef<KeyPathTypeCheckCompletionCallback::Result> Results,
+    DeclContext *DC, SourceLoc DotLoc,
+    ide::CodeCompletionContext &CompletionCtx,
+    CodeCompletionConsumer &Consumer) {
+  ASTContext &Ctx = DC->getASTContext();
+  CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
+                          &CompletionCtx);
+
+  if (DotLoc.isValid()) {
+    Lookup.setHaveDot(DotLoc);
+  }
+  Lookup.shouldCheckForDuplicates(Results.size() > 1);
+
+  for (auto &Result : Results) {
+    Lookup.setIsSwiftKeyPathExpr(Result.OnRoot);
+    Lookup.getValueExprCompletions(Result.BaseType);
+  }
+
+  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+}
+
 void deliverDotExprResults(
     ArrayRef<DotExprTypeCheckCompletionCallback::Result> Results,
     Expr *BaseExpr, DeclContext *DC, SourceLoc DotLoc, bool IsInSelector,
@@ -6818,6 +6840,21 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
                                    CompletionContext, Consumer);
     return true;
   }
+  case CompletionKind::KeyPathExprSwift: {
+    assert(CurDeclContext);
+
+    // CodeCompletionCallbacks::completeExprKeyPath takes a \c KeyPathExpr,
+    // so we can safely cast the \c ParsedExpr back to a \c KeyPathExpr.
+    auto KeyPath = cast<KeyPathExpr>(ParsedExpr);
+    KeyPathTypeCheckCompletionCallback Lookup(KeyPath);
+    llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
+        Context.CompletionCallback, &Lookup);
+    typeCheckContextAt(CurDeclContext, CompletionLoc);
+
+    deliverKeyPathResults(Lookup.getResults(), CurDeclContext, DotLoc,
+                          CompletionContext, Consumer);
+    return true;
+  }
   default:
     return false;
   }
@@ -6938,60 +6975,10 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::None:
   case CompletionKind::DotExpr:
   case CompletionKind::UnresolvedMember:
+  case CompletionKind::KeyPathExprSwift:
     llvm_unreachable("should be already handled");
     return;
 
-  case CompletionKind::KeyPathExprSwift: {
-    auto KPE = dyn_cast<KeyPathExpr>(ParsedExpr);
-    auto BGT = (*ExprType)->getAs<BoundGenericType>();
-    if (!KPE || !BGT || BGT->getGenericArgs().size() != 2)
-      break;
-    assert(!KPE->isObjC());
-
-    if (DotLoc.isValid())
-      Lookup.setHaveDot(DotLoc);
-
-    bool OnRoot = !KPE->getComponents().front().isValid();
-    Lookup.setIsSwiftKeyPathExpr(OnRoot);
-
-    Type baseType = BGT->getGenericArgs()[OnRoot ? 0 : 1];
-    if (OnRoot && baseType->is<UnresolvedType>()) {
-      // Infer the root type of the keypath from the context type.
-      ExprContextInfo ContextInfo(CurDeclContext, ParsedExpr);
-      for (auto T : ContextInfo.getPossibleTypes()) {
-        if (auto unwrapped = T->getOptionalObjectType())
-          T = unwrapped;
-
-        // If the context type is any of the KeyPath types, use it.
-        if (T->getAnyNominal() && T->getAnyNominal()->getKeyPathTypeKind() &&
-            !T->hasUnresolvedType() && T->is<BoundGenericType>()) {
-          baseType = T->castTo<BoundGenericType>()->getGenericArgs()[0];
-          break;
-        }
-
-        // KeyPath can be used as a function that receives its root type.
-        if (T->is<AnyFunctionType>()) {
-          auto *fnType = T->castTo<AnyFunctionType>();
-          if (fnType->getNumParams() == 1) {
-            const AnyFunctionType::Param &param = fnType->getParams()[0];
-            baseType = param.getParameterType();
-            break;
-          }
-        }
-      }
-    }
-    if (!OnRoot && KPE->getComponents().back().getKind() ==
-                       KeyPathExpr::Component::Kind::OptionalWrap) {
-      // KeyPath expr with '?' (e.g. '\Ty.[0].prop?.another').
-      // Although expected type is optional, we should unwrap it because it's
-      // unwrapped.
-      baseType = baseType->getOptionalObjectType();
-    }
-
-    Lookup.getValueExprCompletions(baseType);
-    break;
-  }
-
   case CompletionKind::StmtOrExpr:
   case CompletionKind::ForEachSequence:
   case CompletionKind::PostfixExprBeginning: {

diff --git a/lib/IDE/SourceEntityWalker.cpp b/lib/IDE/SourceEntityWalker.cpp
@@ -494,6 +494,7 @@ std::pair<bool, Expr *> SemaAnnotator::walkToExprPre(Expr *E) {
       case KeyPathExpr::Component::Kind::OptionalForce:
       case KeyPathExpr::Component::Kind::Identity:
       case KeyPathExpr::Component::Kind::DictionaryKey:
+      case KeyPathExpr::Component::Kind::CodeCompletion:
         break;
       }
     }

diff --git a/lib/IDE/SwiftSourceDocInfo.cpp b/lib/IDE/SwiftSourceDocInfo.cpp
@@ -440,6 +440,7 @@ std::pair<bool, Expr*> NameMatcher::walkToExprPre(Expr *E) {
             break;
           case KeyPathExpr::Component::Kind::DictionaryKey:
           case KeyPathExpr::Component::Kind::Invalid:
+          case KeyPathExpr::Component::Kind::CodeCompletion:
             break;
           case KeyPathExpr::Component::Kind::OptionalForce:
           case KeyPathExpr::Component::Kind::OptionalChain:

diff --git a/lib/Parse/ParseExpr.cpp b/lib/Parse/ParseExpr.cpp
@@ -659,21 +659,26 @@ ParserResult<Expr> Parser::parseExprKeyPath() {
   if (rootResult.isNull() && pathResult.isNull())
     return nullptr;
 
-  auto keypath =
-      new (Context) KeyPathExpr(backslashLoc, rootResult.getPtrOrNull(),
-                                pathResult.getPtrOrNull(), hasLeadingDot);
-
   // Handle code completion.
   if ((Tok.is(tok::code_complete) && !Tok.isAtStartOfLine()) ||
       (Tok.is(tok::period) && peekToken().isAny(tok::code_complete))) {
     SourceLoc DotLoc;
     consumeIf(tok::period, DotLoc);
+
+    // Add the code completion expression to the path result.
+    CodeCompletionExpr *CC = new (Context)
+        CodeCompletionExpr(pathResult.getPtrOrNull(), Tok.getLoc());
+    auto keypath = new (Context)
+        KeyPathExpr(backslashLoc, rootResult.getPtrOrNull(), CC, hasLeadingDot);
     if (CodeCompletion)
       CodeCompletion->completeExprKeyPath(keypath, DotLoc);
     consumeToken(tok::code_complete);
     return makeParserCodeCompletionResult(keypath);
   }
 
+  auto keypath =
+      new (Context) KeyPathExpr(backslashLoc, rootResult.getPtrOrNull(),
+                                pathResult.getPtrOrNull(), hasLeadingDot);
   return makeParserResult(parseStatus, keypath);
 }
 

diff --git a/lib/SILGen/SILGenExpr.cpp b/lib/SILGen/SILGenExpr.cpp
@@ -3748,6 +3748,7 @@ RValue RValueEmitter::visitKeyPathExpr(KeyPathExpr *E, SGFContext C) {
     case KeyPathExpr::Component::Kind::Invalid:
     case KeyPathExpr::Component::Kind::UnresolvedProperty:
     case KeyPathExpr::Component::Kind::UnresolvedSubscript:
+    case KeyPathExpr::Component::Kind::CodeCompletion:
       llvm_unreachable("not resolved");
       break;
 

diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -377,6 +377,7 @@ static bool buildObjCKeyPathString(KeyPathExpr *E,
     case KeyPathExpr::Component::Kind::Invalid:
     case KeyPathExpr::Component::Kind::UnresolvedProperty:
     case KeyPathExpr::Component::Kind::UnresolvedSubscript:
+    case KeyPathExpr::Component::Kind::CodeCompletion:
       // Don't bother building the key path string if the key path didn't even
       // resolve.
       return false;
@@ -4938,7 +4939,8 @@ namespace {
           }
           break;
         }
-        case KeyPathExpr::Component::Kind::Invalid: {
+        case KeyPathExpr::Component::Kind::Invalid:
+        case KeyPathExpr::Component::Kind::CodeCompletion: {
           auto component = origComponent;
           component.setComponentType(leafTy);
           resolvedComponents.push_back(component);

diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -3122,6 +3122,7 @@ namespace {
         if (!rootObjectTy || rootObjectTy->hasError())
           return Type();
 
+        CS.setType(rootRepr, rootObjectTy);
         // Allow \Derived.property to be inferred as \Base.property to
         // simulate a sort of covariant conversion from
         // KeyPath<Derived, T> to KeyPath<Base, T>.
@@ -3143,7 +3144,13 @@ namespace {
         switch (auto kind = component.getKind()) {
         case KeyPathExpr::Component::Kind::Invalid:
           break;
-
+        case KeyPathExpr::Component::Kind::CodeCompletion:
+          // We don't know what the code completion might resolve to, so we are
+          // creating a new type variable for its result, which might be a hole.
+          base = CS.createTypeVariable(
+              resultLocator,
+              TVO_CanBindToLValue | TVO_CanBindToNoEscape | TVO_CanBindToHole);
+          break;
         case KeyPathExpr::Component::Kind::UnresolvedProperty:
         // This should only appear in resolved ASTs, but we may need to
         // re-type-check the constraints during failure diagnosis.
@@ -3239,7 +3246,8 @@ namespace {
       // If there was an optional chaining component, the end result must be
       // optional.
       if (didOptionalChain) {
-        auto objTy = CS.createTypeVariable(locator, TVO_CanBindToNoEscape);
+        auto objTy = CS.createTypeVariable(locator, TVO_CanBindToNoEscape |
+                                                        TVO_CanBindToHole);
         componentTypeVars.push_back(objTy);
 
         auto optTy = OptionalType::get(objTy);
@@ -3250,8 +3258,8 @@ namespace {
 
       auto baseLocator =
           CS.getConstraintLocator(E, ConstraintLocator::KeyPathValue);
-      auto rvalueBase = CS.createTypeVariable(baseLocator,
-                                              TVO_CanBindToNoEscape);
+      auto rvalueBase = CS.createTypeVariable(
+          baseLocator, TVO_CanBindToNoEscape | TVO_CanBindToHole);
       CS.addConstraint(ConstraintKind::Equal, base, rvalueBase, locator);
 
       // The result is a KeyPath from the root to the end component.

diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -9488,7 +9488,12 @@ ConstraintSystem::simplifyKeyPathConstraint(
     case KeyPathExpr::Component::Kind::Invalid:
     case KeyPathExpr::Component::Kind::Identity:
       break;
-
+
+    case KeyPathExpr::Component::Kind::CodeCompletion: {
+      anyComponentsUnresolved = true;
+      capability = ReadOnly;
+      break;
+    }
     case KeyPathExpr::Component::Kind::Property:
     case KeyPathExpr::Component::Kind::Subscript:
     case KeyPathExpr::Component::Kind::UnresolvedProperty: