Augment maybeDiagnoseCallToKeyValueObserveMethod() and add support for a @_semantics attribute to trigger it rather than just hard-coding the name of the Foundation method.

Author: grynspan

include/swift/AST/SemanticAttrs.def

@@ -97,6 +97,7 @@ SEMANTICS_ATTR(PROGRAMTERMINATION_POINT, "programtermination_point")
 SEMANTICS_ATTR(CONVERT_TO_OBJECTIVE_C, "convertToObjectiveC")
 
 SEMANTICS_ATTR(KEYPATH_KVC_KEY_PATH_STRING, "keypath.kvcKeyPathString")
+SEMANTICS_ATTR(KEYPATH_MUST_BE_VALID_FOR_KVO, "keypath.mustBeValidForKVO")
 
 /// The prefix used to force opt-remarks to be emitted in a specific function.
 ///

lib/Sema/MiscDiagnostics.cpp

@@ -25,6 +25,7 @@
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/Pattern.h"
+#include "swift/AST/SemanticAttrs.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeCheckRequests.h"
@@ -5053,31 +5054,40 @@ static void maybeDiagnoseCallToKeyValueObserveMethod(const Expr *E,
       auto fn = expr->getCalledValue();
       if (!fn)
         return;
-      if (fn->getModuleContext()->getName() != C.Id_Foundation)
-        return;
-      if (!fn->getName().isCompoundName("observe",
-                                        {"", "options", "changeHandler"}))
-        return;
+      SmallVector<KeyPathExpr *, 1> keyPathArgs;
       auto *args = expr->getArgs();
-      auto firstArg = dyn_cast<KeyPathExpr>(args->getExpr(0));
-      if (!firstArg)
-        return;
-      auto lastComponent = firstArg->getComponents().back();
-      if (lastComponent.getKind() != KeyPathExpr::Component::Kind::Property)
-        return;
-      auto property = lastComponent.getDeclRef().getDecl();
-      if (!property)
-        return;
-      auto propertyVar = cast<VarDecl>(property);
-      if (propertyVar->shouldUseObjCDispatch() ||
-          (propertyVar->isObjC() &&
-           propertyVar->getParsedAccessor(AccessorKind::Set)))
-        return;
-      C.Diags
+      if (fn->getModuleContext()->getName() == C.Id_Foundation &&
+          fn->getName().isCompoundName("observe",
+                                       {"", "options", "changeHandler"})) {
+        if (auto keyPathArg = dyn_cast<KeyPathExpr>(args->getExpr(0))) {
+          keyPathArgs.push_back(keyPathArg);
+        }
+      } else if (fn->getAttrs()
+                 .hasSemanticsAttr(semantics::KEYPATH_MUST_BE_VALID_FOR_KVO)) {
+        for (const auto& arg: *args) {
+          if (auto keyPathArg = dyn_cast<KeyPathExpr>(arg.getExpr())) {
+            keyPathArgs.push_back(keyPathArg);
+          }
+        }
+      }
+      for (auto *keyPathArg : keyPathArgs) {
+        auto lastComponent = keyPathArg->getComponents().back();
+        if (lastComponent.getKind() != KeyPathExpr::Component::Kind::Property)
+          continue;
+        auto property = lastComponent.getDeclRef().getDecl();
+        if (!property)
+          continue;
+        auto propertyVar = cast<VarDecl>(property);
+        if (propertyVar->shouldUseObjCDispatch() ||
+            (propertyVar->isObjC() &&
+             propertyVar->getParsedAccessor(AccessorKind::Set)))
+          continue;
+        C.Diags
           .diagnose(expr->getLoc(),
                     diag::observe_keypath_property_not_objc_dynamic,
                     property->getName(), fn->getName())
           .highlight(lastComponent.getLoc());
+      }
     }
 
     PreWalkResult<Expr *> walkToExprPre(Expr *E) override {

test/expr/primary/keypath/keypath-observe-objc.swift

@@ -45,3 +45,11 @@ class Bar: NSObject {
     }
   }
 }
+
+@_semantics("keypath.mustBeValidForKVO")
+func quux<T, V, U>(_ object: T, at keyPath: KeyPath<T, V>, _ keyPath2: KeyPath<T, U>) { }
+
+// The presence of a valid keypath should not prevent detection of invalid ones
+// later in the argument list, so start with a valid one here.
+quux(Foo(), at: \.number4, \.number1)
+// expected-warning@-1 {{passing reference to non-'@objc dynamic' property 'number1' to KVO method 'quux(_:at:_:)' may lead to unexpected behavior or runtime trap}}