Improve @derivative type-checking diagnostics order.

Attempt to look up original function before checking whether the `value:` result
conforms to `Differentiable`.

This improves diagnostics: "original function not found" should be diagnosed as
early as possible.

Author: dan-zheng

lib/Sema/TypeCheckAttr.cpp

@@ -4054,19 +4054,6 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
     return true;
   }
   attr->setDerivativeKind(kind);
-  // `value: R` result tuple element must conform to `Differentiable`.
-  auto diffableProto = Ctx.getProtocol(KnownProtocolKind::Differentiable);
-  auto valueResultType = valueResultElt.getType();
-  if (valueResultType->hasTypeParameter())
-    valueResultType = derivative->mapTypeIntoContext(valueResultType);
-  auto valueResultConf = TypeChecker::conformsToProtocol(
-      valueResultType, diffableProto, derivative->getDeclContext(), None);
-  if (!valueResultConf) {
-    diags.diagnose(attr->getLocation(),
-                   diag::derivative_attr_result_value_not_differentiable,
-                   valueResultElt.getType());
-    return true;
-  }
 
   // Compute expected original function type and look up original function.
   auto *originalFnType =
@@ -4221,6 +4208,7 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
                                     diffParamTypes);
 
   // Get the differentiability parameters' `TangentVector` associated types.
+  auto *diffableProto = Ctx.getProtocol(KnownProtocolKind::Differentiable);
   auto diffParamTanTypes =
       map<SmallVector<TupleTypeElt, 4>>(diffParamTypes, [&](Type paramType) {
         if (paramType->hasTypeParameter())
@@ -4234,7 +4222,19 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
         return TupleTypeElt(paramAssocType);
       });
 
+  // `value: R` result tuple element must conform to `Differentiable`.
   // Get the `TangentVector` associated type of the `value:` result type.
+  auto valueResultType = valueResultElt.getType();
+  if (valueResultType->hasTypeParameter())
+    valueResultType = derivative->mapTypeIntoContext(valueResultType);
+  auto valueResultConf = TypeChecker::conformsToProtocol(
+      valueResultType, diffableProto, derivative->getDeclContext(), None);
+  if (!valueResultConf) {
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_result_value_not_differentiable,
+                   valueResultElt.getType());
+    return true;
+  }
   auto resultTanType = valueResultConf.getTypeWitnessByName(
       valueResultType, Ctx.Id_TangentVector);
 

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -61,6 +61,22 @@ func vjpSubtractWrt1(x: Float, y: Float) -> (value: Float, pullback: (Float) ->
   return (x - y, { $0 })
 }
 
+// Test invalid original function.
+
+// expected-error @+1 {{use of unresolved identifier 'nonexistentFunction'}}
+@derivative(of: nonexistentFunction)
+func vjpOriginalFunctionNotFound(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+// Test `@derivative` attribute where `value:` result does not conform to `Differentiable`.
+// Invalid original function should be diagnosed first.
+// expected-error @+1 {{use of unresolved identifier 'nonexistentFunction'}}
+@derivative(of: nonexistentFunction)
+func vjpOriginalFunctionNotFound2(_ x: Float) -> (value: Int, pullback: (Float) -> Float) {
+  fatalError()
+}
+
 // Test incorrect `@derivative` declaration type.
 
 // expected-note @+1 {{'incorrectDerivativeType' defined here}}
@@ -83,7 +99,7 @@ func vjpResultIncorrectFirstLabel(x: Float) -> (Float, (Float) -> Float) {
 func vjpResultIncorrectSecondLabel(x: Float) -> (value: Float, (Float) -> Float) {
   return (x, { $0 })
 }
-// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple; first element type 'Int' must conform to 'Differentiable'}}
+// expected-error @+1 {{could not find function 'incorrectDerivativeType' with expected type '(Int) -> Int'}}
 @derivative(of: incorrectDerivativeType)
 func vjpResultNotDifferentiable(x: Int) -> (
   value: Int, pullback: (Int) -> Int