[ConstraintSystem] Detect and diagnoses use of members with mutating getters in key path

lib/Sema/CSApply.cpp

@@ -4566,11 +4566,7 @@ namespace {
         } else {
           // Key paths don't work with mutating-get properties.
           auto varDecl = cast<VarDecl>(property);
-          if (varDecl->isGetterMutating()) {
-            cs.TC.diagnose(componentLoc, diag::expr_keypath_mutating_getter,
-                           property->getFullName());
-          }
-
+          assert(!varDecl->isGetterMutating());
           // Key paths don't currently support static members.
           // There is a fix which diagnoses such situation already.
           assert(!varDecl->isStatic());
@@ -4605,10 +4601,7 @@ namespace {
         SelectedOverload &overload, SourceLoc componentLoc, Expr *indexExpr,
         ArrayRef<Identifier> labels, ConstraintLocator *locator) {
       auto subscript = cast<SubscriptDecl>(overload.choice.getDecl());
-      if (subscript->isGetterMutating()) {
-        cs.TC.diagnose(componentLoc, diag::expr_keypath_mutating_getter,
-                       subscript->getFullName());
-      }
+      assert(!subscript->isGetterMutating());
 
       cs.TC.requestMemberLayout(subscript);
 

lib/Sema/CSDiagnostics.cpp

@@ -2449,9 +2449,8 @@ bool KeyPathSubscriptIndexHashableFailure::diagnoseAsError() {
   return true;
 }
 
-bool InvalidStaticMemberRefInKeyPath::diagnoseAsError() {
+SourceLoc InvalidMemberRefInKeyPath::getLoc() const {
   auto *anchor = getRawAnchor();
-  auto loc = anchor->getLoc();
 
   if (auto *KPE = dyn_cast<KeyPathExpr>(anchor)) {
     auto *locator = getLocator();
@@ -2461,9 +2460,18 @@ bool InvalidStaticMemberRefInKeyPath::diagnoseAsError() {
         });
 
     assert(component != locator->getPath().end());
-    loc = KPE->getComponents()[component->getValue()].getLoc();
+    return KPE->getComponents()[component->getValue()].getLoc();
   }
 
-  emitDiagnostic(loc, diag::expr_keypath_static_member, Member->getBaseName());
+  return anchor->getLoc();
+}
+
+bool InvalidStaticMemberRefInKeyPath::diagnoseAsError() {
+  emitDiagnostic(getLoc(), diag::expr_keypath_static_member, getName());
+  return true;
+}
+
+bool InvalidMemberWithMutatingGetterInKeyPath::diagnoseAsError() {
+  emitDiagnostic(getLoc(), diag::expr_keypath_mutating_getter, getName());
   return true;
 }

lib/Sema/CSDiagnostics.h

@@ -1031,6 +1031,26 @@ class KeyPathSubscriptIndexHashableFailure final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+class InvalidMemberRefInKeyPath : public FailureDiagnostic {
+  ValueDecl *Member;
+
+public:
+  InvalidMemberRefInKeyPath(Expr *root, ConstraintSystem &cs, ValueDecl *member,
+                            ConstraintLocator *locator)
+      : FailureDiagnostic(root, cs, locator), Member(member) {
+    assert(member->hasName());
+    assert(locator->isForKeyPathComponent());
+  }
+
+  DeclName getName() const { return Member->getFullName(); }
+
+  bool diagnoseAsError() override = 0;
+
+protected:
+  /// Compute location of the failure for diagnostic.
+  SourceLoc getLoc() const;
+};
+
 /// Diagnose an attempt to reference a static member as a key path component
 /// e.g.
 ///
@@ -1041,16 +1061,39 @@ class KeyPathSubscriptIndexHashableFailure final : public FailureDiagnostic {
 ///
 /// _ = \S.Type.foo
 /// ```
-class InvalidStaticMemberRefInKeyPath final : public FailureDiagnostic {
-  ValueDecl *Member;
-
+class InvalidStaticMemberRefInKeyPath final : public InvalidMemberRefInKeyPath {
 public:
   InvalidStaticMemberRefInKeyPath(Expr *root, ConstraintSystem &cs,
                                   ValueDecl *member, ConstraintLocator *locator)
-      : FailureDiagnostic(root, cs, locator), Member(member) {
-    assert(member->hasName());
-    assert(locator->isForKeyPathComponent());
-  }
+      : InvalidMemberRefInKeyPath(root, cs, member, locator) {}
+
+  bool diagnoseAsError() override;
+};
+
+/// Diagnose an attempt to reference a member which has a mutating getter as a
+/// key path component e.g.
+///
+/// ```swift
+/// struct S {
+///   var foo: Int {
+///     mutating get { return 42 }
+///   }
+///
+///   subscript(_: Int) -> Bool {
+///     mutating get { return false }
+///   }
+/// }
+///
+/// _ = \S.foo
+/// _ = \S.[42]
+/// ```
+class InvalidMemberWithMutatingGetterInKeyPath final
+    : public InvalidMemberRefInKeyPath {
+public:
+  InvalidMemberWithMutatingGetterInKeyPath(Expr *root, ConstraintSystem &cs,
+                                           ValueDecl *member,
+                                           ConstraintLocator *locator)
+      : InvalidMemberRefInKeyPath(root, cs, member, locator) {}
 
   bool diagnoseAsError() override;
 };

lib/Sema/CSFix.cpp

@@ -423,15 +423,26 @@ TreatKeyPathSubscriptIndexAsHashable::create(ConstraintSystem &cs, Type type,
       TreatKeyPathSubscriptIndexAsHashable(cs, type, locator);
 }
 
-bool AllowStaticMemberRefInKeyPath::diagnose(Expr *root, bool asNote) const {
-  InvalidStaticMemberRefInKeyPath failure(root, getConstraintSystem(), Member,
-                                          getLocator());
-  return failure.diagnose(asNote);
+bool AllowInvalidRefInKeyPath::diagnose(Expr *root, bool asNote) const {
+  switch (Kind) {
+  case RefKind::StaticMember: {
+    InvalidStaticMemberRefInKeyPath failure(root, getConstraintSystem(), Member,
+                                            getLocator());
+    return failure.diagnose(asNote);
+  }
+
+  case RefKind::MutatingGetter: {
+    InvalidMemberWithMutatingGetterInKeyPath failure(
+        root, getConstraintSystem(), Member, getLocator());
+    return failure.diagnose(asNote);
+  }
+  }
 }
 
-AllowStaticMemberRefInKeyPath *
-AllowStaticMemberRefInKeyPath::create(ConstraintSystem &cs, ValueDecl *member,
-                                      ConstraintLocator *locator) {
+AllowInvalidRefInKeyPath *
+AllowInvalidRefInKeyPath::create(ConstraintSystem &cs, RefKind kind,
+                                 ValueDecl *member,
+                                 ConstraintLocator *locator) {
   return new (cs.getAllocator())
-      AllowStaticMemberRefInKeyPath(cs, member, locator);
+      AllowInvalidRefInKeyPath(cs, kind, member, locator);
 }

lib/Sema/CSFix.h

@@ -144,8 +144,9 @@ enum class FixKind : uint8_t {
   /// of the index arguments to be Hashable.
   TreatKeyPathSubscriptIndexAsHashable,
 
-  /// Allow a reference to a static member as a key path component.
-  AllowStaticMemberRefInKeyPath,
+  /// Allow an invalid reference to a member declaration as part
+  /// of a key path component.
+  AllowInvalidRefInKeyPath,
 };
 
 class ConstraintFix {
@@ -783,23 +784,51 @@ class TreatKeyPathSubscriptIndexAsHashable final : public ConstraintFix {
   create(ConstraintSystem &cs, Type type, ConstraintLocator *locator);
 };
 
-class AllowStaticMemberRefInKeyPath final : public ConstraintFix {
+class AllowInvalidRefInKeyPath final : public ConstraintFix {
+  enum RefKind {
+    // Allow a reference to a static member as a key path component.
+    StaticMember,
+    // Allow a reference to a declaration with mutating getter as
+    // a key path component.
+    MutatingGetter,
+  } Kind;
+
   ValueDecl *Member;
 
-  AllowStaticMemberRefInKeyPath(ConstraintSystem &cs, ValueDecl *member,
-                                ConstraintLocator *locator)
-      : ConstraintFix(cs, FixKind::AllowStaticMemberRefInKeyPath, locator),
-        Member(member) {}
+  AllowInvalidRefInKeyPath(ConstraintSystem &cs, RefKind kind,
+                           ValueDecl *member, ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::AllowInvalidRefInKeyPath, locator),
+        Kind(kind), Member(member) {}
 
 public:
   std::string getName() const override {
-    return "allow reference to a static member as a key path component";
+    switch (Kind) {
+    case RefKind::StaticMember:
+      return "allow reference to a static member as a key path component";
+    case RefKind::MutatingGetter:
+      return "allow reference to a member with mutating getter as a key "
+             "path component";
+    }
   }
 
   bool diagnose(Expr *root, bool asNote = false) const override;
 
-  static AllowStaticMemberRefInKeyPath *
-  create(ConstraintSystem &cs, ValueDecl *member, ConstraintLocator *locator);
+  static AllowInvalidRefInKeyPath *forStaticMember(ConstraintSystem &cs,
+                                                   ValueDecl *member,
+                                                   ConstraintLocator *locator) {
+    return create(cs, RefKind::StaticMember, member, locator);
+  }
+
+  static AllowInvalidRefInKeyPath *
+  forMutatingGetter(ConstraintSystem &cs, ValueDecl *member,
+                    ConstraintLocator *locator) {
+    return create(cs, RefKind::MutatingGetter, member, locator);
+  }
+
+private:
+  static AllowInvalidRefInKeyPath *create(ConstraintSystem &cs, RefKind kind,
+                                          ValueDecl *member,
+                                          ConstraintLocator *locator);
 };
 
 } // end namespace constraints

lib/Sema/CSSimplify.cpp

@@ -5022,14 +5022,21 @@ ConstraintSystem::simplifyKeyPathConstraint(Type keyPathTy,
       }
 
       // Referencing static members in key path is not currently allowed.
-      if (storage->isStatic()) {
+      if (storage->isStatic() || storage->isGetterMutating()) {
         if (!shouldAttemptFixes())
           return SolutionKind::Error;
 
-        auto componentLoc =
-            locator.withPathElement(LocatorPathElt::getKeyPathComponent(i));
-        auto *fix = AllowStaticMemberRefInKeyPath::create(
-            *this, choices[i].getDecl(), getConstraintLocator(componentLoc));
+        auto *componentLoc = getConstraintLocator(
+            locator.withPathElement(LocatorPathElt::getKeyPathComponent(i)));
+
+        ConstraintFix *fix = nullptr;
+        if (storage->isStatic()) {
+          fix = AllowInvalidRefInKeyPath::forStaticMember(
+              *this, choices[i].getDecl(), componentLoc);
+        } else {
+          fix = AllowInvalidRefInKeyPath::forMutatingGetter(
+              *this, choices[i].getDecl(), componentLoc);
+        }
 
         if (recordFix(fix))
           return SolutionKind::Error;
@@ -6329,7 +6336,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::AllowInaccessibleMember:
   case FixKind::AllowAnyObjectKeyPathRoot:
   case FixKind::TreatKeyPathSubscriptIndexAsHashable:
-  case FixKind::AllowStaticMemberRefInKeyPath:
+  case FixKind::AllowInvalidRefInKeyPath:
     llvm_unreachable("handled elsewhere");
   }
 

test/expr/unary/keypath/keypath.swift

@@ -732,7 +732,7 @@ func test_keypath_with_static_members(_ p: P_With_Static_Members) {
     static var baz: Int = 42
   }
 
-func foo(_ s: S) {
+  func foo(_ s: S) {
     let _ = \S.Type.foo
     // expected-error@-1 {{key path cannot refer to static member 'foo'}}
     let _ = s[keyPath: \.foo]
@@ -748,6 +748,34 @@ func foo(_ s: S) {
   }
 }
 
+func test_keypath_with_mutating_getter() {
+  struct S {
+    var foo: Int {
+      mutating get { return 42 }
+    }
+
+    subscript(_: Int) -> [Int] {
+      mutating get { return [] }
+    }
+  }
+
+  _ = \S.foo
+  // expected-error@-1 {{key path cannot refer to 'foo', which has a mutating getter}}
+  let _: KeyPath<S, Int> = \.foo
+  // expected-error@-1 {{key path cannot refer to 'foo', which has a mutating getter}}
+  _ = \S.[0]
+  // expected-error@-1 {{key path cannot refer to 'subscript(_:)', which has a mutating getter}}
+  _ = \S.[0].count
+  // expected-error@-1 {{key path cannot refer to 'subscript(_:)', which has a mutating getter}}
+
+  func test_via_subscript(_ s: S) {
+    _ = s[keyPath: \.foo]
+    // expected-error@-1 {{key path cannot refer to 'foo', which has a mutating getter}}
+    _ = s[keyPath: \.[0].count]
+    // expected-error@-1 {{key path cannot refer to 'subscript(_:)', which has a mutating getter}}
+  }
+}
+
 func testSyntaxErrors() { // expected-note{{}}
   _ = \.  ; // expected-error{{expected member name following '.'}}
   _ = \.a ;