[AutoDiff][SR-13152] Better diagnostic for static decl modifier mismatch. (#36128)

Improved diagnostic for when the registered derivative function and
the function it derivates (the original) differ in terms of
`static` declaration modifier usage. Suggesting as well a fix-it,
to either remove or add the `static` keyword.

The registered derivative needs to be marked as `static` in two cases:

1. When the original function is a constructor.
2. When the original function is static as well.

When the original function is an instance method, the registered derivative must be as well.

Resolves [SR-13152](https://bugs.swift.org/browse/SR-13152).

Author: vguerra

include/swift/AST/DiagnosticsSema.def

@@ -3232,6 +3232,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,
@@ -3249,6 +3260,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,

lib/Sema/TypeCheckAttr.cpp

@@ -4028,7 +4028,9 @@ static bool checkFunctionSignature(
   if (!std::equal(required->getParams().begin(), required->getParams().end(),
                   candidateFnTy->getParams().begin(),
                   [&](AnyFunctionType::Param x, AnyFunctionType::Param y) {
-                    return x.getOldType()->isEqual(mapType(y.getOldType()));
+                    auto xInstanceTy = x.getOldType()->getMetatypeInstanceType();
+                    auto yInstanceTy = y.getOldType()->getMetatypeInstanceType();
+                    return xInstanceTy->isEqual(mapType(yInstanceTy));
                   }))
     return false;
 
@@ -4827,8 +4829,41 @@ 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 (!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;
+  }
+
   // Returns true if:
   // - Original function and derivative function have the same access level.
   // - Original function is public and derivative function is internal
@@ -5231,8 +5266,9 @@ void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
   // If the transpose function is curried and `self` is a linearity parameter,
   // check that the instance and static `Self` types are equal.
   Type staticSelfType, instanceSelfType;
+  bool doSelfTypesMatch = false;
   if (isCurried && wrtSelf) {
-    bool doSelfTypesMatch = doTransposeStaticAndInstanceSelfTypesMatch(
+    doSelfTypesMatch = doTransposeStaticAndInstanceSelfTypesMatch(
         transposeInterfaceType, staticSelfType, instanceSelfType);
     if (!doSelfTypesMatch) {
       diagnose(attr->getLocation(),
@@ -5343,6 +5379,37 @@ 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 && !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.
   attr->setParameterIndices(linearParamIndices);
 }

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -669,21 +669,6 @@ func jvpInvalid<T: Differentiable>(x: T) -> (
   return (x, { $0 })
 }
 
-// Test invalid derivative type context: instance vs static method mismatch.
-
-struct InvalidTypeContext<T: Differentiable> {
-  // expected-note @+1 {{candidate static method does not have type equal to or less constrained than '<T where T : Differentiable> (InvalidTypeContext<T>) -> (T) -> T'}}
-  static func staticMethod(_ x: T) -> T { x }
-
-  // expected-error @+1 {{referenced declaration 'staticMethod' could not be resolved}}
-  @derivative(of: staticMethod)
-  func jvpStatic(_ x: T) -> (
-    value: T, differential: (T.TangentVector) -> (T.TangentVector)
-  ) {
-    return (x, { $0 })
-  }
-}
-
 // Test stored property original declaration.
 
 struct HasStoredProperty {
@@ -1165,3 +1150,41 @@ func opaqueResult(_ x: Float) -> some Differentiable { x }
 func vjpOpaqueResult(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
   fatalError()
 }
+
+// Test instance vs static method mismatch.
+
+struct StaticMismatch<T: Differentiable> {
+  // expected-note @+1 {{original function 'init(_:)' is a 'static' method}}
+  init(_ x: T) {}
+  // expected-note @+1 {{original function 'instanceMethod' is an instance method}}
+  func instanceMethod(_ x: T) -> T { x }
+  // 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-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-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-note @+1 {{make derivative function 'jvpStatic' a 'static' method}}{{3-3=static }}
+  func jvpStatic(_ x: T) -> (
+    value: T, differential: (T.TangentVector) -> (T.TangentVector)
+  ) {
+    return (x, { $0 })
+  }
+}
+
+

test/AutoDiff/Sema/transpose_attr_type_checking.swift

@@ -474,7 +474,6 @@ extension Float {
 
 // Test non-`func` original declarations.
 
-// expected-note @+1 {{candidate initializer does not have type equal to or less constrained than '<T where T : Differentiable, T == T.TangentVector> (Struct<T>) -> (Float) -> Struct<T>'}}
 struct Struct<T> {}
 extension Struct: Equatable where T: Equatable {}
 extension Struct: Differentiable & AdditiveArithmetic
@@ -499,9 +498,8 @@ extension Struct where T: Differentiable & AdditiveArithmetic {
 
 // Test initializers.
 extension Struct {
-  // expected-note @+1 {{candidate initializer does not have type equal to or less constrained than '<T where T : Differentiable, T == T.TangentVector> (Struct<T>) -> (Float) -> Struct<T>'}}
+  // expected-note @+1 {{original function 'init(_:)' is a 'static' method}}
   init(_ x: Float) {}
-  // expected-note @+1 {{candidate initializer does not have type equal to or less constrained than '<T where T : Differentiable, T == T.TangentVector> (Struct<T>) -> (Float) -> Struct<T>'}}
   init(_ x: T, y: Float) {}
 }
 
@@ -517,9 +515,9 @@ extension Struct where T: Differentiable, T == T.TangentVector {
   }
 
   // Test instance transpose for static original initializer.
-  // TODO(TF-1015): Add improved instance/static member mismatch error.
-  // expected-error @+1 {{referenced declaration 'init' could not be resolved}}
+  // expected-error @+1 {{unexpected transpose function declaration; 'init(_:)' requires the transpose function 'vjpInitStaticMismatch' to be a 'static' method}}
   @transpose(of: init, wrt: 0)
+  // expected-note @+1 {{make transpose function 'vjpInitStaticMismatch' a 'static' method}}{{3-3=static }}
   func vjpInitStaticMismatch(_ x: Self) -> Float {
     fatalError()
   }