Demangler library: add a function swift_demangle_getModuleName to get the module name of a mangled symbol.

rdar://problem/47560963

Author: eeckstein

include/swift/Demangling/Demangle.h

@@ -379,6 +379,14 @@ class Context {
   /// function symbol.
   bool hasSwiftCallingConvention(llvm::StringRef MangledName);
 
+  /// Demangle the given symbol and return the module name of the symbol.
+  ///
+  /// \param MangledName The mangled symbol string, which start a mangling
+  /// prefix: _T, _T0, $S, _$S.
+  ///
+  /// \returns The module name.
+  std::string getModuleName(llvm::StringRef mangledName);
+
   /// Deallocates all nodes.
   ///
   /// The memory which is used for nodes is not freed but recycled for the next
@@ -578,6 +586,14 @@ 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.
+bool isContext(Node::Kind kind);
+
+/// Returns true if the node \p kind refers to a node which is placed before a
+/// function node, e.g. a specialization attribute.
+bool isFunctionAttr(Node::Kind kind);
+
 /// Form a StringRef around the mangled name starting at base, if the name may
 /// contain symbolic references.
 llvm::StringRef makeSymbolicMangledNameStringRef(const char *base);

include/swift/SwiftDemangle/SwiftDemangle.h

@@ -57,6 +57,16 @@ size_t swift_demangle_getSimplifiedDemangledName(const char *MangledName,
                                                  char *OutputBuffer,
                                                  size_t Length);
 
+/// Demangle a Swift symbol and return the module name of the mangled entity.
+///
+/// \returns the length of the demangled module name (even if greater than the
+/// size of the output buffer) or 0 if the input is not a Swift-mangled name
+/// (in which cases \p OutputBuffer is left untouched).
+SWIFT_DEMANGLE_LINKAGE
+size_t swift_demangle_getModuleName(const char *MangledName,
+                                    char *OutputBuffer,
+                                    size_t Length);
+
 /// Demangles a Swift function name and returns true if the function
 /// conforms to the Swift calling convention.
 ///

lib/Demangling/Context.cpp

@@ -171,6 +171,43 @@ bool Context::hasSwiftCallingConvention(llvm::StringRef MangledName) {
   return true;
 }
 
+std::string Context::getModuleName(llvm::StringRef mangledName) {
+  NodePointer node = demangleSymbolAsNode(mangledName);
+  while (node) {
+    switch (node->getKind()) {
+    case Demangle::Node::Kind::Module:
+      return node->getText().str();
+    case Demangle::Node::Kind::TypeMangling:
+    case Demangle::Node::Kind::Type:
+      node = node->getFirstChild();
+      break;
+    case Demangle::Node::Kind::Global: {
+      NodePointer newNode = nullptr;
+      for (NodePointer child : *node) {
+        if (!isFunctionAttr(child->getKind())) {
+          newNode = child;
+          break;
+        }
+      }
+      node = newNode;
+      break;
+    }
+    default:
+      if (isSpecialized(node)) {
+        node = getUnspecialized(node, *D);
+        break;
+      }
+      if (isContext(node->getKind())) {
+        node = node->getFirstChild();
+        break;
+      }
+      return std::string();
+    }
+  }
+  return std::string();
+}
+
+
 //////////////////////////////////
 // Public utility functions     //
 //////////////////////////////////

lib/Demangling/Demangler.cpp

@@ -47,18 +47,6 @@ static bool isDeclName(Node::Kind kind) {
   }
 }
 
-static bool isContext(Node::Kind kind) {
-  switch (kind) {
-#define NODE(ID)
-#define CONTEXT_NODE(ID)                                                \
-    case Node::Kind::ID:
-#include "swift/Demangling/DemangleNodes.def"
-      return true;
-    default:
-      return false;
-  }
-}
-
 static bool isAnyGeneric(Node::Kind kind) {
   switch (kind) {
     case Node::Kind::Structure:
@@ -93,7 +81,25 @@ static bool isRequirement(Node::Kind kind) {
   }
 }
 
-static bool isFunctionAttr(Node::Kind kind) {
+} // anonymous namespace
+
+//////////////////////////////////
+// Public utility functions    //
+//////////////////////////////////
+
+bool swift::Demangle::isContext(Node::Kind kind) {
+  switch (kind) {
+#define NODE(ID)
+#define CONTEXT_NODE(ID)                                                \
+    case Node::Kind::ID:
+#include "swift/Demangling/DemangleNodes.def"
+      return true;
+    default:
+      return false;
+  }
+}
+
+bool swift::Demangle::isFunctionAttr(Node::Kind kind) {
   switch (kind) {
     case Node::Kind::FunctionSignatureSpecialization:
     case Node::Kind::GenericSpecialization:
@@ -120,12 +126,6 @@ static bool isFunctionAttr(Node::Kind kind) {
   }
 }
 
-} // anonymous namespace
-
-//////////////////////////////////
-// Public utility functions    //
-//////////////////////////////////
-
 llvm::StringRef
 swift::Demangle::makeSymbolicMangledNameStringRef(const char *base) {
   if (!base)

lib/SwiftDemangle/SwiftDemangle.cpp

@@ -58,6 +58,22 @@ size_t swift_demangle_getSimplifiedDemangledName(const char *MangledName,
                                                  Length, Opts);
 }
 
+size_t swift_demangle_getModuleName(const char *MangledName,
+                                    char *OutputBuffer,
+                                    size_t Length) {
+
+  swift::Demangle::Context DCtx;
+  std::string Result = DCtx.getModuleName(llvm::StringRef(MangledName));
+
+  // Copy the result to an output buffer and ensure '\0' termination.
+  if (OutputBuffer && Length > 0) {
+    auto Dest = strncpy(OutputBuffer, Result.c_str(), Length);
+    Dest[Length - 1] = '\0';
+  }
+  return Result.length();
+}
+
+
 int swift_demangle_hasSwiftCallingConvention(const char *MangledName) {
   swift::Demangle::Context DCtx;
   if (DCtx.hasSwiftCallingConvention(llvm::StringRef(MangledName)))

unittests/SwiftDemangle/DemangleTest.cpp

@@ -87,3 +87,21 @@ TEST(FunctionNameDemangleTests, IgnoresNonMangledInputs) {
   EXPECT_STREQ("0123456789abcdef", OutputBuffer);
 }
 
+TEST(FunctionNameDemangleTests, ModuleName) {
+  const char *Sym1 = "_TtCs5Class";
+  const char *ModuleName1 = "Swift";
+  const char *Sym2 = "_TtCC3Mod7ExampleP33_211017DA67536A354F5F5EB94C7AC12E2Pv";
+  const char *ModuleName2 = "Mod";
+  char OutputBuffer[128];
+
+  size_t Result = swift_demangle_getModuleName(Sym1, OutputBuffer,
+                                               sizeof(OutputBuffer));
+  EXPECT_STREQ(ModuleName1, OutputBuffer);
+  EXPECT_EQ(Result, strlen(ModuleName1));
+
+  Result = swift_demangle_getModuleName(Sym2, OutputBuffer,
+                                        sizeof(OutputBuffer));
+  EXPECT_STREQ(ModuleName2, OutputBuffer);
+  EXPECT_EQ(Result, strlen(ModuleName2));
+}
+