[AutoDiff] Improve getConstrainedDerivativeGenericSignature helper.

include/swift/AST/AutoDiff.h

@@ -29,6 +29,7 @@
 namespace swift {
 
 class AnyFunctionType;
+class SILFunctionType;
 class TupleType;
 
 /// A function type differentiability kind.
@@ -231,6 +232,18 @@ void getSubsetParameterTypes(IndexSubset *indices, AnyFunctionType *type,
                              SmallVectorImpl<Type> &results,
                              bool reverseCurryLevels = false);
 
+/// "Constrained" derivative generic signatures require all differentiability
+/// parameters to conform to the `Differentiable` protocol.
+///
+/// Returns the "constrained" derivative generic signature given:
+/// - An original SIL function type.
+/// - Differentiability parameter indices.
+/// - A possibly "unconstrained" derivative generic signature.
+GenericSignature
+getConstrainedDerivativeGenericSignature(SILFunctionType *originalFnTy,
+                                         IndexSubset *diffParamIndices,
+                                         GenericSignature derivativeGenSig);
+
 } // end namespace autodiff
 
 } // end namespace swift

lib/AST/AutoDiff.cpp

@@ -11,6 +11,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/AST/AutoDiff.h"
+#include "swift/AST/ASTContext.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 
 using namespace swift;
@@ -67,6 +69,31 @@ void autodiff::getSubsetParameterTypes(IndexSubset *subset,
   }
 }
 
+GenericSignature autodiff::getConstrainedDerivativeGenericSignature(
+    SILFunctionType *originalFnTy, IndexSubset *diffParamIndices,
+    GenericSignature derivativeGenSig) {
+  if (!derivativeGenSig)
+    derivativeGenSig = originalFnTy->getSubstGenericSignature();
+  if (!derivativeGenSig)
+    return nullptr;
+  // Constrain all differentiability parameters to `Differentiable`.
+  auto &ctx = originalFnTy->getASTContext();
+  auto *diffableProto = ctx.getProtocol(KnownProtocolKind::Differentiable);
+  SmallVector<Requirement, 4> requirements;
+  for (unsigned paramIdx : diffParamIndices->getIndices()) {
+    auto paramType = originalFnTy->getParameters()[paramIdx].getInterfaceType();
+    Requirement req(RequirementKind::Conformance, paramType,
+                    diffableProto->getDeclaredType());
+    requirements.push_back(req);
+  }
+  return evaluateOrDefault(
+      ctx.evaluator,
+      AbstractGenericSignatureRequest{derivativeGenSig.getPointer(),
+                                      /*addedGenericParams*/ {},
+                                      std::move(requirements)},
+      nullptr);
+}
+
 Type TangentSpace::getType() const {
   switch (kind) {
   case Kind::TangentVector:

lib/SIL/SILFunctionType.cpp

@@ -191,36 +191,6 @@ SILFunctionType::getWitnessMethodClass(SILModule &M) const {
   return nullptr;
 }
 
-// Returns the canonical generic signature for an autodiff derivative function
-// given an existing derivative function generic signature. All
-// differentiability parameters are required to conform to `Differentiable`.
-static CanGenericSignature getAutoDiffDerivativeFunctionGenericSignature(
-    CanGenericSignature derivativeFnGenSig,
-    ArrayRef<SILParameterInfo> originalParameters,
-    IndexSubset *parameterIndices, ModuleDecl *module) {
-  if (!derivativeFnGenSig)
-    return nullptr;
-  auto &ctx = module->getASTContext();
-  GenericSignatureBuilder builder(ctx);
-  // Add derivative function generic signature.
-  builder.addGenericSignature(derivativeFnGenSig);
-  // All differentiability parameters are required to conform to
-  // `Differentiable`.
-  auto source =
-      GenericSignatureBuilder::FloatingRequirementSource::forAbstract();
-  auto *differentiableProtocol =
-      ctx.getProtocol(KnownProtocolKind::Differentiable);
-  for (unsigned paramIdx : parameterIndices->getIndices()) {
-    auto paramType = originalParameters[paramIdx].getInterfaceType();
-    Requirement req(RequirementKind::Conformance, paramType,
-                    differentiableProtocol->getDeclaredType());
-    builder.addRequirement(req, source, module);
-  }
-  return std::move(builder)
-      .computeGenericSignature(SourceLoc(), /*allowConcreteGenericParams*/ true)
-      ->getCanonicalSignature();
-}
-
 CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     IndexSubset *parameterIndices, unsigned resultIndex,
     AutoDiffDerivativeFunctionKind kind, TypeConverter &TC,
@@ -243,8 +213,8 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
   // Get the canonical derivative function generic signature.
   if (!derivativeFnGenSig)
     derivativeFnGenSig = getSubstGenericSignature();
-  derivativeFnGenSig = getAutoDiffDerivativeFunctionGenericSignature(
-      derivativeFnGenSig, getParameters(), parameterIndices, &TC.M);
+  derivativeFnGenSig = autodiff::getConstrainedDerivativeGenericSignature(
+      this, parameterIndices, derivativeFnGenSig).getCanonicalSignature();
 
   // Given a type, returns its formal SIL parameter info.
   auto getTangentParameterInfoForOriginalResult =