[AutoDiff] Fix @derivative attribute type-checking crash. (swiftlang#30936)

Fix `@derivative` attribute type-checking crash, so far reproducible only via
`-parse-stdlib`.

The crash occurs because it is not sufficient for type-checking to check for
`Differentiable` conformances. We must also check for invalid `TangentVector`
associated types.

Resolves SR-12559.

Author: dan-zheng

lib/Sema/TypeCheckAttr.cpp

@@ -3157,20 +3157,43 @@ DynamicallyReplacedDeclRequest::evaluate(Evaluator &evaluator,
   return nullptr;
 }
 
-/// Returns true if the given type conforms to `Differentiable` in the given
-/// module.
-static bool conformsToDifferentiable(Type type, DeclContext *DC) {
+/// If the given type conforms to `Differentiable` in the given context, returns
+/// the `ProtocolConformanceRef`. Otherwise, returns an invalid
+/// `ProtocolConformanceRef`.
+///
+/// This helper verifies that the `TangentVector` type witness is valid, in case
+/// the conformance has not been fully checked and the type witness cannot be
+/// resolved.
+static ProtocolConformanceRef getDifferentiableConformance(Type type,
+                                                           DeclContext *DC) {
   auto &ctx = type->getASTContext();
   auto *differentiableProto =
       ctx.getProtocol(KnownProtocolKind::Differentiable);
-  auto conf = TypeChecker::conformsToProtocol(
-      type, differentiableProto, DC, ConformanceCheckFlags::InExpression);
+  auto conf =
+      TypeChecker::conformsToProtocol(type, differentiableProto, DC, None);
   if (!conf)
-    return false;
+    return ProtocolConformanceRef();
   // Try to get the `TangentVector` type witness, in case the conformance has
-  // not been fully checked and the type witness cannot be resolved.
+  // not been fully checked.
+  Type tanType = conf.getTypeWitnessByName(type, ctx.Id_TangentVector);
+  if (tanType.isNull() || tanType->hasError())
+    return ProtocolConformanceRef();
+  return conf;
+};
+
+/// Returns true if the given type conforms to `Differentiable` in the given
+/// contxt. If `tangentVectorEqualsSelf` is true, also check whether the given
+/// type satisfies `TangentVector == Self`.
+static bool conformsToDifferentiable(Type type, DeclContext *DC,
+                                     bool tangentVectorEqualsSelf = false) {
+  auto conf = getDifferentiableConformance(type, DC);
+  if (conf.isInvalid())
+    return false;
+  if (!tangentVectorEqualsSelf)
+    return true;
+  auto &ctx = type->getASTContext();
   Type tanType = conf.getTypeWitnessByName(type, ctx.Id_TangentVector);
-  return !tanType.isNull() && !tanType->hasError();
+  return type->isEqual(tanType);
 };
 
 IndexSubset *TypeChecker::inferDifferentiabilityParameters(
@@ -4364,9 +4387,8 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
   auto originalResult = originalResults.front();
   auto originalResultType = originalResult.type;
   // Check that the original semantic result conforms to `Differentiable`.
-  auto *diffableProto = Ctx.getProtocol(KnownProtocolKind::Differentiable);
-  auto valueResultConf = TypeChecker::conformsToProtocol(
-      originalResultType, diffableProto, derivative->getDeclContext(), None);
+  auto valueResultConf = getDifferentiableConformance(
+      originalResultType, derivative->getDeclContext());
   if (!valueResultConf) {
     diags.diagnose(attr->getLocation(),
                    diag::derivative_attr_result_value_not_differentiable,
@@ -4467,21 +4489,6 @@ DerivativeAttrOriginalDeclRequest::evaluate(Evaluator &evaluator,
   return nullptr;
 }
 
-/// Returns true if the given type's `TangentVector` is equal to itself in the
-/// given module.
-static bool tangentVectorEqualsSelf(Type type, DeclContext *DC) {
-  assert(conformsToDifferentiable(type, DC));
-  auto &ctx = type->getASTContext();
-  auto *differentiableProto =
-      ctx.getProtocol(KnownProtocolKind::Differentiable);
-  auto conf = TypeChecker::conformsToProtocol(
-                  type, differentiableProto, DC,
-                  ConformanceCheckFlags::InExpression);
-  auto tanType = conf.getTypeWitnessByName(type, ctx.Id_TangentVector);
-  return type->getCanonicalType() == tanType->getCanonicalType();
-};
-
-
 // Computes the linearity parameter indices from the given parsed linearity
 // parameters for the given transpose function. On error, emits diagnostics and
 // returns `nullptr`.
@@ -4600,8 +4607,8 @@ static bool checkLinearityParameters(
         parsedLinearParams.empty() ? attrLoc : parsedLinearParams[i].getLoc();
     // Parameter must conform to `Differentiable` and satisfy
     // `Self == Self.TangentVector`.
-    if (!conformsToDifferentiable(linearParamType, originalAFD) ||
-        !tangentVectorEqualsSelf(linearParamType, originalAFD)) {
+    if (!conformsToDifferentiable(linearParamType, originalAFD,
+                                  /*tangentVectorEqualsSelf*/ true)) {
       diags.diagnose(loc,
                      diag::transpose_attr_invalid_linearity_parameter_or_result,
                      linearParamType.getString(), /*isParameter*/ true);
@@ -4713,8 +4720,8 @@ void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
   if (expectedOriginalResultType->hasTypeParameter())
     expectedOriginalResultType = transpose->mapTypeIntoContext(
         expectedOriginalResultType);
-  if (!conformsToDifferentiable(expectedOriginalResultType, transpose) ||
-      !tangentVectorEqualsSelf(expectedOriginalResultType, transpose)) {
+  if (!conformsToDifferentiable(expectedOriginalResultType, transpose,
+                                /*tangentVectorEqualsSelf*/ true)) {
     diagnoseAndRemoveAttr(
         attr, diag::transpose_attr_invalid_linearity_parameter_or_result,
         expectedOriginalResultType.getString(), /*isParameter*/ false);

test/AutoDiff/compiler_crashers_fixed/sr12559-derivative-attr-type-checking-parse-stdlib.swift

@@ -0,0 +1,30 @@
+// RUN: not %target-swift-frontend-typecheck -parse-stdlib %s
+
+// SR-12559: `@derivative` attribute type-checking crash.
+// This program is not valid, but the compiler should not crash nonetheless.
+// Reproducible only with `-parse-stdlib`.
+
+// The crash occurs because it is not sufficient for attribute type-checking to
+// check for `Differentiable` conformances. We must also check for invalid
+// associated types:
+//
+//   (lldb) p valueResultConf.dump()
+//   (normal_conformance type=AnyDerivative protocol=Differentiable
+//     (assoc_type req=TangentVector type=<<error type>>))
+
+import _Differentiation
+
+struct AnyDerivative: Differentiable {
+  init<T>(_ base: T) {}
+
+  @derivative(of: init)
+  static func _vjpInit<T: Differentiable>(
+    _ base: T
+  ) -> (value: AnyDerivative, pullback: (AnyDerivative) -> T.TangentVector) {
+    fatalError()
+  }
+
+  typealias TangentVector = AnyDerivative
+}
+
+// Assertion failed: (resultTan && "Original result has no tangent space?"), function getAutoDiffDerivativeFunctionLinearMapType, file /Users/danielzheng/swift-merge/swift/lib/AST/Type.cpp, line 5190.