[TypeChecker] Add ReferenceWritableKeyPath support to keypath dynamic member lookup

Author: xedin

lib/Sema/CSDiag.cpp

@@ -812,6 +812,7 @@ void FailureDiagnosis::diagnoseUnviableLookupResults(
     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/CSSimplify.cpp

@@ -3853,6 +3853,17 @@ performMemberLookup(ConstraintKind constraintKind, DeclName memberName,
               MemberLookupResult::UR_WritableKeyPathOnReadOnlyMember);
           return;
         }
+
+        // A nonmutating setter indicates a reference-writable base,
+        // on the other hand if setter is mutating there is no point
+        // of attempting `ReferenceWritableKeyPath` overload.
+        if (storage->isSetterMutating() &&
+            keyPath == getASTContext().getReferenceWritableKeyPathDecl()) {
+          result.addUnviable(
+              candidate,
+              MemberLookupResult::UR_ReferenceWritableKeyPathOnMutatingMember);
+          return;
+        }
       }
     }
 
@@ -4245,10 +4256,11 @@ fixMemberRef(ConstraintSystem &cs, Type baseTy,
     case MemberLookupResult::UR_LabelMismatch:
     case MemberLookupResult::UR_UnavailableInExistential:
     // TODO(diagnostics): Add a new fix that is suggests to
-    // add `subscript(dynamicMember: KeyPath<T, U>)`
+    // add `subscript(dynamicMember: {Writable}KeyPath<T, U>)`
     // overload here, that would help if such subscript has
     // not been provided.
     case MemberLookupResult::UR_WritableKeyPathOnReadOnlyMember:
+    case MemberLookupResult::UR_ReferenceWritableKeyPathOnMutatingMember:
       break;
     }
   }

lib/Sema/ConstraintSystem.cpp

@@ -1894,10 +1894,12 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
       refType = fnType->getResult();
 
       auto *keyPathDecl = keyPathTy->getAnyNominal();
-      assert(keyPathDecl &&
-             (keyPathDecl == getASTContext().getKeyPathDecl() ||
-              keyPathDecl == getASTContext().getWritableKeyPathDecl()) &&
-             "parameter is supposed to be a keypath");
+      assert(
+          keyPathDecl &&
+          (keyPathDecl == getASTContext().getKeyPathDecl() ||
+           keyPathDecl == getASTContext().getWritableKeyPathDecl() ||
+           keyPathDecl == getASTContext().getReferenceWritableKeyPathDecl()) &&
+          "parameter is supposed to be a keypath");
 
       auto *keyPathLoc = getConstraintLocator(
           locator, LocatorPathElt::getKeyPathDynamicMember(keyPathDecl));
@@ -1965,7 +1967,8 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
         // form additional `applicable fn` constraint here and bind it to a
         // function type, but it would create inconsistency with how properties
         // are handled, which means more special handling in CSApply.
-        if (keyPathDecl == getASTContext().getWritableKeyPathDecl())
+        if (keyPathDecl == getASTContext().getWritableKeyPathDecl() ||
+            keyPathDecl == getASTContext().getReferenceWritableKeyPathDecl())
           dynamicResultTy->getImpl().setCanBindToLValue(getSavedBindings(),
                                                         /*enabled=*/true);
 

lib/Sema/ConstraintSystem.h

@@ -882,6 +882,11 @@ struct MemberLookupResult {
     /// because it's not known upfront what access capability would the
     /// member have.
     UR_WritableKeyPathOnReadOnlyMember,
+    /// This is a `ReferenceWritableKeyPath` being used to look up mutating
+    /// member, used in situations involving dynamic member lookup via keypath,
+    /// because it's not known upfront what access capability would the
+    /// member have.
+    UR_ReferenceWritableKeyPathOnMutatingMember,
   };
 
   /// This is a list of considered (but rejected) candidates, along with a

lib/Sema/TypeCheckAttr.cpp

@@ -1026,7 +1026,8 @@ bool swift::isValidKeyPathDynamicMemberLookup(SubscriptDecl *decl,
   const auto *param = decl->getIndices()->get(0);
   if (auto NTD = param->getType()->getAnyNominal()) {
     return NTD == TC.Context.getKeyPathDecl() ||
-           NTD == TC.Context.getWritableKeyPathDecl();
+           NTD == TC.Context.getWritableKeyPathDecl() ||
+           NTD == TC.Context.getReferenceWritableKeyPathDecl();
   }
   return false;
 }

test/Constraints/keypath_dynamic_member_lookup.swift

@@ -137,3 +137,45 @@ func keypath_with_subscripts(_ arr: SubscriptLens<[Int]>,
   // CHECK: keypath $WritableKeyPath<Dictionary<String, Int>, Optional<Int>>, (root $Dictionary<String, Int>; settable_property $Optional<Int>,  id @$sSDyq_Sgxcig : {{.*}})
   dict["ultimate question"] = 42
 }
+
+struct DotStruct {
+  var x, y: Int
+}
+
+class DotClass {
+  var x, y: Int
+
+  init(x: Int, y: Int) {
+    self.x = x
+    self.y = y
+  }
+}
+
+@dynamicMemberLookup
+struct DotLens<T> {
+  var value: T
+
+  subscript<U>(dynamicMember member: WritableKeyPath<T, U>) -> U {
+    get { return value[keyPath: member] }
+    set { value[keyPath: member] = newValue }
+  }
+
+  subscript<U>(dynamicMember member: ReferenceWritableKeyPath<T, U>) -> U {
+    get { return value[keyPath: member] }
+    set { value[keyPath: member] = newValue }
+  }
+}
+
+func dot_struct_test(_ lens: inout DotLens<DotStruct>) {
+  // CHECK: keypath $WritableKeyPath<DotStruct, Int>, (root $DotStruct; stored_property #DotStruct.x : $Int)
+  lens.x = 1
+  // CHECK: keypath $WritableKeyPath<DotStruct, Int>, (root $DotStruct; stored_property #DotStruct.y : $Int)
+  let _ = lens.y
+}
+
+func dot_class_test(_ lens: inout DotLens<DotClass>) {
+  // CHECK: keypath $ReferenceWritableKeyPath<DotClass, Int>, (root $DotClass; settable_property $Int,  id #DotClass.x!getter.1 : (DotClass) -> () -> Int, getter @$s29keypath_dynamic_member_lookup8DotClassC1xSivpACTK : {{.*}})
+  lens.x = 1
+  // CHECK: keypath $ReferenceWritableKeyPath<DotClass, Int>, (root $DotClass; settable_property $Int,  id #DotClass.y!getter.1 : (DotClass) -> () -> Int, getter @$s29keypath_dynamic_member_lookup8DotClassC1ySivpACTK : {{.*}})
+  let _ = lens.y
+}