[AutoDiff] NFC: add ASTMangler entry points for AutoDiff-generated decls.

Add `ASTMangler::mangleAutoDiffGeneratedDeclaration`.

Author: dan-zheng

include/swift/AST/ASTMangler.h

@@ -172,6 +172,19 @@ class ASTMangler : public Mangler {
                                             AutoDiffLinearMapKind kind,
                                             AutoDiffConfig config);
 
+  /// Mangle the AutoDiff generated declaration for the given:
+  /// - Generated declaration kind: linear map struct or branching trace enum.
+  /// - Mangled original function name.
+  /// - Basic block number.
+  /// - Linear map kind: differential or pullback.
+  /// - Derivative function configuration: parameter/result indices and
+  ///   derivative generic signature.
+  std::string
+  mangleAutoDiffGeneratedDeclaration(AutoDiffGeneratedDeclarationKind declKind,
+                                     StringRef origFnName, unsigned bbId,
+                                     AutoDiffLinearMapKind linearMapKind,
+                                     AutoDiffConfig config);
+
   /// Mangle a SIL differentiability witness key:
   /// - Mangled original function name.
   /// - Parameter indices.

include/swift/AST/AutoDiff.h

@@ -164,6 +164,12 @@ struct DifferentiabilityWitnessFunctionKind {
   Optional<AutoDiffDerivativeFunctionKind> getAsDerivativeFunctionKind() const;
 };
 
+/// The kind of a declaration generated by the differentiation transform.
+enum class AutoDiffGeneratedDeclarationKind : uint8_t {
+  LinearMapStruct,
+  BranchingTraceEnum
+};
+
 /// SIL-level automatic differentiation indices. Consists of:
 /// - Parameter indices: indices of parameters to differentiate with respect to.
 /// - Result index: index of the result to differentiate from.

lib/AST/ASTMangler.cpp

@@ -424,6 +424,45 @@ std::string ASTMangler::mangleAutoDiffLinearMapHelper(
   return result;
 }
 
+std::string ASTMangler::mangleAutoDiffGeneratedDeclaration(
+    AutoDiffGeneratedDeclarationKind declKind, StringRef origFnName,
+    unsigned bbId, AutoDiffLinearMapKind linearMapKind, AutoDiffConfig config) {
+  beginManglingWithoutPrefix();
+
+  Buffer << "_AD__" << origFnName << "_bb" + std::to_string(bbId);
+  switch (declKind) {
+  case AutoDiffGeneratedDeclarationKind::LinearMapStruct:
+    switch (linearMapKind) {
+    case AutoDiffLinearMapKind::Differential:
+      Buffer << "__DF__";
+      break;
+    case AutoDiffLinearMapKind::Pullback:
+      Buffer << "__PB__";
+      break;
+    }
+    break;
+  case AutoDiffGeneratedDeclarationKind::BranchingTraceEnum:
+    switch (linearMapKind) {
+    case AutoDiffLinearMapKind::Differential:
+      Buffer << "__Succ__";
+      break;
+    case AutoDiffLinearMapKind::Pullback:
+      Buffer << "__Pred__";
+      break;
+    }
+    break;
+  }
+  Buffer << config.getSILAutoDiffIndices().mangle();
+  if (config.derivativeGenericSignature) {
+    Buffer << '_';
+    appendGenericSignature(config.derivativeGenericSignature);
+  }
+
+  auto result = Storage.str().str();
+  Storage.clear();
+  return result;
+}
+
 std::string ASTMangler::mangleSILDifferentiabilityWitnessKey(
     SILDifferentiabilityWitnessKey key) {
   // TODO(TF-20): Make the mangling scheme robust. Support demangling.

lib/SILOptimizer/Utils/Differentiation/LinearMapInfo.cpp

@@ -130,19 +130,15 @@ EnumDecl *LinearMapInfo::createBranchingTraceDecl(
   auto *moduleDecl = original->getModule().getSwiftModule();
   auto &file = getDeclarationFileUnit();
   // Create a branching trace enum.
-  std::string enumName;
-  switch (kind) {
-  case AutoDiffLinearMapKind::Differential:
-    enumName = "_AD__" + original->getName().str() + "_bb" +
-               std::to_string(originalBB->getDebugID()) + "__Succ__" +
-               indices.mangle();
-    break;
-  case AutoDiffLinearMapKind::Pullback:
-    enumName = "_AD__" + original->getName().str() + "_bb" +
-               std::to_string(originalBB->getDebugID()) + "__Pred__" +
-               indices.mangle();
-    break;
-  }
+  Mangle::ASTMangler mangler;
+  auto *resultIndices = IndexSubset::get(
+      original->getASTContext(),
+      original->getLoweredFunctionType()->getNumResults(), indices.source);
+  auto *parameterIndices = indices.parameters;
+  AutoDiffConfig config(parameterIndices, resultIndices, genericSig);
+  auto enumName = mangler.mangleAutoDiffGeneratedDeclaration(
+      AutoDiffGeneratedDeclarationKind::BranchingTraceEnum,
+      original->getName().str(), originalBB->getDebugID(), kind, config);
   auto enumId = astCtx.getIdentifier(enumName);
   auto loc = original->getLocation().getSourceLoc();
   GenericParamList *genericParams = nullptr;
@@ -205,19 +201,16 @@ LinearMapInfo::createLinearMapStruct(SILBasicBlock *originalBB,
   auto *original = originalBB->getParent();
   auto &astCtx = original->getASTContext();
   auto &file = getDeclarationFileUnit();
-  std::string structName;
-  switch (kind) {
-  case swift::AutoDiffLinearMapKind::Differential:
-    structName = "_AD__" + original->getName().str() + "_bb" +
-                 std::to_string(originalBB->getDebugID()) + "__DF__" +
-                 indices.mangle();
-    break;
-  case swift::AutoDiffLinearMapKind::Pullback:
-    structName = "_AD__" + original->getName().str() + "_bb" +
-                 std::to_string(originalBB->getDebugID()) + "__PB__" +
-                 indices.mangle();
-    break;
-  }
+  // Create a linear map struct.
+  Mangle::ASTMangler mangler;
+  auto *resultIndices = IndexSubset::get(
+      original->getASTContext(),
+      original->getLoweredFunctionType()->getNumResults(), indices.source);
+  auto *parameterIndices = indices.parameters;
+  AutoDiffConfig config(parameterIndices, resultIndices, genericSig);
+  auto structName = mangler.mangleAutoDiffGeneratedDeclaration(
+      AutoDiffGeneratedDeclarationKind::LinearMapStruct,
+      original->getName().str(), originalBB->getDebugID(), kind, config);
   auto structId = astCtx.getIdentifier(structName);
   GenericParamList *genericParams = nullptr;
   if (genericSig)