Allow for passing a custom Archetype naming schema to the demangler.

include/swift/Demangling/Demangle.h

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -35,6 +35,11 @@ namespace Demangle {
 
 enum class SymbolicReferenceKind : uint8_t;
 
+/// A simple default implementation that assigns letters to archetypes in
+/// alphabetic order.
+std::string archetypeName(uint64_t index, uint64_t depth);
+
+/// Display style options for the demangler.
 struct DemangleOptions {
   bool SynthesizeSugarOnTypes = false;
   bool DisplayDebuggerGeneratedModule = true;
@@ -52,6 +57,7 @@ struct DemangleOptions {
   bool ShortenArchetype = false;
   bool ShowPrivateDiscriminators = true;
   bool ShowFunctionArgumentTypes = true;
+  std::function<std::string(uint64_t, uint64_t)> ArchetypeName = archetypeName;
 
   DemangleOptions() {}
 
@@ -346,17 +352,19 @@ class Context {
   /// prefix: _T, _T0, $S, _$S.
   ///
   /// \returns The demangled string.
-  std::string demangleSymbolAsString(llvm::StringRef MangledName,
-                            const DemangleOptions &Options = DemangleOptions());
+  std::string demangleSymbolAsString(
+      llvm::StringRef MangledName,
+      const DemangleOptions &Options = DemangleOptions());
 
   /// Demangle the given type and return the readable name.
   ///
   /// \param MangledName The mangled type string, which does _not_ start with
   /// a mangling prefix.
   ///
   /// \returns The demangled string.
-  std::string demangleTypeAsString(llvm::StringRef MangledName,
-                            const DemangleOptions &Options = DemangleOptions());
+  std::string
+  demangleTypeAsString(llvm::StringRef MangledName,
+                       const DemangleOptions &Options = DemangleOptions());
 
   /// Returns true if the mangledName refers to a thunk function.
   ///
@@ -584,7 +592,6 @@ bool nodeConsumesGenericArgs(Node *node);
 bool isSpecialized(Node *node);
 
 NodePointer getUnspecialized(Node *node, NodeFactory &Factory);
-std::string archetypeName(Node::IndexType index, Node::IndexType depth);
 
 /// Returns true if the node \p kind refers to a context node, e.g. a nominal
 /// type or a function.

lib/Demangling/NodePrinter.cpp

@@ -50,8 +50,7 @@ DemanglerPrinter &DemanglerPrinter::operator<<(long long n) & {
   return *this;
 }
 
-std::string Demangle::archetypeName(Node::IndexType index,
-                                    Node::IndexType depth) {
+std::string Demangle::archetypeName(uint64_t index, uint64_t depth) {
   DemanglerPrinter name;
   do {
     name << (char)('A' + (index % 26));
@@ -1959,7 +1958,7 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
         }
         // FIXME: Depth won't match when a generic signature applies to a
         // method in generic type context.
-        Printer << archetypeName(index, depth);
+        Printer << Options.ArchetypeName(index, depth);
       }
     }
 
@@ -2036,13 +2035,8 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
   }
   case Node::Kind::DependentGenericParamType: {
     unsigned index = Node->getChild(1)->getIndex();
-    do {
-      Printer << (char)('A' + (index % 26));
-      index /= 26;
-    } while (index);
     unsigned depth = Node->getChild(0)->getIndex();
-    if (depth != 0)
-      Printer << depth;
+    Printer << Options.ArchetypeName(index, depth);
     return nullptr;
   }
   case Node::Kind::DependentGenericType: {

unittests/Basic/DemangleTest.cpp

@@ -31,3 +31,15 @@ TEST(Demangle, IsObjCSymbol) {
             isObjCSymbol(llvm::StringRef("_$s3pat7inlinedSo8NSNumberCvp")));
   EXPECT_EQ(true, isObjCSymbol(llvm::StringRef("_$sSC3fooyS2d_SdtFTO")));
 }
+
+TEST(Demangle, CustomArchetypes) {
+  std::string SymbolName = "_$s1a1gyq_q__xt_tr0_lF";
+  std::string DemangledName = "a.g<Q, U>((U, Q)) -> U";
+
+  DemangleOptions Options;
+  Options.ArchetypeName = [](uint64_t index, uint64_t depth) {
+    return index ? "U" : "Q";
+  };
+  std::string Result = demangleSymbolAsString(SymbolName, Options);
+  EXPECT_STREQ(DemangledName.c_str(), Result.c_str());
+}