Consolidated repeated result index generation into a central function.

Author: BradLarson

include/swift/AST/AutoDiff.h

@@ -32,6 +32,7 @@
 
 namespace swift {
 
+class AbstractFunctionDecl;
 class AnyFunctionType;
 class SourceFile;
 class SILFunctionType;
@@ -575,6 +576,10 @@ void getFunctionSemanticResultTypes(
     SmallVectorImpl<AutoDiffSemanticFunctionResultType> &result,
     GenericEnvironment *genericEnv = nullptr);
 
+/// Returns the indices of all semantic results for a given function.
+IndexSubset *getAllFunctionSemanticResultIndices(
+    const AbstractFunctionDecl *AFD);
+
 /// Returns the lowered SIL parameter indices for the given AST parameter
 /// indices and `AnyfunctionType`.
 ///

lib/AST/AutoDiff.cpp

@@ -211,6 +211,16 @@ void autodiff::getFunctionSemanticResultTypes(
   }
 }
 
+IndexSubset *
+autodiff::getAllFunctionSemanticResultIndices(const AbstractFunctionDecl *AFD) {
+  auto originalFn = AFD->getInterfaceType()->castTo<AnyFunctionType>();
+  SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+  autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
+  auto numResults = semanticResults.size();
+  return IndexSubset::getDefault(
+    AFD->getASTContext(), numResults, /*includeAll*/ true);
+}
+
 // TODO(TF-874): Simplify this helper. See TF-874 for WIP.
 IndexSubset *
 autodiff::getLoweredParameterIndices(IndexSubset *parameterIndices,

lib/IRGen/IRGenMangler.h

@@ -57,9 +57,11 @@ class IRGenMangler : public Mangle::ASTMangler {
       AutoDiffDerivativeFunctionIdentifier *derivativeId) {
     beginManglingWithAutoDiffOriginalFunction(func);
     auto kind = Demangle::getAutoDiffFunctionKind(derivativeId->getKind());
+    auto *resultIndices =
+        autodiff::getAllFunctionSemanticResultIndices(func);
     AutoDiffConfig config(
         derivativeId->getParameterIndices(),
-        IndexSubset::get(func->getASTContext(), 1, {0}),
+        resultIndices,
         derivativeId->getDerivativeGenericSignature());
     appendAutoDiffFunctionParts("TJ", kind, config);
     appendOperator("Tj");
@@ -86,9 +88,11 @@ class IRGenMangler : public Mangle::ASTMangler {
       AutoDiffDerivativeFunctionIdentifier *derivativeId) {
     beginManglingWithAutoDiffOriginalFunction(func);
     auto kind = Demangle::getAutoDiffFunctionKind(derivativeId->getKind());
+    auto *resultIndices =
+        autodiff::getAllFunctionSemanticResultIndices(func);
     AutoDiffConfig config(
         derivativeId->getParameterIndices(),
-        IndexSubset::get(func->getASTContext(), 1, {0}),
+        resultIndices,
         derivativeId->getDerivativeGenericSignature());
     appendAutoDiffFunctionParts("TJ", kind, config);
     appendOperator("Tq");

lib/SIL/IR/SILDeclRef.cpp

@@ -852,13 +852,8 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
     auto *silParameterIndices = autodiff::getLoweredParameterIndices(
         derivativeFunctionIdentifier->getParameterIndices(),
         getDecl()->getInterfaceType()->castTo<AnyFunctionType>());
-    auto originalFn =
-        getDecl()->getInterfaceType()->castTo<AnyFunctionType>();
-    SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
-    autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
-    auto numResults = semanticResults.size();
-    auto *resultIndices = IndexSubset::getDefault(
-        getDecl()->getASTContext(), numResults, /*includeAll*/ true);
+    auto *resultIndices = autodiff::getAllFunctionSemanticResultIndices(
+        asAutoDiffOriginalFunction().getAbstractFunctionDecl());
     AutoDiffConfig silConfig(
         silParameterIndices, resultIndices,
         derivativeFunctionIdentifier->getDerivativeGenericSignature());

lib/Sema/TypeCheckAttr.cpp

@@ -5026,12 +5026,8 @@ IndexSubset *DifferentiableAttributeTypeCheckRequest::evaluate(
     }
     getterDecl->getAttrs().add(newAttr);
     // Register derivative function configuration.
-    auto originalFn = getterDecl->getInterfaceType()->castTo<AnyFunctionType>();
-    SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
-    autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
-    auto numResults = semanticResults.size();
-    auto *resultIndices = IndexSubset::getDefault(
-        ctx, numResults, /*includeAll*/ true);
+    auto *resultIndices =
+        autodiff::getAllFunctionSemanticResultIndices(getterDecl);
     getterDecl->addDerivativeFunctionConfiguration(
         {resolvedDiffParamIndices, resultIndices, derivativeGenSig});
     return resolvedDiffParamIndices;
@@ -5472,12 +5468,8 @@ static bool typeCheckDerivativeAttr(DerivativeAttr *attr) {
   }
 
   // Register derivative function configuration.
-  auto originalFn = originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
-  SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
-  autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
-  auto numResults = semanticResults.size();
-  auto *resultIndices = IndexSubset::getDefault(
-      Ctx, numResults, /*includeAll*/ true);
+  auto *resultIndices =
+    autodiff::getAllFunctionSemanticResultIndices(originalAFD);
   originalAFD->addDerivativeFunctionConfiguration(
       {resolvedDiffParamIndices, resultIndices,
        derivative->getGenericSignature()});

lib/Sema/TypeCheckProtocol.cpp

@@ -491,13 +491,8 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
           witness->getAttrs().add(newAttr);
           success = true;
           // Register derivative function configuration.
-          auto originalFn =
-              witnessAFD->getInterfaceType()->castTo<AnyFunctionType>();
-          SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
-          autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
-          auto numResults = semanticResults.size();
-          auto *resultIndices = IndexSubset::getDefault(
-              ctx, numResults, /*includeAll*/ true);
+          auto *resultIndices =
+              autodiff::getAllFunctionSemanticResultIndices(witnessAFD);
           witnessAFD->addDerivativeFunctionConfiguration(
               {newAttr->getParameterIndices(), resultIndices,
                newAttr->getDerivativeGenericSignature()});

lib/Serialization/ModuleFile.cpp

@@ -668,13 +668,8 @@ void ModuleFile::loadDerivativeFunctionConfigurations(
     auto derivativeGenSig = derivativeGenSigOrError.get();
     // NOTE(TF-1038): Result indices are currently unsupported in derivative
     // registration attributes. In the meantime, always use all results.
-    auto originalFn =
-        originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
-    SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
-    autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
-    auto numResults = semanticResults.size();
-    auto *resultIndices = IndexSubset::getDefault(
-        ctx, numResults, /*includeAll*/ true);
+    auto *resultIndices =
+        autodiff::getAllFunctionSemanticResultIndices(originalAFD);
     results.insert({parameterIndices, resultIndices, derivativeGenSig});
   }
 }