add new methods to the NodePrinter to enable range tracking possibilities when demangling a name

Author: charles-zablit

include/swift/Demangling/Demangle.h

@@ -19,6 +19,7 @@
 #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"
@@ -100,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
@@ -235,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)
@@ -466,16 +480,26 @@ class Context {
   /// The lifetime of the returned node tree ends with the lifetime of the
   /// context or with a call of clear().
   NodePointer demangleTypeAsNode(llvm::StringRef MangledName);
-  
+
   /// Demangle the given symbol and return the readable name.
   ///
   /// \param MangledName The mangled symbol string, which start a mangling
   /// 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 symbol and store the result in the `printer`.
+  ///
+  /// \param MangledName The mangled symbol string, which start a mangling
+  /// prefix: _T, _T0, $S, _$S.
+  /// \param printer The NodePrinter that will be used to demangle the symbol.
+  ///
+  /// \returns The demangled string.
+  void demangleSymbolAsString(llvm::StringRef MangledName,
+                              NodePrinter *printer);
 
   /// Demangle the given type and return the readable name.
   ///
@@ -534,6 +558,17 @@ std::string
 demangleSymbolAsString(const char *mangledName, size_t mangledNameLength,
                        const DemangleOptions &options = DemangleOptions());
 
+/// Standalone utility function to demangle the given symbol as string. The
+/// demangled string is stored in the `printer`.
+///
+/// If performance is an issue when demangling multiple symbols,
+/// `Context::demangleSymbolAsString` should be used instead.
+/// \param mangledName The mangled name string pointer.
+/// \param mangledNameLength The length of the mangledName string.
+/// \param printer The NodePrinter that will be used to demangle the symbol.
+void demangleSymbolAsString(const char *mangledName, size_t mangledNameLength,
+                            NodePrinter *printer);
+
 /// Standalone utility function to demangle the given symbol as string.
 ///
 /// If performance is an issue when demangling multiple symbols,
@@ -546,7 +581,7 @@ demangleSymbolAsString(const std::string &mangledName,
   return demangleSymbolAsString(mangledName.data(), mangledName.size(),
                                 options);
 }
-  
+
 /// Standalone utility function to demangle the given symbol as string.
 ///
 /// If performance is an issue when demangling multiple symbols,
@@ -560,6 +595,17 @@ demangleSymbolAsString(llvm::StringRef MangledName,
                                 MangledName.size(), Options);
 }
 
+/// Standalone utility function to demangle the given symbol as string. The
+/// result is stored in the `printer`.
+///
+/// If performance is an issue when demangling multiple symbols,
+/// Context::demangleSymbolAsString should be used instead.
+/// \param MangledName The mangled name string.
+inline void demangleSymbolAsString(llvm::StringRef MangledName,
+                                   NodePrinter *printer) {
+  demangleSymbolAsString(MangledName.data(), MangledName.size(), printer);
+}
+
 /// Standalone utility function to demangle the given type as string.
 ///
 /// If performance is an issue when demangling multiple symbols,
@@ -726,13 +772,19 @@ ManglingErrorOr<const char *> mangleNodeAsObjcCString(NodePointer node,
 /// \endcode
 ///
 /// \param Root A pointer to a parse tree generated by the demangler.
-/// \param Options An object encapsulating options to use to perform this demangling.
+/// \param Options An object encapsulating options to use to perform this
+/// demangling.
 ///
 /// \returns A string representing the demangled name.
-///
 std::string nodeToString(NodePointer Root,
                          const DemangleOptions &Options = DemangleOptions());
 
+/// Transform the node structure to a string, which is stored in the `Printer`.
+///
+/// \param Root A pointer to a parse tree generated by the demangler.
+/// \param Printer A NodePrinter used to pretty print the demangled Node.
+void nodeToString(NodePointer Root, NodePrinter *Printer);
+
 /// Transforms a mangled key path accessor thunk helper
 /// into the identfier/subscript that would be used to invoke it in swift code.
 std::string keyPathSourceString(const char *MangledName,
@@ -778,11 +830,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;
 };
@@ -819,8 +874,17 @@ std::string mangledNameForTypeMetadataAccessor(
     llvm::StringRef moduleName, llvm::StringRef typeName, Node::Kind typeKind,
     Mangle::ManglingFlavor Flavor = Mangle::ManglingFlavor::Default);
 
+/// Base class for printing a Swift demangled node tree.
+///
+/// NodePrinter is used to convert demangled Swift symbol nodes into
+/// human-readable string representations. It handles formatting, indentation,
+/// and Swift-specific syntax.
+///
+/// The virtual methods in this class are meant to be overriden to allow
+/// external consumers (e.g lldb) to track the ranges of components of the
+/// demangled name.
 class NodePrinter {
-private:
+protected:
   DemanglerPrinter Printer;
   DemangleOptions Options;
   bool SpecializationPrefixPrinted = false;
@@ -829,17 +893,24 @@ class NodePrinter {
 public:
   NodePrinter(DemangleOptions options) : Options(options) {}
 
-  std::string printRoot(NodePointer root) {
+  virtual ~NodePrinter() = default;
+
+  void printRoot(NodePointer root) {
     isValid = true;
     print(root, 0);
+  }
+
+  std::string takeString() {
     if (isValid)
       return std::move(Printer).str();
     return "";
   }
 
-private:
+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
@@ -864,13 +935,13 @@ class NodePrinter {
             node->getText() == STDLIB_NAME);
   }
 
-  bool printContext(NodePointer Context);
-
   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,
@@ -893,8 +964,9 @@ class NodePrinter {
 
   NodePointer getChildIf(NodePointer Node, Node::Kind Kind);
 
-  void printFunctionParameters(NodePointer LabelList, NodePointer ParameterType,
-                               unsigned depth, bool showTypes);
+  virtual void printFunctionParameters(NodePointer LabelList,
+                                       NodePointer ParameterType,
+                                       unsigned depth, bool showTypes);
 
   void printFunctionType(NodePointer LabelList, NodePointer node,
                          unsigned depth);
@@ -938,6 +1010,12 @@ class NodePrinter {
                           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.

lib/Demangling/Context.cpp

@@ -66,6 +66,12 @@ std::string Context::demangleSymbolAsString(llvm::StringRef MangledName,
   return demangling;
 }
 
+void Context::demangleSymbolAsString(llvm::StringRef MangledName,
+                                     NodePrinter *printer) {
+  NodePointer root = demangleSymbolAsNode(MangledName);
+  nodeToString(root, printer);
+}
+
 std::string Context::demangleTypeAsString(llvm::StringRef MangledName,
                                           const DemangleOptions &Options) {
   NodePointer root = demangleTypeAsNode(MangledName);
@@ -275,6 +281,13 @@ std::string demangleSymbolAsString(const char *MangledName,
                                     Options);
 }
 
+void demangleSymbolAsString(const char *MangledName, size_t MangledNameLength,
+                            NodePrinter *printer) {
+  Context Ctx;
+  return Ctx.demangleSymbolAsString(StringRef(MangledName, MangledNameLength),
+                                    printer);
+}
+
 std::string demangleTypeAsString(const char *MangledName,
                                  size_t MangledNameLength,
                                  const DemangleOptions &Options) {

lib/Demangling/NodePrinter.cpp

@@ -3569,6 +3569,43 @@ NodePointer NodePrinter::printEntity(NodePointer Entity, unsigned depth,
   return PostfixContext;
 }
 
+void NodePrinter::printFunctionName(bool hasName,
+                                    llvm::StringRef &OverwriteName,
+                                    llvm::StringRef &ExtraName,
+                                    bool MultiWordName, int &ExtraIndex,
+                                    swift::Demangle::NodePointer Entity,
+                                    unsigned int depth) {
+  if (hasName || !OverwriteName.empty()) {
+    if (!ExtraName.empty() && MultiWordName) {
+      Printer << ExtraName;
+      if (ExtraIndex >= 0)
+        Printer << ExtraIndex;
+
+      Printer << " of ";
+      ExtraName = "";
+      ExtraIndex = -1;
+    }
+    size_t CurrentPos = Printer.getStringRef().size();
+    if (!OverwriteName.empty()) {
+      Printer << OverwriteName;
+    } else {
+      auto Name = Entity->getChild(1);
+      if (Name->getKind() != Node::Kind::PrivateDeclName)
+        print(Name, depth + 1);
+
+      if (auto PrivateName = getChildIf(Entity, Node::Kind::PrivateDeclName))
+        print(PrivateName, depth + 1);
+    }
+    if (Printer.getStringRef().size() != CurrentPos && !ExtraName.empty())
+      Printer << '.';
+  }
+  if (!ExtraName.empty()) {
+    Printer << ExtraName;
+    if (ExtraIndex >= 0)
+      Printer << ExtraIndex;
+  }
+}
+
 void NodePrinter::printEntityType(NodePointer Entity, NodePointer type,
                                   NodePointer genericFunctionTypeList,
                                   unsigned depth) {
@@ -3747,12 +3784,32 @@ std::string Demangle::keyPathSourceString(const char *MangledName,
   return invalid;
 }
 
+/// Converts a demangled node to a string.
+///
+/// \param root The root of the AST to demangle.
+/// @param options The `DemangleOptions` which will be used to create the
+/// NodePrinter.
+///
+/// \return The demangled node as a string.
 std::string Demangle::nodeToString(NodePointer root,
                                    const DemangleOptions &options) {
   if (!root)
     return "";
 
-  return NodePrinter(options).printRoot(root);
+  NodePrinter printer = NodePrinter(options);
+  nodeToString(root, &printer);
+  return printer.takeString();
+}
+
+/// Converts a demangled node to a string, which is stored in the `printer`.
+///
+/// \param root The root of the AST to demangle.
+/// @param printer The `NodePrinter` which will be used to print the AST to a
+/// string.
+void Demangle::nodeToString(NodePointer root, NodePrinter *printer) {
+  if (!root)
+    return;
+  printer->printRoot(root);
 }
 
 #endif