[Sema/Index] Resolve #keyPath components so they get handled by indexing, semantic highlighting, etc.

include/swift/AST/Expr.h

@@ -5209,6 +5209,7 @@ class KeyPathExpr : public Expr {
       OptionalWrap,
       Identity,
       TupleElement,
+      DictionaryKey,
     };
 
   private:
@@ -5317,6 +5318,16 @@ class KeyPathExpr : public Expr {
                        propertyType,
                        loc);
     }
+
+    /// Create a component for a dictionary key (#keyPath only).
+    static Component forDictionaryKey(DeclNameRef UnresolvedName,
+                                      Type valueType,
+                                      SourceLoc loc) {
+      return Component(nullptr, UnresolvedName, nullptr, {}, {},
+                       Kind::DictionaryKey,
+                       valueType,
+                       loc);
+    }
 
     /// Create a component for a subscript.
     static Component forSubscript(ASTContext &ctx,
@@ -5407,6 +5418,7 @@ class KeyPathExpr : public Expr {
       case Kind::Property:
       case Kind::Identity:
       case Kind::TupleElement:
+      case Kind::DictionaryKey:
         return true;
 
       case Kind::UnresolvedSubscript:
@@ -5431,6 +5443,7 @@ class KeyPathExpr : public Expr {
       case Kind::Property:
       case Kind::Identity:
       case Kind::TupleElement:
+      case Kind::DictionaryKey:
         return nullptr;
       }
       llvm_unreachable("unhandled kind");
@@ -5450,6 +5463,7 @@ class KeyPathExpr : public Expr {
       case Kind::Property:
       case Kind::Identity:
       case Kind::TupleElement:
+      case Kind::DictionaryKey:
         llvm_unreachable("no subscript labels for this kind");
       }
       llvm_unreachable("unhandled kind");
@@ -5472,6 +5486,7 @@ class KeyPathExpr : public Expr {
       case Kind::Property:
       case Kind::Identity:
       case Kind::TupleElement:
+      case Kind::DictionaryKey:
         return {};
       }
       llvm_unreachable("unhandled kind");
@@ -5483,6 +5498,7 @@ class KeyPathExpr : public Expr {
     DeclNameRef getUnresolvedDeclName() const {
       switch (getKind()) {
       case Kind::UnresolvedProperty:
+      case Kind::DictionaryKey:
         return Decl.UnresolvedName;
 
       case Kind::Invalid:
@@ -5513,6 +5529,7 @@ class KeyPathExpr : public Expr {
       case Kind::OptionalForce:
       case Kind::Identity:
       case Kind::TupleElement:
+      case Kind::DictionaryKey:
         llvm_unreachable("no decl ref for this kind");
       }
       llvm_unreachable("unhandled kind");
@@ -5532,6 +5549,7 @@ class KeyPathExpr : public Expr {
         case Kind::Identity:
         case Kind::Property:
         case Kind::Subscript:
+        case Kind::DictionaryKey:
           llvm_unreachable("no field number for this kind");
       }
       llvm_unreachable("unhandled kind");

lib/AST/ASTDumper.cpp

@@ -2821,6 +2821,11 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
         PrintWithColorRAII(OS, DiscriminatorColor)
           << "#" << component.getTupleIndex();
         break;
+      case KeyPathExpr::Component::Kind::DictionaryKey:
+        PrintWithColorRAII(OS, ASTNodeColor) << "dict_key";
+        PrintWithColorRAII(OS, IdentifierColor)
+          << "  key='" << component.getUnresolvedDeclName() << "'";
+        break;
       }
       PrintWithColorRAII(OS, TypeColor)
         << " type='" << GetTypeOfKeyPathComponent(E, i) << "'";

lib/AST/ASTWalker.cpp

@@ -1120,6 +1120,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       case KeyPathExpr::Component::Kind::Invalid:
       case KeyPathExpr::Component::Kind::Identity:
       case KeyPathExpr::Component::Kind::TupleElement:
+      case KeyPathExpr::Component::Kind::DictionaryKey:
         // No subexpr to visit.
         break;
       }

lib/AST/Expr.cpp

@@ -2367,6 +2367,7 @@ void KeyPathExpr::Component::setSubscriptIndexHashableConformances(
   case Kind::Property:
   case Kind::Identity:
   case Kind::TupleElement:
+  case Kind::DictionaryKey:
     llvm_unreachable("no hashable conformances for this kind");
   }
 }

lib/IDE/SourceEntityWalker.cpp

@@ -415,6 +415,7 @@ std::pair<bool, Expr *> SemaAnnotator::walkToExprPre(Expr *E) {
       case KeyPathExpr::Component::Kind::OptionalWrap:
       case KeyPathExpr::Component::Kind::OptionalForce:
       case KeyPathExpr::Component::Kind::Identity:
+      case KeyPathExpr::Component::Kind::DictionaryKey:
         break;
       }
     }

lib/SILGen/SILGenExpr.cpp

@@ -3724,6 +3724,11 @@ RValue RValueEmitter::visitKeyPathExpr(KeyPathExpr *E, SGFContext C) {
     case KeyPathExpr::Component::Kind::UnresolvedProperty:
     case KeyPathExpr::Component::Kind::UnresolvedSubscript:
       llvm_unreachable("not resolved");
+      break;
+
+    case KeyPathExpr::Component::Kind::DictionaryKey:
+      llvm_unreachable("DictionaryKey only valid in #keyPath");
+      break;
     }
   }
 

lib/Sema/CSApply.cpp

@@ -237,6 +237,9 @@ static bool buildObjCKeyPathString(KeyPathExpr *E,
       // Don't bother building the key path string if the key path didn't even
       // resolve.
       return false;
+    case KeyPathExpr::Component::Kind::DictionaryKey:
+      llvm_unreachable("DictionaryKey only valid in #keyPath expressions.");
+      return false;
     }
   }
 
@@ -4680,6 +4683,10 @@ namespace {
         case KeyPathExpr::Component::Kind::OptionalWrap:
         case KeyPathExpr::Component::Kind::TupleElement:
           llvm_unreachable("already resolved");
+          break;
+        case KeyPathExpr::Component::Kind::DictionaryKey:
+          llvm_unreachable("DictionaryKey only valid in #keyPath");
+          break;
         }
 
         // Update "componentTy" with the result type of the last component.
@@ -7593,9 +7600,8 @@ namespace {
             componentType = solution.simplifyType(cs.getType(kp, i));
             assert(!componentType->hasTypeVariable() &&
                    "Should not write type variable into key-path component");
+            kp->getMutableComponents()[i].setComponentType(componentType);
           }
-
-          kp->getMutableComponents()[i].setComponentType(componentType);
         }
       }
 

lib/Sema/CSGen.cpp

@@ -3527,6 +3527,9 @@ namespace {
         }
         case KeyPathExpr::Component::Kind::Identity:
           continue;
+        case KeyPathExpr::Component::Kind::DictionaryKey:
+          llvm_unreachable("DictionaryKey only valid in #keyPath");
+          break;
         }
 
         // By now, `base` is the result type of this component. Set it in the

lib/Sema/CSSimplify.cpp

@@ -8051,6 +8051,9 @@ ConstraintSystem::simplifyKeyPathConstraint(
     case KeyPathExpr::Component::Kind::TupleElement:
       llvm_unreachable("not implemented");
       break;
+    case KeyPathExpr::Component::Kind::DictionaryKey:
+      llvm_unreachable("DictionaryKey only valid in #keyPath");
+      break;
     }
   }
 

lib/Sema/ConstraintSystem.cpp

@@ -493,6 +493,7 @@ ConstraintLocator *ConstraintSystem::getCalleeLocator(
     case ComponentKind::OptionalChain:
     case ComponentKind::OptionalWrap:
     case ComponentKind::Identity:
+    case ComponentKind::DictionaryKey:
       // These components don't have any callee associated, so just continue.
       break;
     }

lib/Sema/TypeCheckAvailability.cpp

@@ -2515,6 +2515,7 @@ class AvailabilityWalker : public ASTWalker {
       case KeyPathExpr::Component::Kind::OptionalWrap:
       case KeyPathExpr::Component::Kind::OptionalForce:
       case KeyPathExpr::Component::Kind::Identity:
+      case KeyPathExpr::Component::Kind::DictionaryKey:
         break;
       }
     }

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -622,9 +622,21 @@ static Optional<Type> getTypeOfCompletionContextExpr(
 
   case CompletionTypeCheckKind::KeyPath:
     referencedDecl = nullptr;
-    if (auto keyPath = dyn_cast<KeyPathExpr>(parsedExpr))
-      return TypeChecker::checkObjCKeyPathExpr(DC, keyPath,
-                                               /*requireResultType=*/true);
+    if (auto keyPath = dyn_cast<KeyPathExpr>(parsedExpr)) {
+      auto components = keyPath->getComponents();
+      if (!components.empty()) {
+        auto &last = components.back();
+        if (last.isResolved()) {
+          if (last.getKind() == KeyPathExpr::Component::Kind::Property)
+            referencedDecl = last.getDeclRef();
+          Type lookupTy = last.getComponentType();
+          ASTContext &Ctx = DC->getASTContext();
+          if (auto bridgedClass = Ctx.getBridgedToObjC(DC, lookupTy))
+            return bridgedClass;
+          return lookupTy;
+        }
+      }
+    }
 
     return None;
   }

lib/Sema/TypeCheckExprObjC.cpp

@@ -221,6 +221,7 @@ Optional<Type> TypeChecker::checkObjCKeyPathExpr(DeclContext *dc,
     case KeyPathExpr::Component::Kind::OptionalWrap:
     case KeyPathExpr::Component::Kind::Property:
     case KeyPathExpr::Component::Kind::Subscript:
+    case KeyPathExpr::Component::Kind::DictionaryKey:
       llvm_unreachable("already resolved!");
     }
 
@@ -241,6 +242,9 @@ Optional<Type> TypeChecker::checkObjCKeyPathExpr(DeclContext *dc,
       // From here, we're resolving a property. Use the current type.
       updateState(/*isProperty=*/true, currentType);
 
+      auto resolved = KeyPathExpr::Component::
+        forDictionaryKey(componentName, currentType, componentNameLoc);
+      resolvedComponents.push_back(resolved);
       continue;
     }
 
@@ -321,10 +325,14 @@ Optional<Type> TypeChecker::checkObjCKeyPathExpr(DeclContext *dc,
     if (auto var = dyn_cast<VarDecl>(found)) {
       // Resolve this component to the variable we found.
       auto varRef = ConcreteDeclRef(var);
-      auto resolved =
-        KeyPathExpr::Component::forProperty(varRef, Type(), componentNameLoc);
+      Type varTy = var->getInterfaceType();
+
+      // Updates currentType
+      updateState(/*isProperty=*/true, varTy);
+
+      auto resolved = KeyPathExpr::Component::forProperty(varRef, currentType,
+                                                          componentNameLoc);
       resolvedComponents.push_back(resolved);
-      updateState(/*isProperty=*/true, var->getInterfaceType());
 
       // Check that the property is @objc.
       if (!var->isObjC()) {
@@ -390,7 +398,15 @@ Optional<Type> TypeChecker::checkObjCKeyPathExpr(DeclContext *dc,
         break;
       }
 
+      // Updates currentType based on newType.
       updateState(/*isProperty=*/false, newType);
+
+      // Resolve this component to the type we found.
+      auto typeRef = ConcreteDeclRef(type);
+      auto resolved = KeyPathExpr::Component::forProperty(typeRef, currentType,
+                                                          componentNameLoc);
+      resolvedComponents.push_back(resolved);
+
       continue;
     }
 

test/IDE/complete_pound_keypath.swift

@@ -6,6 +6,16 @@
 
 // RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -code-completion -source-filename %s -code-completion-token=IN_KEYPATH_2 | %FileCheck -check-prefix=CHECK-IN_KEYPATH %s
 
+// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -code-completion -source-filename %s -code-completion-token=IN_KEYPATH_3 | %FileCheck -check-prefix=CHECK-IN_KEYPATH_BRIDGED_STRING %s
+
+// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -code-completion -source-filename %s -code-completion-token=IN_KEYPATH_4 | %FileCheck -check-prefix=CHECK-IN_KEYPATH_BRIDGED_STRING %s
+
+// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -code-completion -source-filename %s -code-completion-token=IN_KEYPATH_5 | %FileCheck -check-prefixes=CHECK-IN_KEYPATH,CHECK-IN_KEYPATH_OPT %s
+
+// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -code-completion -source-filename %s -code-completion-token=IN_KEYPATH_6 | %FileCheck -check-prefixes=CHECK-IN_KEYPATH,CHECK-IN_KEYPATH_OPT %s
+
+// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -code-completion -source-filename %s -code-completion-token=IN_KEYPATH_7 | %FileCheck -check-prefixes=CHECK-IN_KEYPATH_BRIDGED_STRING %s
+
 
 // REQUIRES: objc_interop
 
@@ -21,9 +31,11 @@ func selectorArg1(obj: NSObject) {
   acceptKeyPath(#^KEYPATH_ARG^#
 }
 
-class ObjCClass : NSObject {
+@objcMembers class ObjCClass : NSObject {
   var prop1: String = ""
   var prop2: ObjCClass?
+  var prop3: [ObjCClass]? = []
+  var prop4: [String: String] = [:]
 
   func completeInKeyPath1() {
     _ = #keyPath(#^IN_KEYPATH_1^#
@@ -34,12 +46,42 @@ func completeInKeyPath2() {
   _ = #keyPath(ObjCClass.#^IN_KEYPATH_2^#
 }
 
+func completeInKeyPath3() {
+   _ = #keyPath(ObjCClass.prop1.#^IN_KEYPATH_3^#
+}
+func completeInKeyPath3() {
+     _ = #keyPath(String.#^IN_KEYPATH_4^#
+}
+
+func completeInKeyPath4() {
+  _ = #keyPath(ObjCClass.prop2.#^IN_KEYPATH_5^#
+}
+
+func completeInKeyPath5() {
+  _ = #keyPath(ObjCClass.prop3.#^IN_KEYPATH_6^#
+}
+
+func completeInKeyPath6() {
+  _ = #keyPath(ObjCClass.prop4.anythingHere.#^IN_KEYPATH_7^#
+}
+
 // CHECK-AFTER_POUND-NOT: keyPath
 
 // CHECK-KEYPATH_ARG: Keyword/None/TypeRelation[Identical]: #keyPath({#@objc property sequence#})[#String#]; name=#keyPath(@objc property sequence)
 
 // CHECK-IN_KEYPATH: Decl[InstanceVar]/CurrNominal:      prop1[#String#]; name=prop1
 // CHECK-IN_KEYPATH: Decl[InstanceVar]/CurrNominal:      prop2[#ObjCClass?#]; name=prop2
+// CHECK-IN_KEYPATH: Decl[InstanceVar]/CurrNominal:      prop3[#[ObjCClass]?#]; name=prop3
 // CHECK-IN_KEYPATH: Decl[InstanceVar]/Super:            hashValue[#Int#]; name=hashValue
 
+// Make sure we unwrap optionals (members of Optional itself are invalid in this context)
+//
+// CHECK-IN_KEYPATH_OPT-NOT: name=map
+
+// Make sure we handle bridged types (i.e. show NSString members rather than String members)
+//
+// CHECK-IN_KEYPATH_BRIDGED_STRING: Decl[InstanceVar]/CurrNominal/IsSystem: urlsInText[#[URL]#]; name=urlsInText
+// CHECK-IN_KEYPATH_BRIDGED_STRING: Decl[InstanceVar]/CurrNominal/IsSystem: uppercased[#String!#]; name=uppercased
+// CHECK-IN_KEYPATH_BRIDGED_STRING-NOT: name=count
+