[CodeCompletion] Migrate key path completion to be solver based

This commit essentially consistes of the following steps:
- Add a new code completion key path component that represents the code completion token inside a key path. Previously, the key path would have an invalid component at the end if it contained a code completion token.
- When type checking the key path, model the code completion token’s result type by a new type variable that is unrelated to the previous components (because the code completion token might resolve to anything).
- Since the code completion token is now properly modelled in the constraint system, we can use the solver based code completion implementation and inspect any solution determined by the constraint solver. The base type for code completion is now the result type of the key path component that preceeds the code completion component.

This resolves bugs where code completion was not working correctly if the key path’s type had a generic base or result type. It’s also nice to have moved another completion type over to the solver-based implementation.

Resolves rdar://78779234 [SR-14685] and rdar://78779335 [SR-14703]

Author: ahoppen

include/swift/AST/Expr.h

@@ -5462,6 +5462,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 {
@@ -5476,6 +5477,7 @@ class KeyPathExpr : public Expr {
       Identity,
       TupleElement,
       DictionaryKey,
+      CodeCompletion,
     };
 
   private:
@@ -5651,8 +5653,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;
     }
@@ -5690,6 +5696,7 @@ class KeyPathExpr : public Expr {
       case Kind::UnresolvedSubscript:
       case Kind::UnresolvedProperty:
       case Kind::Invalid:
+      case Kind::CodeCompletion:
         return false;
       }
       llvm_unreachable("unhandled kind");
@@ -5710,6 +5717,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         return nullptr;
       }
       llvm_unreachable("unhandled kind");
@@ -5730,6 +5738,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");
@@ -5753,6 +5762,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         return {};
       }
       llvm_unreachable("unhandled kind");
@@ -5776,6 +5786,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");
@@ -5796,6 +5807,7 @@ class KeyPathExpr : public Expr {
       case Kind::Identity:
       case Kind::TupleElement:
       case Kind::DictionaryKey:
+      case Kind::CodeCompletion:
         return false;
       }
       llvm_unreachable("unhandled kind");
@@ -5816,6 +5828,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");
@@ -5836,6 +5849,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");

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

lib/AST/ASTDumper.cpp

@@ -2877,6 +2877,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) << "'";

lib/AST/ASTWalker.cpp

@@ -1126,6 +1126,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;
       }

lib/AST/Expr.cpp

@@ -2418,6 +2418,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");
   }
 }

lib/IDE/CodeCompletion.cpp

@@ -6742,6 +6742,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,
@@ -6833,6 +6855,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;
   }
@@ -6953,60 +6990,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: {

lib/IDE/SourceEntityWalker.cpp

@@ -435,6 +435,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;
       }
     }

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:

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);
 }
 

lib/SILGen/SILGenExpr.cpp

@@ -3741,6 +3741,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;