Address review comments.

- Add doc comments to `getAutoDiffFunctionLinkage`.
- Optimize "curry level unwrapping" in
  `AnyFunctionType::getAutoDiffAssociatedFunctionType`.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -572,16 +572,20 @@ getOffsetForAutoDiffAssociatedFunction(unsigned order,
 unsigned
 getNumAutoDiffAssociatedFunctions(unsigned differentiationOrder);
 
-// Retrieve config from the function name of a variant of
-// `Builtin.autodiffApply`, e.g. `Builtin.autodiffApply_jvp_arity2_order1`.
-// Returns true if the function name is parsed successfully.
+/// Retrieve config from the function name of a variant of
+/// `Builtin.autodiffApply`, e.g. `Builtin.autodiffApply_jvp_arity2_order1`.
+/// Returns true if the function name is parsed successfully.
 bool getBuiltinAutoDiffApplyConfig(StringRef operationName,
                                    AutoDiffAssociatedFunctionKind &kind,
                                    unsigned &arity, unsigned &order,
                                    bool &rethrows);
 
+/// Computes the correct linkage for associated functions given the linkage of
+/// the original function. If the original linkage is not external and
+/// `isAssocFnExported` is true, use the original function's linkage. Otherwise,
+/// return hidden linkage.
 SILLinkage getAutoDiffFunctionLinkage(SILLinkage originalLinkage,
-                                      bool isExported);
+                                      bool isAssocFnExported);
 
 } // end namespace autodiff
 

lib/AST/AutoDiff.cpp

@@ -90,20 +90,22 @@ bool autodiff::getBuiltinAutoDiffApplyConfig(
 }
 
 SILLinkage autodiff::getAutoDiffFunctionLinkage(SILLinkage originalLinkage,
-                                                bool isExported) {
+                                                bool isAssocFnExported) {
   // If the original is defined externally, then the AD pass is just generating
   // associated functions for use in the current module and therefore these
   // associated functions should not be visible outside the module.
   if (isAvailableExternally(originalLinkage))
     return SILLinkage::Hidden;
 
   // If the original is public, then external modules may need to link the
-  // associated function. Make the associated function public unless
-  // differentiation is not explicitly requested.
+  // associated function. Return the linkage of the original function, unless
+  // the associated function is not exported (i.e. differentiation is not
+  // explicitly requested via a `[differentiable]` attribute on the original
+  // function).
   if (originalLinkage == SILLinkage::Public ||
       originalLinkage == SILLinkage::PublicNonABI ||
       originalLinkage == SILLinkage::Shared)
-    return isExported ? originalLinkage : SILLinkage::Hidden;
+    return isAssocFnExported ? originalLinkage : SILLinkage::Hidden;
 
   // Otherwise, the original function is defined and used only in the current
   // module, so external modules will never try to access the associated

lib/AST/Type.cpp

@@ -4463,7 +4463,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
   indices->getSubsetParameterTypes(this, wrtParamTypes);
 
   // Unwrap curry levels. At most, two parameter lists are necessary, for
-  // curried method types.
+  // curried method types with a `(Self)` parameter list.
   SmallVector<AnyFunctionType *, 2> curryLevels;
   auto *currentLevel = eraseDynamicSelfType()->castTo<AnyFunctionType>();
   for (unsigned i : range(2)) {
@@ -4576,7 +4576,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
 AnyFunctionType *
 AnyFunctionType::getAutoDiffOriginalFunctionType() {
   // Unwrap curry levels. At most, two parameter lists are necessary, for
-  // curried method types.
+  // curried method types with a `(Self)` parameter list.
   SmallVector<AnyFunctionType *, 2> curryLevels;
   auto *currentLevel = this;
   for (unsigned i : range(2)) {

lib/SILGen/SILGenPoly.cpp

@@ -3531,7 +3531,7 @@ SILGenModule::getOrCreateAutoDiffAssociatedFunctionReorderingThunk(
   auto loc = assocFn->getLocation();
   SILGenFunctionBuilder fb(*this);
   auto linkage = autodiff::getAutoDiffFunctionLinkage(
-      original->getLinkage(), /*isExported*/ true);
+      original->getLinkage(), /*isAssocFnExported*/ true);
   auto *thunk = fb.getOrCreateFunction(
       loc, name, linkage, targetType, IsBare, IsNotTransparent,
       assocFn->isSerialized(), assocFn->isDynamicallyReplaceable(),