[SE-0252][TypeChecker] Keypath Dynamic Member Lookup

include/swift/AST/DiagnosticsSema.def

@@ -1057,8 +1057,8 @@ ERROR(missing_dynamic_callable_kwargs_method,none,
 
 ERROR(invalid_dynamic_member_lookup_type,none,
       "@dynamicMemberLookup attribute requires %0 to have a "
-      "'subscript(dynamicMember:)' method with an "
-      "'ExpressibleByStringLiteral' parameter", (Type))
+      "'subscript(dynamicMember:)' method that accepts either "
+      "'ExpressibleByStringLiteral' or a keypath", (Type))
 
 ERROR(string_index_not_integer,none,
       "String must not be indexed with %0, it has variable size elements",

lib/Sema/CSDiag.cpp

@@ -811,6 +811,8 @@ void FailureDiagnosis::diagnoseUnviableLookupResults(
   if (sameProblem) {
     switch (firstProblem) {
     case MemberLookupResult::UR_LabelMismatch:
+    case MemberLookupResult::UR_WritableKeyPathOnReadOnlyMember:
+    case MemberLookupResult::UR_ReferenceWritableKeyPathOnMutatingMember:
       break;
     case MemberLookupResult::UR_UnavailableInExistential:
       diagnose(loc, diag::could_not_use_member_on_existential,

lib/Sema/CSDiagnostics.cpp

@@ -776,7 +776,7 @@ bool MissingOptionalUnwrapFailure::diagnoseAsError() {
 bool RValueTreatedAsLValueFailure::diagnoseAsError() {
   Diag<StringRef> subElementDiagID;
   Diag<Type> rvalueDiagID = diag::assignment_lhs_not_lvalue;
-  Expr *diagExpr = getLocator()->getAnchor();
+  Expr *diagExpr = getRawAnchor();
   SourceLoc loc = diagExpr->getLoc();
 
   if (auto assignExpr = dyn_cast<AssignExpr>(diagExpr)) {
@@ -854,10 +854,18 @@ bool RValueTreatedAsLValueFailure::diagnoseAsError() {
       }
     }
 
-    if (auto resolvedOverload = getResolvedOverload(getLocator()))
+    if (auto resolvedOverload = getResolvedOverload(getLocator())) {
       if (resolvedOverload->Choice.getKind() ==
           OverloadChoiceKind::DynamicMemberLookup)
         subElementDiagID = diag::assignment_dynamic_property_has_immutable_base;
+
+      if (resolvedOverload->Choice.getKind() ==
+          OverloadChoiceKind::KeyPathDynamicMemberLookup) {
+        if (!getType(member->getBase())->hasLValueType())
+          subElementDiagID =
+              diag::assignment_dynamic_property_has_immutable_base;
+      }
+    }
   } else if (auto sub = dyn_cast<SubscriptExpr>(diagExpr)) {
       subElementDiagID = diag::assignment_subscript_has_immutable_base;
   } else {
@@ -1202,7 +1210,7 @@ AssignmentFailure::resolveImmutableBase(Expr *expr) const {
     if (!member) {
       auto loc =
           cs.getConstraintLocator(SE, ConstraintLocator::SubscriptMember);
-      member = dyn_cast_or_null<SubscriptDecl>(cs.findResolvedMemberRef(loc));
+      member = dyn_cast_or_null<SubscriptDecl>(getMemberRef(loc));
     }
 
     // If it isn't settable, return it.
@@ -1231,9 +1239,10 @@ AssignmentFailure::resolveImmutableBase(Expr *expr) const {
     // If we found a decl for the UDE, check it.
     auto loc = cs.getConstraintLocator(UDE, ConstraintLocator::Member);
 
+    auto *member = getMemberRef(loc);
     // If we can resolve a member, we can determine whether it is settable in
     // this context.
-    if (auto *member = cs.findResolvedMemberRef(loc)) {
+    if (member) {
       auto *memberVD = dyn_cast<VarDecl>(member);
 
       // If the member isn't a vardecl (e.g. its a funcdecl), or it isn't
@@ -1281,6 +1290,43 @@ AssignmentFailure::resolveImmutableBase(Expr *expr) const {
   return {expr, nullptr};
 }
 
+ValueDecl *AssignmentFailure::getMemberRef(ConstraintLocator *locator) const {
+  auto member = getOverloadChoiceIfAvailable(locator);
+  if (!member || !member->choice.isDecl())
+    return nullptr;
+
+  auto *DC = getDC();
+  auto &TC = getTypeChecker();
+
+  auto *decl = member->choice.getDecl();
+  if (isa<SubscriptDecl>(decl) &&
+      isValidDynamicMemberLookupSubscript(cast<SubscriptDecl>(decl), DC, TC)) {
+    auto *subscript = cast<SubscriptDecl>(decl);
+    // If this is a keypath dynamic member lookup, we have to
+    // adjust the locator to find member referred by it.
+    if (isValidKeyPathDynamicMemberLookup(subscript, TC)) {
+      auto &cs = getConstraintSystem();
+      // Type has a following format:
+      // `(Self) -> (dynamicMember: {Writable}KeyPath<T, U>) -> U`
+      auto *fullType = member->openedFullType->castTo<FunctionType>();
+      auto *fnType = fullType->getResult()->castTo<FunctionType>();
+
+      auto paramTy = fnType->getParams()[0].getPlainType();
+      auto keyPath = paramTy->getAnyNominal();
+      auto memberLoc = cs.getConstraintLocator(
+          locator, LocatorPathElt::getKeyPathDynamicMember(keyPath));
+
+      auto memberRef = getOverloadChoiceIfAvailable(memberLoc);
+      return memberRef ? memberRef->choice.getDecl() : nullptr;
+    }
+
+    // If this is a string based dynamic lookup, there is no member declaration.
+    return nullptr;
+  }
+
+  return decl;
+}
+
 Diag<StringRef> AssignmentFailure::findDeclDiagonstic(ASTContext &ctx,
                                                       Expr *destExpr) {
   if (isa<ApplyExpr>(destExpr) || isa<SelfApplyExpr>(destExpr))

lib/Sema/CSDiagnostics.h

@@ -647,6 +647,10 @@ class AssignmentFailure final : public FailureDiagnostic {
              isLoadedLValue(ifExpr->getElseExpr());
     return false;
   }
+
+  /// Retrive an member reference associated with given member
+  /// looking through dynamic member lookup on the way.
+  ValueDecl *getMemberRef(ConstraintLocator *locator) const;
 };
 
 /// Intended to diagnose any possible contextual failure

lib/Sema/CSGen.cpp

@@ -3274,6 +3274,38 @@ namespace {
             expr = sanitizeArgumentList(expr);
         }
 
+        // If this expression represents keypath based dynamic member
+        // lookup, let's convert it back to the original form of
+        // member or subscript reference.
+        if (auto *SE = dyn_cast<SubscriptExpr>(expr)) {
+          if (auto *TE = dyn_cast<TupleExpr>(SE->getIndex())) {
+            auto isImplicitKeyPathExpr = [](Expr *argExpr) -> bool {
+              if (auto *KP = dyn_cast<KeyPathExpr>(argExpr))
+                return KP->isImplicit();
+              return false;
+            };
+
+            if (TE->isImplicit() && TE->getNumElements() == 1 &&
+                TE->getElementName(0) == TC.Context.Id_dynamicMember &&
+                isImplicitKeyPathExpr(TE->getElement(0))) {
+              auto *keyPathExpr = cast<KeyPathExpr>(TE->getElement(0));
+              auto *componentExpr = keyPathExpr->getParsedPath();
+
+              if (auto *UDE = dyn_cast<UnresolvedDotExpr>(componentExpr)) {
+                UDE->setBase(SE->getBase());
+                return {true, UDE};
+              }
+
+              if (auto *subscript = dyn_cast<SubscriptExpr>(componentExpr)) {
+                subscript->setBase(SE->getBase());
+                return {true, subscript};
+              }
+
+              llvm_unreachable("unknown keypath component type");
+            }
+          }
+        }
+
         // Now, we're ready to walk into sub expressions.
         return {true, expr};
       }
@@ -3898,4 +3930,4 @@ swift::getOriginalArgumentList(Expr *expr) {
   }
 
   return result;
-}
+}

lib/Sema/CSRanking.cpp

@@ -151,6 +151,7 @@ static bool sameOverloadChoice(const OverloadChoice &x,
   case OverloadChoiceKind::DeclViaBridge:
   case OverloadChoiceKind::DeclViaUnwrappedOptional:
   case OverloadChoiceKind::DynamicMemberLookup:
+  case OverloadChoiceKind::KeyPathDynamicMemberLookup:
     return sameDecl(x.getDecl(), y.getDecl());
 
   case OverloadChoiceKind::TupleIndex:
@@ -875,6 +876,20 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
         continue;
       }
 
+      // Dynamic member lookup through a keypath is better than one using string
+      // because it carries more type information.
+      if (choice1.getKind() == OverloadChoiceKind::KeyPathDynamicMemberLookup &&
+          choice2.getKind() == OverloadChoiceKind::DynamicMemberLookup) {
+        score1 += weight;
+        continue;
+      }
+
+      if (choice1.getKind() == OverloadChoiceKind::DynamicMemberLookup &&
+          choice2.getKind() == OverloadChoiceKind::KeyPathDynamicMemberLookup) {
+        score2 += weight;
+        continue;
+      }
+
       continue;
     }
 
@@ -893,6 +908,7 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
     case OverloadChoiceKind::DeclViaBridge:
     case OverloadChoiceKind::DeclViaUnwrappedOptional:
     case OverloadChoiceKind::DynamicMemberLookup:
+    case OverloadChoiceKind::KeyPathDynamicMemberLookup:
       break;
     }