Emit @noDerivative fixit only when valid.

Emit `@noDerivative` fixit only when there is at least one valid
differentiability/linearity parameter. Otherwise, adding `@noDerivative`
results in an ill-formed `@differentiable` function type.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -4097,8 +4097,10 @@ ERROR(attr_only_on_parameters_of_differentiable,none,
 ERROR(differentiable_function_type_invalid_parameter,none,
       "parameter type '%0' does not conform to 'Differentiable'"
       "%select{| and satisfy '%0 == %0.TangentVector'}1, but the enclosing "
-      "function type is '@differentiable%select{|(linear)}1'; did you want to "
-      "add '@noDerivative' to this parameter?", (StringRef, bool))
+      "function type is '@differentiable%select{|(linear)}1'"
+      "%select{|; did you want to add '@noDerivative' to this parameter?}2",
+      (StringRef, /*tangentVectorEqualsSelf*/ bool,
+       /*hasValidDifferentiabilityParameter*/ bool))
 ERROR(differentiable_function_type_invalid_result,none,
       "result type '%0' does not conform to 'Differentiable'"
       "%select{| and satisfy '%0 == %0.TangentVector'}1, but the enclosing "

lib/Sema/TypeCheckType.cpp

@@ -2581,29 +2581,48 @@ bool TypeResolver::resolveASTFunctionTypeParams(
       // SWIFT_ENABLE_TENSORFLOW END
     }
 
-    // SWIFT_ENABLE_TENSORFLOW
-    // All non-`@noDerivative` parameters of `@differentiable` and
-    // `@differentiable(linear)` function types must be differentiable.
-    if (diffKind != DifferentiabilityKind::NonDifferentiable &&
-        resolution.getStage() != TypeResolutionStage::Structural) {
-      if (!noDerivative) {
-        bool isLinear = diffKind == DifferentiabilityKind::Linear;
-        if (!isDifferentiable(ty, /*tangentVectorEqualsSelf*/ isLinear)) {
-          diagnose(eltTypeRepr->getLoc(),
-                   diag::differentiable_function_type_invalid_parameter,
-                   ty->getString(), isLinear)
-              .fixItInsert(eltTypeRepr->getLoc(), "@noDerivative ");
-        }
-      }
-    }
-    // SWIFT_ENABLE_TENSORFLOW END
-
     auto paramFlags = ParameterTypeFlags::fromParameterType(
         ty, variadic, autoclosure, /*isNonEphemeral*/ false, ownership,
         noDerivative);
     elements.emplace_back(ty, Identifier(), paramFlags);
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  // All non-`@noDerivative` parameters of `@differentiable` and
+  // `@differentiable(linear)` function types must be differentiable.
+  if (diffKind != DifferentiabilityKind::NonDifferentiable &&
+      resolution.getStage() != TypeResolutionStage::Structural) {
+    bool isLinear = diffKind == DifferentiabilityKind::Linear;
+    // Emit `@noDerivative` fixit only if there is at least one valid
+    // differentiability/linearity parameter. Otherwise, adding `@noDerivative`
+    // produces an ill-formed function type.
+    auto hasValidDifferentiabilityParam =
+        llvm::find_if(elements, [&](AnyFunctionType::Param param) {
+          if (param.isNoDerivative())
+            return false;
+          return isDifferentiable(param.getPlainType(),
+                                  /*tangentVectorEqualsSelf*/ isLinear);
+        }) != elements.end();
+    // });
+    for (unsigned i = 0, end = inputRepr->getNumElements(); i != end; ++i) {
+      auto *eltTypeRepr = inputRepr->getElementType(i);
+      auto param = elements[i];
+      if (param.isNoDerivative())
+        continue;
+      auto paramType = param.getPlainType();
+      if (isDifferentiable(paramType, /*tangentVectorEqualsSelf*/ isLinear))
+        continue;
+      auto paramTypeString = paramType->getString();
+      auto diagnostic =
+          diagnose(eltTypeRepr->getLoc(),
+                   diag::differentiable_function_type_invalid_parameter,
+                   paramTypeString, isLinear, hasValidDifferentiabilityParam);
+      if (hasValidDifferentiabilityParam)
+        diagnostic.fixItInsert(eltTypeRepr->getLoc(), "@noDerivative ");
+    }
+  }
+  // SWIFT_ENABLE_TENSORFLOW END
+
   return false;
 }
 

test/AutoDiff/differentiable_func_type_type_checking.swift

@@ -8,12 +8,28 @@ let _: @differentiable (Float) throws -> Float
 //===----------------------------------------------------------------------===//
 
 struct NonDiffType { var x: Int }
+
 // FIXME: Properly type-check parameters and the result's differentiability
-// expected-error @+1 {{parameter type 'NonDiffType' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'; did you want to add '@noDerivative' to this parameter?}} {{25-25=@noDerivative }}
+// expected-error @+1 {{parameter type 'NonDiffType' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'}}
 let _: @differentiable (NonDiffType) -> Float
+
+// Emit `@noDerivative` fixit iff there is at least one valid differentiability parameter.
+// expected-error @+1 {{parameter type 'NonDiffType' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'; did you want to add '@noDerivative' to this parameter?}} {{32-32=@noDerivative }}
+let _: @differentiable (Float, NonDiffType) -> Float
+
+// expected-error @+1 {{result type 'NonDiffType' does not conform to 'Differentiable' and satisfy 'NonDiffType == NonDiffType.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
+let _: @differentiable(linear) (Float) -> NonDiffType
+
+// Emit `@noDerivative` fixit iff there is at least one valid linearity parameter.
+// expected-error @+1 {{parameter type 'NonDiffType' does not conform to 'Differentiable' and satisfy 'NonDiffType == NonDiffType.TangentVector', but the enclosing function type is '@differentiable(linear)'; did you want to add '@noDerivative' to this parameter?}} {{40-40=@noDerivative }}
+let _: @differentiable(linear) (Float, NonDiffType) -> Float
+
 // expected-error @+1 {{result type 'NonDiffType' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'}}
 let _: @differentiable (Float) -> NonDiffType
 
+// expected-error @+1 {{result type 'NonDiffType' does not conform to 'Differentiable' and satisfy 'NonDiffType == NonDiffType.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
+let _: @differentiable(linear) (Float) -> NonDiffType
+
 let _: @differentiable(linear) (Float) -> Float
 
 //===----------------------------------------------------------------------===//
@@ -115,7 +131,7 @@ func inferredConformancesGeneric<T, U>(_: @differentiable (Vector<T>) -> Vector<
 
 // expected-error @+5 {{generic signature requires types 'Vector<T>' and 'Vector<T>.TangentVector' to be the same}}
 // expected-error @+4 {{generic signature requires types 'Vector<U>' and 'Vector<U>.TangentVector' to be the same}}
-// expected-error @+3 {{parameter type 'Vector<T>' does not conform to 'Differentiable' and satisfy 'Vector<T> == Vector<T>.TangentVector', but the enclosing function type is '@differentiable(linear)'; did you want to add '@noDerivative' to this parameter?}}
+// expected-error @+3 {{parameter type 'Vector<T>' does not conform to 'Differentiable' and satisfy 'Vector<T> == Vector<T>.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
 // expected-error @+2 {{result type 'Vector<U>' does not conform to 'Differentiable' and satisfy 'Vector<U> == Vector<U>.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
 // expected-note @+1 2 {{where 'T' = 'Int'}}
 func inferredConformancesGenericLinear<T, U>(_: @differentiable(linear) (Vector<T>) -> Vector<U>) {}
@@ -132,7 +148,7 @@ inferredConformancesGeneric(diff) // okay!
 func inferredConformancesGenericResult<T, U>() -> @differentiable (Vector<T>) -> Vector<U> {}
 // expected-error @+4 {{generic signature requires types 'Vector<T>' and 'Vector<T>.TangentVector' to be the same}}
 // expected-error @+3 {{generic signature requires types 'Vector<U>' and 'Vector<U>.TangentVector' to be the same}}
-// expected-error @+2 {{parameter type 'Vector<T>' does not conform to 'Differentiable' and satisfy 'Vector<T> == Vector<T>.TangentVector', but the enclosing function type is '@differentiable(linear)'; did you want to add '@noDerivative' to this parameter?}}
+// expected-error @+2 {{parameter type 'Vector<T>' does not conform to 'Differentiable' and satisfy 'Vector<T> == Vector<T>.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
 // expected-error @+1 {{result type 'Vector<U>' does not conform to 'Differentiable' and satisfy 'Vector<U> == Vector<U>.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
 func inferredConformancesGenericResultLinear<T, U>() -> @differentiable(linear) (Vector<T>) -> Vector<U> {}