[AutoDiff] Disallow JVP/VJP in protocol requirement @differentiable attribute. (#26113)

Protocol requirements cannot specify `jvp:` or `vjp:` in `@differentiable`
attribute. This is because all witnesses must specify a `@differentiable`
attribute with the same configuration, so all witnesses will have their own
JVP/VJP (either specified or generated).

Reorganize tests.

Resolves TF-651.
Exposes TF-650: diagnose protocol default implementations missing
`@differentiable` attribute.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -2784,8 +2784,12 @@ ERROR(differentiable_attr_invalid_access,none,
 ERROR(differentiable_attr_result_not_differentiable,none,
       "can only differentiate functions with results that conform to "
       "'Differentiable', but %0 does not conform to 'Differentiable'", (Type))
-ERROR(differentiable_attr_protocol_where_clause,none,
-      "'where' clauses cannot be used in a '@differentiable' attribute on a protocol requirement", ())
+ERROR(differentiable_attr_protocol_req_where_clause,none,
+      "'@differentiable' attribute on protocol requirement cannot specify "
+      "'where' clause", ())
+ERROR(differentiable_attr_protocol_req_assoc_func,none,
+      "'@differentiable' attribute on protocol requirement cannot specify "
+      "'jvp:' or 'vjp:'", ())
 ERROR(differentiable_attr_empty_where_clause,none,
       "empty 'where' clause in '@differentiable' attribute", ())
 ERROR(differentiable_attr_nongeneric_trailing_where,none,

lib/Sema/TypeCheckAttr.cpp

@@ -3267,6 +3267,12 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   // Start type-checking the arguments of the @differentiable attribute. This
   // covers 'wrt:', 'jvp:', 'vjp:', and 'where', all of which are optional.
 
+  // `@differentiable` attributes on protocol requirements do not support
+  // JVP/VJP or 'where' clauses.
+  bool isOriginalProtocolRequirement =
+      isa<ProtocolDecl>(original->getDeclContext()) &&
+      original->isProtocolRequirement();
+
   // Handle 'where' clause, if it exists.
   // - Resolve attribute where clause requirements and store in the attribute
   //   for serialization.
@@ -3276,12 +3282,11 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   GenericSignature *whereClauseGenSig = nullptr;
   GenericEnvironment *whereClauseGenEnv = nullptr;
   if (auto whereClause = attr->getWhereClause()) {
-    // 'where' clauses in '@differentiable' attributes of protocol
-    // requirements are not supported.
-    if (isa<ProtocolDecl>(original->getDeclContext()) &&
-        original->isProtocolRequirement()) {
+    // `@differentiable` attributes on protocol requirements do not support
+    // 'where' clauses.
+    if (isOriginalProtocolRequirement) {
       TC.diagnose(attr->getLocation(),
-                  diag::differentiable_attr_protocol_where_clause);
+                  diag::differentiable_attr_protocol_req_where_clause);
       attr->setInvalid();
       return;
     }
@@ -3398,6 +3403,15 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
     return;
   }
 
+  // `@differentiable` attributes on protocol requirements do not support
+  // JVP/VJP.
+  if (isOriginalProtocolRequirement && (attr->getJVP() || attr->getVJP())) {
+    TC.diagnose(attr->getLocation(),
+                diag::differentiable_attr_protocol_req_assoc_func);
+    attr->setInvalid();
+    return;
+  }
+
   // Checks that the `candidate` function type equals the `required` function
   // type, disregarding parameter labels and tuple result labels.
   std::function<bool(CanAnyFunctionType, CanType)> checkFunctionSignature;

test/AutoDiff/differentiable_attr_type_checking.swift

@@ -612,7 +612,7 @@ func invalidRequirementLayout<Scalar>(x: Scalar) -> Scalar {
   return x
 }
 
-protocol DifferentiableAttrRequirements : Differentiable {
+protocol ProtocolRequirements : Differentiable {
   // expected-note @+2 {{protocol requires initializer 'init(x:y:)' with type '(x: Float, y: Float)'}}
   @differentiable
   init(x: Float, y: Float)
@@ -638,8 +638,13 @@ protocol DifferentiableAttrRequirements : Differentiable {
   func f2(_ x: Float, _ y: Float) -> Float
 }
 
-// expected-error @+1 {{does not conform to protocol 'DifferentiableAttrRequirements'}}
-struct DiffAttrConformanceErrors : DifferentiableAttrRequirements {
+protocol ProtocolRequirementsRefined : ProtocolRequirements {
+  // expected-note @+1 {{candidate is missing attribute '@differentiable'}}
+  func f1(_ x: Float) -> Float
+}
+
+// expected-error @+1 {{does not conform to protocol 'ProtocolRequirements'}}
+struct DiffAttrConformanceErrors : ProtocolRequirements {
   var x: Float
   var y: Float
 
@@ -683,6 +688,31 @@ struct DiffAttrConformanceErrors : DifferentiableAttrRequirements {
   }
 }
 
+protocol ProtocolRequirementsWithDefault_NoConformingTypes {
+  @differentiable
+  func f1(_ x: Float) -> Float
+}
+extension ProtocolRequirementsWithDefault_NoConformingTypes {
+  // TODO(TF-650): It would be nice to diagnose protocol default implementation
+  // with missing `@differentiable` attribute.
+  func f1(_ x: Float) -> Float { x }
+}
+
+protocol ProtocolRequirementsWithDefault {
+  // expected-note @+2 {{protocol requires function 'f1'}}
+  @differentiable
+  func f1(_ x: Float) -> Float
+}
+extension ProtocolRequirementsWithDefault {
+  // expected-note @+1 {{candidate is missing attribute '@differentiable'}}
+  func f1(_ x: Float) -> Float { x }
+}
+// expected-error @+1 {{type 'DiffAttrConformanceErrors2' does not conform to protocol 'ProtocolRequirementsWithDefault'}}
+struct DiffAttrConformanceErrors2 : ProtocolRequirementsWithDefault {
+  // expected-note @+1 {{candidate is missing attribute '@differentiable'}}
+  func f1(_ x: Float) -> Float { x }
+}
+
 protocol NotRefiningDiffable {
   @differentiable(wrt: x)
   // expected-note @+1 {{protocol requires function 'a' with type '(Float) -> Float'; do you want to add a stub?}}
@@ -846,17 +876,7 @@ func inout2(x: Float, y: inout Float) -> Float {
   let _ = x + y
 }
 
-
-// Missing `@differentiable` attribute, without printing the 'wrt' arguments.
-
-protocol DifferentiableWhereClause: Differentiable {
-  associatedtype Scalar
-
-  @differentiable(where Scalar: Differentiable) // expected-error {{'where' clauses cannot be used in a '@differentiable' attribute on a protocol requirement}}
-  func test(value: Scalar) -> Float
-}
-
-// Missing a `@differentiable` attribute.
+// Test refining protocol requirements with `@differentiable` attribute.
 
 public protocol Distribution {
   associatedtype Value
@@ -870,18 +890,25 @@ public protocol DifferentiableDistribution: Differentiable, Distribution {
 
 public protocol MissingDifferentiableDistribution: DifferentiableDistribution
   where Value: Differentiable {
-  func logProbability(of value: Value) -> Float // expected-note {{candidate is missing attribute '@differentiable(wrt: self)'}}
+  // expected-note @+1 {{candidate is missing attribute '@differentiable(wrt: self)'}}
+  func logProbability(of value: Value) -> Float
 }
 
-// Missing `@differentiable` attribute, without printing the 'wrt' arguments.
+// Test protocol requirement `@differentiable` attribute unsupported features.
 
-protocol Example: Differentiable {
-  associatedtype Scalar: Differentiable
+protocol ProtocolRequirementUnsupported : Differentiable {
+  associatedtype Scalar
 
-  @differentiable
-  func test(value: Scalar) -> Float
-}
+  // expected-error @+1 {{'@differentiable' attribute on protocol requirement cannot specify 'where' clause}}
+  @differentiable(where Scalar: Differentiable)
+  func unsupportedWhereClause(value: Scalar) -> Float
 
-protocol MissingDifferentiableTest: Example {
-  func test(value: Scalar) -> Float // expected-note {{candidate is missing attribute '@differentiable'}}
+  // expected-error @+1 {{'@differentiable' attribute on protocol requirement cannot specify 'jvp:' or 'vjp:'}}
+  @differentiable(wrt: x, jvp: dfoo, vjp: dfoo)
+  func unsupportedDerivatives(_ x: Float) -> Float
+}
+extension ProtocolRequirementUnsupported {
+  func dfoo(_ x: Float) -> (Float, (Float) -> Float) {
+    (x, { $0 })
+  }
 }