addressing feedback on PR

* splitting diagnostic messages for `@derivative` and `@transpose` attr.
* duplicating the diagnostic logic using lambda functions within the
  respective attr checking methods.
* change wording of error/notes to use non-`static`|`static`.

Author: vguerra

include/swift/AST/DiagnosticsSema.def

@@ -3216,6 +3216,17 @@ NOTE(derivative_attr_fix_access,none,
      "mark the derivative function as "
      "'%select{private|fileprivate|internal|@usableFromInline|@usableFromInline}0' "
      "to match the original function", (AccessLevel))
+ERROR(derivative_attr_static_method_mismatch_original,none,
+      "unexpected derivative function declaration; "
+      "%0 requires the derivative function %1 to be %select{an instance|a 'static'}2 method",
+      (/*original*/DeclName, /*derivative*/ DeclName,
+       /*originalIsStatic*/bool))
+NOTE(derivative_attr_static_method_mismatch_original_note,none,
+      "original function %0 is %select{an instance|a 'static'}1 method",
+      (/*original*/ DeclName, /*originalIsStatic*/bool))
+NOTE(derivative_attr_static_method_mismatch_fix,none,
+      "make derivative function %0 %select{an instance|a 'static'}1 method",
+      (/*derivative*/ DeclName, /*mustBeStatic*/bool))
 
 // @transpose
 ERROR(transpose_attr_invalid_linearity_parameter_or_result,none,
@@ -3233,6 +3244,17 @@ ERROR(transpose_attr_wrt_self_must_be_static,none,
 NOTE(transpose_attr_wrt_self_self_type_mismatch_note,none,
      "the transpose is declared in %0 but the original function is declared in "
      "%1", (Type, Type))
+ERROR(transpose_attr_static_method_mismatch_original,none,
+      "unexpected transpose function declaration; "
+      "%0 requires the transpose function %1 to be %select{an instance|a 'static'}2 method",
+      (/*original*/DeclName, /*transpose*/ DeclName,
+       /*originalIsStatic*/bool))
+NOTE(transpose_attr_static_method_mismatch_original_note,none,
+      "original function %0 is %select{an instance|a 'static'}1 method",
+      (/*original*/ DeclName, /*originalIsStatic*/bool))
+NOTE(transpose_attr_static_method_mismatch_fix,none,
+      "make transpose function %0 %select{an instance|a 'static'}1 method",
+      (/*transpose*/ DeclName, /*mustBeStatic*/bool))
 
 // Automatic differentiation attributes
 ERROR(autodiff_attr_original_decl_ambiguous,none,
@@ -3264,18 +3286,6 @@ ERROR(autodiff_attr_result_not_differentiable,none,
       "'Differentiable', but %0 does not conform to 'Differentiable'", (Type))
 ERROR(autodiff_attr_opaque_result_type_unsupported,none,
       "cannot differentiate functions returning opaque result types", ())
-ERROR(autodiff_attr_static_decl_mismatch,none,
-      "derivative function %0 operates on %select{an instance|a}1 type, "
-      "not on %select{an instance|a}2 type as required",
-      (/*derivative*/ DeclName, /*derivativeOperatesOnTy*/bool,
-       /*originalOperatesOnTy*/bool))
-NOTE(autodiff_attr_static_decl_original,none,
-      "original function %0 operates on %select{an instance|a}1 type",
-      (/*original*/ DeclName, /*originalOperatesOnTy*/bool))
-NOTE(autodiff_attr_static_decl_mismatch_fix,none,
-      "derivative function %0 must %select{not be|be}1 'static'",
-      (/*derivative*/ DeclName, /*mustBeStatic*/bool))
-
 
 // differentiation `wrt` parameters clause
 ERROR(diff_function_no_parameters,none,

lib/Sema/TypeCheckAttr.cpp

@@ -34,8 +34,10 @@
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/StorageImpl.h"
+#include "swift/AST/TypeAlignments.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
+#include "swift/Basic/SourceLoc.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/Sema/IDETypeChecking.h"
 #include "clang/Basic/CharInfo.h"
@@ -4583,39 +4585,6 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   (void)attr->getParameterIndices();
 }
 
-/// Checks if original candidate and registered derivative match in terms
-/// of static declaration. If the original candidate is a constructor or is
-/// defined as static method then the registered derivative is expected to be static.
-/// Otherwise the registered derivative should be an instance method.
-/// \returns true if mismatch is found, false otherwise.
-static bool checkStaticDeclMismatch(AbstractFunctionDecl *originalCandidate,
-                                    AbstractFunctionDecl *registered) {
-  return (isa<ConstructorDecl>(originalCandidate) ||
-          originalCandidate->isStatic()) != registered->isStatic();
-}
-
-/// Produces diagnostics for mismatch in static/instance method declaration
-/// between original candidate and registered derivative.
-static void diagnoseStaticDeclMismatch(AbstractFunctionDecl *originalCandidate,
-                                       FuncDecl *registered) {
-  auto &diags = originalCandidate->getASTContext().Diags;
-  diags.diagnose(
-      registered->getNameLoc(), diag::autodiff_attr_static_decl_mismatch,
-      registered->getName(), registered->isStatic(), !registered->isStatic());
-  diags.diagnose(
-      originalCandidate->getNameLoc(), diag::autodiff_attr_static_decl_original,
-      originalCandidate->getName(),
-      isa<ConstructorDecl>(originalCandidate) || originalCandidate->isStatic());
-  auto fixItDiag = diags.diagnose(
-      registered->getStartLoc(), diag::autodiff_attr_static_decl_mismatch_fix,
-      registered->getName(), !registered->isStatic());
-  if (registered->isStatic()) {
-    fixItDiag.fixItRemove(registered->getStaticLoc());
-  } else {
-    fixItDiag.fixItInsert(registered->getStartLoc(), "static ");
-  }
-}
-
 /// Type-checks the given `@derivative` attribute `attr` on declaration `D`.
 ///
 /// Effects are:
@@ -4842,12 +4811,39 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
     }
   }
 
+  attr->setOriginalFunction(originalAFD);
+
+  // Returns true if:
+  // - Original function and derivative function are static methods.
+  // - Original function and derivative function are non-static methods.
+  // - Original function is a Constructor declaration and derivative function is
+  // a static method.
+  auto compatibleStaticDecls = [&]() {
+    return (isa<ConstructorDecl>(originalAFD) || originalAFD->isStatic()) ==
+           derivative->isStatic();
+  };
+
   // Diagnose if original function and derivative differ in terms of static declaration.
-  if (checkStaticDeclMismatch(originalAFD, derivative)) {
-    diagnoseStaticDeclMismatch(originalAFD, derivative);
+  if (!compatibleStaticDecls()) {
+    bool derivativeMustBeStatic = !derivative->isStatic();
+    diags.diagnose(attr->getOriginalFunctionName().Loc.getBaseNameLoc(),
+                   diag::derivative_attr_static_method_mismatch_original,
+                   originalAFD->getName(), derivative->getName(),
+                   derivativeMustBeStatic);
+    diags.diagnose(originalAFD->getNameLoc(),
+                   diag::derivative_attr_static_method_mismatch_original_note,
+                   originalAFD->getName(), derivativeMustBeStatic);
+    auto fixItDiag =
+        diags.diagnose(derivative->getStartLoc(),
+                       diag::derivative_attr_static_method_mismatch_fix,
+                       derivative->getName(), derivativeMustBeStatic);
+    if (derivativeMustBeStatic) {
+      fixItDiag.fixItInsert(derivative->getStartLoc(), "static ");
+    } else {
+      fixItDiag.fixItRemove(derivative->getStaticLoc());
+    }
     return true;
   }
-  attr->setOriginalFunction(originalAFD);
 
   // Returns true if:
   // - Original function and derivative function have the same access level.
@@ -5364,11 +5360,35 @@ void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
     return;
   }
 
+  // Returns true if:
+  // - Original function and transpose function are static methods.
+  // - Original function and transpose function are non-static methods.
+  // - Original function is a Constructor declaration and transpose function is
+  // a static method.
+  auto compatibleStaticDecls = [&]() {
+    return (isa<ConstructorDecl>(originalAFD) || originalAFD->isStatic()) ==
+           transpose->isStatic();
+  };
+
   // Diagnose if original function and transpose differ in terms of static declaration.
-  if (!doSelfTypesMatch && checkStaticDeclMismatch(originalAFD, transpose)) {
-      diagnoseStaticDeclMismatch(originalAFD, transpose);
-      attr->setInvalid();
-      return;
+  if (!doSelfTypesMatch && !compatibleStaticDecls()) {
+    bool transposeMustBeStatic = !transpose->isStatic();
+    diagnose(attr->getOriginalFunctionName().Loc.getBaseNameLoc(),
+             diag::transpose_attr_static_method_mismatch_original,
+             originalAFD->getName(), transpose->getName(),
+             transposeMustBeStatic);
+    diagnose(originalAFD->getNameLoc(),
+             diag::transpose_attr_static_method_mismatch_original_note,
+             originalAFD->getName(), transposeMustBeStatic);
+    auto fixItDiag = diagnose(transpose->getStartLoc(),
+                              diag::transpose_attr_static_method_mismatch_fix,
+                              transpose->getName(), transposeMustBeStatic);
+    if (transposeMustBeStatic) {
+      fixItDiag.fixItInsert(transpose->getStartLoc(), "static ");
+    } else {
+      fixItDiag.fixItRemove(transpose->getStaticLoc());
+    }
+    return;
   }
 
   // Set the resolved linearity parameter indices in the attribute.

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -1154,32 +1154,32 @@ func vjpOpaqueResult(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
 // Test instance vs static method mismatch.
 
 struct StaticMismatch<T: Differentiable> {
-  // expected-note @+1 {{original function 'init(_:)' operates on a type}}
+  // expected-note @+1 {{original function 'init(_:)' is a 'static' method}}
   init(_ x: T) {}
-  // expected-note @+1 {{original function 'instanceMethod' operates on an instance type}}
+  // expected-note @+1 {{original function 'instanceMethod' is an instance method}}
   func instanceMethod(_ x: T) -> T { x }
-  // expected-note @+1 {{original function 'staticMethod' operates on a type}}
+  // expected-note @+1 {{original function 'staticMethod' is a 'static' method}}
   static func staticMethod(_ x: T) -> T { x }
 
+  // expected-error @+1 {{unexpected derivative function declaration; 'init(_:)' requires the derivative function 'vjpInit' to be a 'static' method}}
   @derivative(of: init)
-  // expected-error @+2 {{derivative function 'vjpInit' operates on an instance type, not on a type as required}}
-  // expected-note @+1 {{derivative function 'vjpInit' must be 'static'}}{{3-3=static }}
+  // expected-note @+1 {{make derivative function 'vjpInit' a 'static' method}}{{3-3=static }}
   func vjpInit(_ x: T) -> (value: Self, pullback: (T.TangentVector) -> T.TangentVector) {
     fatalError()
   }
 
+  // expected-error @+1 {{unexpected derivative function declaration; 'instanceMethod' requires the derivative function 'jvpInstance' to be an instance method}}
   @derivative(of: instanceMethod)
-  // expected-error @+2 {{derivative function 'jvpInstance' operates on a type, not on an instance type as required}}
-  // expected-note @+1 {{derivative function 'jvpInstance' must not be 'static'}}{{3-10=}}
+  // expected-note @+1 {{make derivative function 'jvpInstance' an instance method}}{{3-10=}}
   static func jvpInstance(_ x: T) -> (
     value: T, differential: (T.TangentVector) -> (T.TangentVector)
   ) {
     return (x, { $0 })
   }
 
+  // expected-error @+1 {{unexpected derivative function declaration; 'staticMethod' requires the derivative function 'jvpStatic' to be a 'static' method}}
   @derivative(of: staticMethod)
-  // expected-error @+2 {{derivative function 'jvpStatic' operates on an instance type, not on a type as required}}
-  // expected-note @+1 {{derivative function 'jvpStatic' must be 'static'}}{{3-3=static }}
+  // expected-note @+1 {{make derivative function 'jvpStatic' a 'static' method}}{{3-3=static }}
   func jvpStatic(_ x: T) -> (
     value: T, differential: (T.TangentVector) -> (T.TangentVector)
   ) {

test/AutoDiff/Sema/transpose_attr_type_checking.swift

@@ -498,7 +498,7 @@ extension Struct where T: Differentiable & AdditiveArithmetic {
 
 // Test initializers.
 extension Struct {
-  // expected-note @+1 {{original function 'init(_:)' operates on a type}}
+  // expected-note @+1 {{original function 'init(_:)' is a 'static' method}}
   init(_ x: Float) {}
   init(_ x: T, y: Float) {}
 }
@@ -515,9 +515,9 @@ extension Struct where T: Differentiable, T == T.TangentVector {
   }
 
   // Test instance transpose for static original initializer.
+  // expected-error @+1 {{unexpected transpose function declaration; 'init(_:)' requires the transpose function 'vjpInitStaticMismatch' to be a 'static' method}}
   @transpose(of: init, wrt: 0)
-  // expected-error @+2 {{derivative function 'vjpInitStaticMismatch' operates on an instance type, not on a type as required}}
-  // expected-note @+1 {{derivative function 'vjpInitStaticMismatch' must be 'static'}}{{3-3=static }}
+  // expected-note @+1 {{make transpose function 'vjpInitStaticMismatch' a 'static' method}}{{3-3=static }}
   func vjpInitStaticMismatch(_ x: Self) -> Float {
     fatalError()
   }