Address review feedback.

Author: dan-zheng

lib/SIL/SILPrinter.cpp

@@ -3168,7 +3168,7 @@ void SILDifferentiabilityWitness::print(
       }
     }
     if (!requirements.empty()) {
-      OS << " [where ";
+      OS << "[where ";
       auto subPrinter = PrintOptions::printSIL();
       subPrinter.GenericEnv = origGenEnv;
       interleave(requirements,
@@ -3189,12 +3189,12 @@ void SILDifferentiabilityWitness::print(
   if (jvp) {
     OS << "  jvp: ";
     printSILFunctionNameAndType(OS, jvp);
-    OS << "\n";
+    OS << '\n';
   }
   if (vjp) {
     OS << "  vjp: ";
     printSILFunctionNameAndType(OS, vjp);
-    OS << "\n";
+    OS << '\n';
   }
   OS << "}\n\n";
 }

lib/SILGen/SILGen.cpp

@@ -772,22 +772,22 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
         jvp = getFunction(SILDeclRef(jvpDecl), NotForDefinition);
       if (auto *vjpDecl = diffAttr->getVJPFunction())
         vjp = getFunction(SILDeclRef(vjpDecl), NotForDefinition);
-      emitDifferentiabilityWitness(AFD, F, diffAttr->getParameterIndices(), jvp,
-                                   vjp,
-                                   diffAttr->getDerivativeGenericSignature());
+      auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
+      AutoDiffConfig config{diffAttr->getParameterIndices(), resultIndices,
+                            diffAttr->getDerivativeGenericSignature()};
+      emitDifferentiabilityWitness(AFD, F, config, jvp, vjp);
     }
   }
   F->verify();
 }
 
 void SILGenModule::emitDifferentiabilityWitness(
     AbstractFunctionDecl *originalAFD, SILFunction *originalFunction,
-    IndexSubset *parameterIndices, SILFunction *jvp, SILFunction *vjp,
-    GenericSignature *derivativeGenSig) {
+    const AutoDiffConfig &config, SILFunction *jvp, SILFunction *vjp) {
   auto *origFnType = originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
   auto origSilFnType = originalFunction->getLoweredFunctionType();
   auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
-      parameterIndices, origFnType);
+      config.parameterIndices, origFnType);
   // NOTE(TF-893): Extending capacity is necessary when `origSilFnType` has
   // parameters corresponding to captured variables. These parameters do not
   // appear in the type of `origFnType`.
@@ -813,7 +813,7 @@ void SILGenModule::emitDifferentiabilityWitness(
 
   // Get or create differentiability witness.
   CanGenericSignature derivativeCanGenSig;
-  if (derivativeGenSig)
+  if (auto *derivativeGenSig = config.derivativeGenericSignature)
     derivativeCanGenSig = derivativeGenSig->getCanonicalSignature();
   // TODO(TF-835): Use simpler derivative generic signature logic below when
   // type-checking no longer generates implicit `@differentiable` attributes.
@@ -830,13 +830,12 @@ void SILGenModule::emitDifferentiabilityWitness(
       derivativeCanGenSig = vjpCanGenSig;
     assert(derivativeCanGenSig == vjpCanGenSig);
   }
-  auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
   // Create new SIL differentiability witness.
   // Witness JVP and VJP are set below.
   auto *diffWitness = SILDifferentiabilityWitness::create(
       M, originalFunction->getLinkage(), originalFunction,
-      loweredParamIndices, resultIndices, derivativeGenSig, /*jvp*/ nullptr,
-      /*vjp*/ nullptr, /*isSerialized*/ true);
+      loweredParamIndices, config.resultIndices, derivativeCanGenSig,
+      /*jvp*/ nullptr, /*vjp*/ nullptr, /*isSerialized*/ true);
 
   // Set derivative function in differentiability witness.
   auto setDerivativeInDifferentiabilityWitness =
@@ -852,8 +851,10 @@ void SILGenModule::emitDifferentiabilityWitness(
       derivativeThunk = getOrCreateAutoDiffDerivativeFunctionThunk(
           originalFunction, indices, derivative, kind, reorderSelf);
     } else {
+      // Note: `AutoDiffDerivativeFunctionIdentifier` must be constructed with
+      // the AST-level parameter indices, not the SIL-level ones.
       auto *id = AutoDiffDerivativeFunctionIdentifier::get(
-          kind, parameterIndices, getASTContext());
+          kind, config.parameterIndices, getASTContext());
       derivativeThunk = getOrCreateAutoDiffThunk(
           SILDeclRef(originalAFD).asAutoDiffDerivativeFunction(id), derivative,
           expectedDerivativeType);

lib/SILGen/SILGen.h

@@ -320,13 +320,12 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
 
   // SWIFT_ENABLE_TENSORFLOW
   /// Emit the differentiability witness for the given original function
-  /// declaration and SIL function, parameter indices, and JVP and VJP
+  /// declaration and SIL function, autodiff configuration, and JVP and VJP
   /// functions (null if undefined).
   void emitDifferentiabilityWitness(AbstractFunctionDecl *originalAFD,
                                     SILFunction *originalFunction,
-                                    IndexSubset *parameterIndices,
-                                    SILFunction *jvp, SILFunction *vjp,
-                                    GenericSignature *derivativeGenSig);
+                                    const AutoDiffConfig &config,
+                                    SILFunction *jvp, SILFunction *vjp);
   // SWIFT_ENABLE_TENSORFLOW END
 
   /// Emit the lazy initializer function for a global pattern binding