[AutoDiff] Disallow differentiation of opaque-result-typed functions. (#32714)

Reject `@differentiable` and `@derivative` attribute for original
functions with opaque result types.

It is not possible to support derivative registration nor the
differentiation transform for such functions.

Resolves SR-12656.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -3151,6 +3151,8 @@ ERROR(autodiff_attr_original_multiple_semantic_results,none,
 ERROR(autodiff_attr_result_not_differentiable,none,
       "can only differentiate functions with results that conform to "
       "'Differentiable', but %0 does not conform to 'Differentiable'", (Type))
+ERROR(autodiff_attr_opaque_result_type_unsupported,none,
+      "cannot differentiate functions returning opaque result types", ())
 
 // differentiation `wrt` parameters clause
 ERROR(diff_function_no_parameters,none,

lib/Sema/TypeCheckAttr.cpp

@@ -4240,6 +4240,15 @@ IndexSubset *DifferentiableAttributeTypeCheckRequest::evaluate(
 
   auto *originalFnTy = original->getInterfaceType()->castTo<AnyFunctionType>();
 
+  // Diagnose if original function has opaque result types.
+  if (auto *opaqueResultTypeDecl = original->getOpaqueResultTypeDecl()) {
+    diags.diagnose(
+        attr->getLocation(),
+        diag::autodiff_attr_opaque_result_type_unsupported);
+    attr->setInvalid();
+    return nullptr;
+  }
+
   // Diagnose if original function is an invalid class member.
   bool isOriginalClassMember = original->getDeclContext() &&
                                original->getDeclContext()->getSelfClassDecl();
@@ -4532,6 +4541,16 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
       return true;
     }
   }
+
+  // Diagnose if original function has opaque result types.
+  if (auto *opaqueResultTypeDecl = originalAFD->getOpaqueResultTypeDecl()) {
+    diags.diagnose(
+        attr->getLocation(),
+        diag::autodiff_attr_opaque_result_type_unsupported);
+    attr->setInvalid();
+    return true;
+  }
+
   // Diagnose if original function is an invalid class member.
   bool isOriginalClassMember =
       originalAFD->getDeclContext() &&
@@ -5083,9 +5102,15 @@ void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
     attr->setInvalid();
     return;
   }
-
   attr->setOriginalFunction(originalAFD);
 
+  // Diagnose if original function has opaque result types.
+  if (auto *opaqueResultTypeDecl = originalAFD->getOpaqueResultTypeDecl()) {
+    diagnose(attr->getLocation(), diag::autodiff_attr_opaque_result_type_unsupported);
+    attr->setInvalid();
+    return;
+  }
+
   // Get the linearity parameter types.
   SmallVector<AnyFunctionType::Param, 4> linearParams;
   expectedOriginalFnType->getSubsetParameters(linearParamIndices, linearParams,

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend-typecheck -verify %s
+// RUN: %target-swift-frontend-typecheck -verify -disable-availability-checking %s
 
 import _Differentiation
 
@@ -1124,3 +1124,13 @@ extension Float {
     fatalError()
   }
 }
+
+// Test original function with opaque result type.
+
+func opaqueResult(_ x: Float) -> some Differentiable { x }
+
+// expected-error @+1 {{could not find function 'opaqueResult' with expected type '(Float) -> Float'}}
+@derivative(of: opaqueResult)
+func vjpOpaqueResult(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}

test/AutoDiff/Sema/differentiable_attr_type_checking.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend-typecheck -verify %s
+// RUN: %target-swift-frontend-typecheck -verify -disable-availability-checking %s
 
 import _Differentiation
 
@@ -697,3 +697,7 @@ struct Accessors: Differentiable {
     _modify { yield &stored }
   }
 }
+
+// expected-error @+1 {{cannot differentiate functions returning opaque result types}}
+@differentiable
+func opaqueResult(_ x: Float) -> some Differentiable { x }

test/AutoDiff/compiler_crashers/sr12656-differentiation-opaque-result-type.swift renamed to test/AutoDiff/compiler_crashers_fixed/sr12656-differentiation-opaque-result-type.swift

@@ -1,4 +1,4 @@
-// RUN: not --crash %target-swift-frontend -disable-availability-checking -emit-sil -verify %s
+// RUN: not %target-swift-frontend -disable-availability-checking -emit-sil -verify %s
 // REQUIRES: asserts
 
 // SR-12656: Differentiation transform crashes for original function with opaque