Converting a last few areas to use multiple result indices.

Author: BradLarson

lib/SILGen/SILGen.cpp

@@ -1145,11 +1145,8 @@ void SILGenModule::emitDifferentiabilityWitnessesForFunction(
               diffAttr->getDerivativeGenericSignature()) &&
              "Type-checking should resolve derivative generic signatures for "
              "all original SIL functions with generic signatures");
-      auto numResults =
-          F->getLoweredFunctionType()->getNumResults() +
-          F->getLoweredFunctionType()->getNumIndirectMutatingParameters();
-      auto *resultIndices = IndexSubset::getDefault(
-          getASTContext(), numResults, /*includeAll*/ true);
+      auto *resultIndices =
+          autodiff::getAllFunctionSemanticResultIndices(AFD);
       auto witnessGenSig =
           autodiff::getDifferentiabilityWitnessGenericSignature(
               AFD->getGenericSignature(),
@@ -1178,11 +1175,8 @@ void SILGenModule::emitDifferentiabilityWitnessesForFunction(
       auto witnessGenSig =
           autodiff::getDifferentiabilityWitnessGenericSignature(
               origAFD->getGenericSignature(), AFD->getGenericSignature());
-      auto numResults =
-          origFn->getLoweredFunctionType()->getNumResults() +
-          origFn->getLoweredFunctionType()->getNumIndirectMutatingParameters();
-      auto *resultIndices = IndexSubset::getDefault(
-          getASTContext(), numResults, /*includeAll*/ true);
+      auto *resultIndices =
+          autodiff::getAllFunctionSemanticResultIndices(origAFD);
       AutoDiffConfig config(derivAttr->getParameterIndices(), resultIndices,
                             witnessGenSig);
       emitDifferentiabilityWitness(origAFD, origFn,

lib/SILGen/SILGenThunk.cpp

@@ -490,11 +490,13 @@ SILFunction *SILGenModule::getOrCreateDerivativeVTableThunk(
   SILGenFunctionBuilder builder(*this);
   auto originalFnDeclRef = derivativeFnDeclRef.asAutoDiffOriginalFunction();
   Mangle::ASTMangler mangler;
+  auto *resultIndices = autodiff::getAllFunctionSemanticResultIndices(
+      originalFnDeclRef.getAbstractFunctionDecl());
   auto name = mangler.mangleAutoDiffDerivativeFunction(
       originalFnDeclRef.getAbstractFunctionDecl(),
       derivativeId->getKind(),
       AutoDiffConfig(derivativeId->getParameterIndices(),
-                     IndexSubset::get(getASTContext(), 1, {0}),
+                     resultIndices,
                      derivativeId->getDerivativeGenericSignature()),
       /*isVTableThunk*/ true);
   auto *thunk = builder.getOrCreateFunction(
@@ -515,7 +517,8 @@ SILFunction *SILGenModule::getOrCreateDerivativeVTableThunk(
   auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
       derivativeId->getParameterIndices(),
       derivativeFnDecl->getInterfaceType()->castTo<AnyFunctionType>());
-  auto *loweredResultIndices = IndexSubset::get(getASTContext(), 1, {0});
+  auto *loweredResultIndices = autodiff::getAllFunctionSemanticResultIndices(
+      originalFnDeclRef.getAbstractFunctionDecl());
   auto diffFn = SGF.B.createDifferentiableFunction(
       loc, loweredParamIndices, loweredResultIndices, originalFn);
   auto derivativeFn = SGF.B.createDifferentiableFunctionExtract(

lib/TBDGen/TBDGen.cpp

@@ -723,22 +723,27 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
 
   // Add derivative function symbols.
   for (const auto *differentiableAttr :
-       AFD->getAttrs().getAttributes<DifferentiableAttr>())
+       AFD->getAttrs().getAttributes<DifferentiableAttr>()) {
+    auto *resultIndices =
+        autodiff::getAllFunctionSemanticResultIndices(AFD);
     addDerivativeConfiguration(
         differentiableAttr->getDifferentiabilityKind(),
         AFD,
         AutoDiffConfig(differentiableAttr->getParameterIndices(),
-                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       resultIndices,
                        differentiableAttr->getDerivativeGenericSignature()));
+  }
   for (const auto *derivativeAttr :
-       AFD->getAttrs().getAttributes<DerivativeAttr>())
+       AFD->getAttrs().getAttributes<DerivativeAttr>()) {
+    auto *resultIndices = autodiff::getAllFunctionSemanticResultIndices(
+        derivativeAttr->getOriginalFunction(AFD->getASTContext()));
     addDerivativeConfiguration(
         DifferentiabilityKind::Reverse,
         derivativeAttr->getOriginalFunction(AFD->getASTContext()),
         AutoDiffConfig(derivativeAttr->getParameterIndices(),
-                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       resultIndices,
                        AFD->getGenericSignature()));
-
+  }
   visitDefaultArguments(AFD, AFD->getParameters());
 
   if (AFD->hasAsync()) {