From 40e5eb2155509f25fd81e54746297a7a96f9a2d7 Mon Sep 17 00:00:00 2001
From: Daniil Kovalev <daniil@kovalev.website>
Date: Mon, 23 Dec 2024 19:29:06 +0300
Subject: [PATCH] [AutoDiff] Allow closure differentiation when used as default
 arg val

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.
---
 .../Mandatory/Differentiation.cpp             | 31 +++++++++++++------
 .../differentiation_diagnostics.swift         |  4 +--
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp
index de3285cdd3f06..1fc394e3920c9 100644
--- a/lib/SILOptimizer/Mandatory/Differentiation.cpp
+++ b/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;
diff --git a/test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift b/test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift
index 5b446df14536d..34f09959d7680 100644
--- a/test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift
+++ b/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