[AutoDiff] Allow closure differentiation when used as default arg val (#78373)

If the default argument generator (and, consequently, the function taking this default argument) has public visibility, it's OK to have a closure (which always has private visibility) as the default value of the argument.

Author: kovdan01

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -681,15 +681,28 @@ DifferentiationTransformer::emitDerivativeFunctionReference(
     }
     assert(minimalWitness);
     if (original->getFunction()->isSerialized()) {
-      // When dealing with curry thunk, look at the function being wrapped
-      // inside implicit closure. If it has public visibility, the corresponding
-      // differentiability witness also has public visibility. It should be OK
-      // for implicit wrapper closure and its witness to have private linkage.
-      SILFunction *unwrappedFn = getUnwrappedCurryThunkFunction(originalFn);
-      bool isWitnessPublic =
-          unwrappedFn == nullptr
-              ? hasPublicVisibility(minimalWitness->getLinkage())
-              : hasPublicVisibility(unwrappedFn->getLinkage());
+      bool isWitnessPublic;
+      if (SILFunction *unwrappedFn =
+              getUnwrappedCurryThunkFunction(originalFn)) {
+        // When dealing with curry thunk, look at the function being wrapped
+        // inside implicit closure. If it has public visibility, the
+        // corresponding differentiability witness also has public visibility.
+        // It should be OK for implicit wrapper closure and its witness to have
+        // private linkage.
+        isWitnessPublic = hasPublicVisibility(unwrappedFn->getLinkage());
+      } else if (originalFn->getDeclRef().getAbstractClosureExpr() &&
+                 originalFRI->getFunction()
+                     ->getDeclRef()
+                     .isDefaultArgGenerator()) {
+        // If we reference a closure from inside default argument generator,
+        // check against generator's visibility. If the function having this
+        // default argument has public visibility, it's OK to have a closure
+        // (which always has private visibility) as its default value.
+        isWitnessPublic =
+            hasPublicVisibility(originalFRI->getFunction()->getLinkage());
+      } else {
+        isWitnessPublic = hasPublicVisibility(minimalWitness->getLinkage());
+      }
       if (!isWitnessPublic) {
         enum { Inlinable = 0, DefaultArgument = 1 };
         unsigned fragileKind = Inlinable;

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -751,9 +751,7 @@ public func hasImplicitlyDifferentiatedTopLevelDefaultArgument(
   _ f: @differentiable(reverse) (Float) -> Float = implicitlyDifferentiableFromFragile
 ) {}
 
-// TODO(TF-1030): This will eventually not be an error.
-// expected-error @+2 {{function is not differentiable}}
-// expected-note @+1 {{differentiated functions in default arguments must be marked '@differentiable' or have a public '@derivative'; this is not possible with a closure, make a top-level function instead}}
+// No error expected
 public func hasImplicitlyDifferentiatedClosureDefaultArgument(_ f: @differentiable(reverse) (Float) -> Float = { $0 }) {}
 
 @inlinable