[Demangling] Add name and parameters range tracking when demangling a name

include/swift/Demangling/Demangle.h

@@ -19,9 +19,11 @@
 #ifndef SWIFT_DEMANGLING_DEMANGLE_H
 #define SWIFT_DEMANGLING_DEMANGLE_H
 
+#include "swift/Demangling/Demangle.h"
 #include "swift/Demangling/Errors.h"
 #include "swift/Demangling/ManglingFlavor.h"
 #include "swift/Demangling/NamespaceMacros.h"
+#include "swift/Strings.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Compiler.h"
 
@@ -99,6 +101,7 @@ struct DemangleOptions {
 
 class Node;
 using NodePointer = Node *;
+class NodePrinter;
 
 enum class FunctionSigSpecializationParamKind : unsigned {
   // Option Flags use bits 0-5. This give us 6 bits implying 64 entries to
@@ -234,6 +237,18 @@ class Node {
 public:
   Kind getKind() const { return NodeKind; }
 
+  bool shouldTrackNameRange() const {
+    switch (getKind()) {
+    case Kind::Function:
+    case Kind::Constructor:
+    case Kind::Allocator:
+    case Kind::ExplicitClosure:
+      return true;
+    default:
+      return false;
+    }
+  }
+
   bool isSimilarTo(const Node *other) const {
     if (NodeKind != other->NodeKind
         || NodePayloadKind != other->NodePayloadKind)
@@ -417,6 +432,10 @@ bool isOldFunctionTypeMangling(llvm::StringRef mangledName);
 
 class Demangler;
 
+class DemanglerPrinter;
+
+class TrackingDemanglerPrinter;
+
 /// The demangler context.
 ///
 /// It owns the allocated nodes which are created during demangling.
@@ -472,9 +491,10 @@ 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(),
+                         NodePrinter *printer = nullptr);
 
   /// Demangle the given type and return the readable name.
   ///
@@ -531,7 +551,8 @@ class Context {
 /// \returns The demangled string.
 std::string
 demangleSymbolAsString(const char *mangledName, size_t mangledNameLength,
-                       const DemangleOptions &options = DemangleOptions());
+                       const DemangleOptions &options = DemangleOptions(),
+                       NodePrinter *printer = nullptr);
 
 /// Standalone utility function to demangle the given symbol as string.
 ///
@@ -541,11 +562,12 @@ demangleSymbolAsString(const char *mangledName, size_t mangledNameLength,
 /// \returns The demangled string.
 inline std::string
 demangleSymbolAsString(const std::string &mangledName,
-                       const DemangleOptions &options = DemangleOptions()) {
-  return demangleSymbolAsString(mangledName.data(), mangledName.size(),
-                                options);
+                       const DemangleOptions &options = DemangleOptions(),
+                       NodePrinter *printer = nullptr) {
+  return demangleSymbolAsString(mangledName.data(), mangledName.size(), options,
+                                printer);
 }
-  
+
 /// Standalone utility function to demangle the given symbol as string.
 ///
 /// If performance is an issue when demangling multiple symbols,
@@ -554,9 +576,10 @@ demangleSymbolAsString(const std::string &mangledName,
 /// \returns The demangled string.
 inline std::string
 demangleSymbolAsString(llvm::StringRef MangledName,
-                       const DemangleOptions &Options = DemangleOptions()) {
-  return demangleSymbolAsString(MangledName.data(),
-                                MangledName.size(), Options);
+                       const DemangleOptions &Options = DemangleOptions(),
+                       NodePrinter *printer = nullptr) {
+  return demangleSymbolAsString(MangledName.data(), MangledName.size(), Options,
+                                printer);
 }
 
 /// Standalone utility function to demangle the given type as string.
@@ -730,7 +753,8 @@ ManglingErrorOr<const char *> mangleNodeAsObjcCString(NodePointer node,
 /// \returns A string representing the demangled name.
 ///
 std::string nodeToString(NodePointer Root,
-                         const DemangleOptions &Options = DemangleOptions());
+                         const DemangleOptions &Options = DemangleOptions(),
+                         NodePrinter *printer = nullptr);
 
 /// Transforms a mangled key path accessor thunk helper
 /// into the identfier/subscript that would be used to invoke it in swift code.
@@ -777,11 +801,14 @@ class DemanglerPrinter {
 
   llvm::StringRef getStringRef() const { return Stream; }
 
+  size_t getStreamLength() { return Stream.length(); }
+
   /// Shrinks the buffer.
   void resetSize(size_t toPos) {
     assert(toPos <= Stream.size());
     Stream.resize(toPos);
   }
+
 private:
   std::string Stream;
 };
@@ -818,6 +845,147 @@ std::string mangledNameForTypeMetadataAccessor(
     llvm::StringRef moduleName, llvm::StringRef typeName, Node::Kind typeKind,
     Mangle::ManglingFlavor Flavor = Mangle::ManglingFlavor::Default);
 
+class NodePrinter {
+protected:
+  DemanglerPrinter Printer;
+  DemangleOptions Options;
+  bool SpecializationPrefixPrinted = false;
+  bool isValid = true;
+
+public:
+  NodePrinter(DemangleOptions options) : Options(options) {}
+
+  virtual ~NodePrinter() {}
+
+  std::string printRoot(NodePointer root) {
+    isValid = true;
+    print(root, 0);
+    if (isValid)
+      return std::move(Printer).str();
+    return "";
+  }
+
+protected:
+  static const unsigned MaxDepth = 768;
+
+  size_t getStreamLength() { return Printer.getStreamLength(); }
+
+  /// Called when the node tree in valid.
+  ///
+  /// The demangler already catches most error cases and mostly produces valid
+  /// node trees. But some cases are difficult to catch in the demangler and
+  /// instead the NodePrinter bails.
+  void setInvalid() { isValid = false; }
+
+  void printChildren(Node::iterator begin, Node::iterator end, unsigned depth,
+                     const char *sep = nullptr);
+
+  void printChildren(NodePointer Node, unsigned depth,
+                     const char *sep = nullptr);
+
+  NodePointer getFirstChildOfKind(NodePointer Node, Node::Kind kind);
+
+  void printBoundGenericNoSugar(NodePointer Node, unsigned depth);
+
+  void printOptionalIndex(NodePointer node);
+
+  static bool isSwiftModule(NodePointer node) {
+    return (node->getKind() == Node::Kind::Module &&
+            node->getText() == STDLIB_NAME);
+  }
+
+  static bool isIdentifier(NodePointer node, StringRef desired) {
+    return (node->getKind() == Node::Kind::Identifier &&
+            node->getText() == desired);
+  }
+
+  bool printContext(NodePointer Context);
+
+  enum class SugarType {
+    None,
+    Optional,
+    ImplicitlyUnwrappedOptional,
+    Array,
+    Dictionary
+  };
+
+  enum class TypePrinting { NoType, WithColon, FunctionStyle };
+
+  /// Determine whether this is a "simple" type, from the type-simple
+  /// production.
+  bool isSimpleType(NodePointer Node);
+
+  void printWithParens(NodePointer type, unsigned depth);
+
+  SugarType findSugar(NodePointer Node);
+
+  void printBoundGeneric(NodePointer Node, unsigned depth);
+
+  NodePointer getChildIf(NodePointer Node, Node::Kind Kind);
+
+  virtual void printFunctionParameters(NodePointer LabelList,
+                                       NodePointer ParameterType,
+                                       unsigned depth, bool showTypes);
+
+  void printFunctionType(NodePointer LabelList, NodePointer node,
+                         unsigned depth);
+
+  void printImplFunctionType(NodePointer fn, unsigned depth);
+
+  void printGenericSignature(NodePointer Node, unsigned depth);
+
+  void printFunctionSigSpecializationParams(NodePointer Node, unsigned depth);
+
+  void printSpecializationPrefix(NodePointer node, StringRef Description,
+                                 unsigned depth,
+                                 StringRef ParamPrefix = StringRef());
+
+  /// The main big print function.
+  NodePointer print(NodePointer Node, unsigned depth,
+                    bool asPrefixContext = false);
+
+  NodePointer printAbstractStorage(NodePointer Node, unsigned depth,
+                                   bool asPrefixContent, StringRef ExtraName);
+
+  /// Utility function to print entities.
+  ///
+  /// \param Entity The entity node to print
+  /// \param depth The depth in the print() call tree.
+  /// \param asPrefixContext Should the entity printed as a context which as a
+  ///        prefix to another entity, e.g. the Abc in Abc.def()
+  /// \param TypePr How should the type of the entity be printed, if at all.
+  ///        E.g. with a colon for properties or as a function type.
+  /// \param hasName Does the entity has a name, e.g. a function in contrast to
+  ///        an initializer.
+  /// \param ExtraName An extra name added to the entity name (if any).
+  /// \param ExtraIndex An extra index added to the entity name (if any),
+  ///        e.g. closure #1
+  /// \param OverwriteName If non-empty, print this name instead of the one
+  ///        provided by the node. Gets printed even if hasName is false.
+  /// \return If a non-null node is returned it's a context which must be
+  ///         printed in postfix-form after the entity: "<entity> in <context>".
+  NodePointer printEntity(NodePointer Entity, unsigned depth,
+                          bool asPrefixContext, TypePrinting TypePr,
+                          bool hasName, StringRef ExtraName = "",
+                          int ExtraIndex = -1, StringRef OverwriteName = "");
+
+  virtual void printFunctionName(bool hasName, llvm::StringRef &OverwriteName,
+                                 llvm::StringRef &ExtraName, bool MultiWordName,
+                                 int &ExtraIndex,
+                                 swift::Demangle::NodePointer Entity,
+                                 unsigned int depth);
+
+  /// Print the type of an entity.
+  ///
+  /// \param Entity The entity.
+  /// \param type The type of the entity.
+  /// \param genericFunctionTypeList If not null, the generic argument types
+  ///           which is printed in the generic signature.
+  /// \param depth The depth in the print() call tree.
+  void printEntityType(NodePointer Entity, NodePointer type,
+                       NodePointer genericFunctionTypeList, unsigned depth);
+};
+
 SWIFT_END_INLINE_NAMESPACE
 } // end namespace Demangle
 } // end namespace swift

lib/Demangling/Context.cpp

@@ -56,11 +56,12 @@ NodePointer Context::demangleTypeAsNode(llvm::StringRef MangledName) {
 #if SWIFT_STDLIB_HAS_TYPE_PRINTING
 
 std::string Context::demangleSymbolAsString(llvm::StringRef MangledName,
-                                            const DemangleOptions &Options) {
+                                            const DemangleOptions &Options,
+                                            NodePrinter *printer) {
   NodePointer root = demangleSymbolAsNode(MangledName);
   if (!root) return MangledName.str();
 
-  std::string demangling = nodeToString(root, Options);
+  std::string demangling = nodeToString(root, Options, printer);
   if (demangling.empty())
     return MangledName.str();
   return demangling;
@@ -269,10 +270,11 @@ std::string Context::getModuleName(llvm::StringRef mangledName) {
 
 std::string demangleSymbolAsString(const char *MangledName,
                                    size_t MangledNameLength,
-                                   const DemangleOptions &Options) {
+                                   const DemangleOptions &Options,
+                                   NodePrinter *printer) {
   Context Ctx;
   return Ctx.demangleSymbolAsString(StringRef(MangledName, MangledNameLength),
-                                    Options);
+                                    Options, printer);
 }
 
 std::string demangleTypeAsString(const char *MangledName,