[SourceKit] Introduce a new request for demangling symbol names.

This is SourceKit support for the swift_demangle_getDemangledName/swift_demangle_getSimplifiedDemangledName APIs.
rdar://23525281

Author: akyrtzi

test/SourceKit/Demangle/demangle.swift

@@ -0,0 +1,13 @@
+// RUN: %sourcekitd-test -req=demangle unmangled _TtBf80_  _TtP3foo3bar_ | FileCheck %s
+// CHECK:      START DEMANGLE
+// CHECK-NEXT: <empty>
+// CHECK-NEXT: Builtin.Float80
+// CHECK-NEXT: foo.bar
+// CHECK-NEXT: END DEMANGLE
+
+// RUN: %sourcekitd-test -req=demangle unmangled _TtBf80_  _TtP3foo3bar_ -simplified-demangling | FileCheck %s -check-prefix=SIMPLIFIED
+// SIMPLIFIED:      START DEMANGLE
+// SIMPLIFIED-NEXT: <empty>
+// SIMPLIFIED-NEXT: Builtin.Float80
+// SIMPLIFIED-NEXT: bar
+// SIMPLIFIED-NEXT: END DEMANGLE

tools/SourceKit/tools/sourcekitd-test/Options.td

@@ -69,3 +69,6 @@ def print_response_as_json : Flag<["-"], "print-response-as-json">,
 
 def group_name : Separate<["-"], "group-name">,
   HelpText<"Module group name to print">;
+
+def simplified_demangling : Flag<["-"], "simplified-demangling">,
+  HelpText<"Use simplified demangling for the 'demangle' request">;

tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp

@@ -103,6 +103,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
     case OPT_req:
       Request = llvm::StringSwitch<SourceKitRequest>(InputArg->getValue())
         .Case("version", SourceKitRequest::ProtocolVersion)
+        .Case("demangle", SourceKitRequest::DemangleNames)
         .Case("index", SourceKitRequest::Index)
         .Case("complete", SourceKitRequest::CodeComplete)
         .Case("complete.open", SourceKitRequest::CodeCompleteOpen)
@@ -131,7 +132,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
         .Default(SourceKitRequest::None);
       if (Request == SourceKitRequest::None) {
         llvm::errs() << "error: invalid request, expected one of "
-            << "version/index/complete/cursor/related-idents/syntax-map/structure/"
+            << "version/demangle/index/complete/cursor/related-idents/syntax-map/structure/"
                "format/expand-placeholder/doc-info/sema/interface-gen/interface-gen-open/"
                "find-usr/find-interface/open/edit/print-annotations/extract-comment/"
                "module-groups\n";
@@ -216,12 +217,17 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
     case OPT_INPUT:
       SourceFile = InputArg->getValue();
       SourceText = llvm::None;
+      Inputs.push_back(InputArg->getValue());
       break;
 
     case OPT_json_request_path:
       JsonRequestPath = InputArg->getValue();
       break;
 
+    case OPT_simplified_demangling:
+      SimplifiedDemangling = true;
+      break;
+
     case OPT_UNKNOWN:
       llvm::errs() << "error: unknown argument: "
                    << InputArg->getAsString(ParsedArgs) << '\n';

tools/SourceKit/tools/sourcekitd-test/TestOptions.h

@@ -22,6 +22,7 @@ namespace sourcekitd_test {
 enum class SourceKitRequest {
   None,
   ProtocolVersion,
+  DemangleNames,
   Index,
   CodeComplete,
   CodeCompleteOpen,
@@ -51,6 +52,7 @@ enum class SourceKitRequest {
 
 struct TestOptions {
   SourceKitRequest Request = SourceKitRequest::None;
+  std::vector<std::string> Inputs;
   std::string SourceFile;
   std::string TextInputFile;
   std::string JsonRequestPath;
@@ -71,6 +73,7 @@ struct TestOptions {
   bool CheckInterfaceIsASCII = false;
   bool UsedSema = false;
   bool PrintResponseAsJSON = false;
+  bool SimplifiedDemangling = false;
 
   bool parseArgs(llvm::ArrayRef<const char *> Args);
 };

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -19,6 +19,7 @@
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Regex.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/FileSystem.h"
@@ -53,6 +54,10 @@ static void expandPlaceholders(llvm::MemoryBuffer *SourceBuf,
 static void printModuleGroupNames(sourcekitd_variant_t Info,
                                   llvm::raw_ostream &OS);
 
+static void prepareDemangleRequest(sourcekitd_object_t Req,
+                                   const TestOptions &Opts);
+static void printDemangleResults(sourcekitd_variant_t Info, raw_ostream &OS);
+
 static unsigned resolveFromLineCol(unsigned Line, unsigned Col,
                                    StringRef Filename);
 static unsigned resolveFromLineCol(unsigned Line, unsigned Col,
@@ -76,6 +81,7 @@ static sourcekitd_uid_t KeySourceFile;
 static sourcekitd_uid_t KeyModuleName;
 static sourcekitd_uid_t KeyGroupName;
 static sourcekitd_uid_t KeyName;
+static sourcekitd_uid_t KeyNames;
 static sourcekitd_uid_t KeyFilePath;
 static sourcekitd_uid_t KeyModuleInterfaceName;
 static sourcekitd_uid_t KeyLength;
@@ -105,8 +111,10 @@ static sourcekitd_uid_t KeyPopular;
 static sourcekitd_uid_t KeyUnpopular;
 static sourcekitd_uid_t KeyTypeInterface;
 static sourcekitd_uid_t KeyModuleGroups;
+static sourcekitd_uid_t KeySimplified;
 
 static sourcekitd_uid_t RequestProtocolVersion;
+static sourcekitd_uid_t RequestDemangle;
 static sourcekitd_uid_t RequestIndex;
 static sourcekitd_uid_t RequestCodeComplete;
 static sourcekitd_uid_t RequestCodeCompleteOpen;
@@ -164,6 +172,7 @@ static int skt_main(int argc, const char **argv) {
   KeyModuleName = sourcekitd_uid_get_from_cstr("key.modulename");
   KeyGroupName = sourcekitd_uid_get_from_cstr("key.groupname");
   KeyName = sourcekitd_uid_get_from_cstr("key.name");
+  KeyNames = sourcekitd_uid_get_from_cstr("key.names");
   KeyFilePath = sourcekitd_uid_get_from_cstr("key.filepath");
   KeyModuleInterfaceName = sourcekitd_uid_get_from_cstr("key.module_interface_name");
   KeyLength = sourcekitd_uid_get_from_cstr("key.length");
@@ -195,6 +204,7 @@ static int skt_main(int argc, const char **argv) {
   KeyUnpopular = sourcekitd_uid_get_from_cstr("key.unpopular");
   KeyTypeInterface = sourcekitd_uid_get_from_cstr("key.typeinterface");
   KeyModuleGroups = sourcekitd_uid_get_from_cstr("key.modulegroups");
+  KeySimplified = sourcekitd_uid_get_from_cstr("key.simplified");
 
   SemaDiagnosticStage = sourcekitd_uid_get_from_cstr("source.diagnostic.stage.swift.sema");
 
@@ -203,6 +213,7 @@ static int skt_main(int argc, const char **argv) {
   semaSemaphore = dispatch_semaphore_create(0);
 
   RequestProtocolVersion = sourcekitd_uid_get_from_cstr("source.request.protocol_version");
+  RequestDemangle = sourcekitd_uid_get_from_cstr("source.request.demangle");
   RequestIndex = sourcekitd_uid_get_from_cstr("source.request.indexsource");
   RequestCodeComplete = sourcekitd_uid_get_from_cstr("source.request.codecomplete");
   RequestCodeCompleteOpen = sourcekitd_uid_get_from_cstr("source.request.codecomplete.open");
@@ -341,6 +352,9 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
   if (Optargc < Args.size())
     Opts.CompilerArgs = Args.slice(Optargc+1);
 
+  if (Opts.Request == SourceKitRequest::DemangleNames)
+    Opts.SourceFile.clear();
+
   std::string SourceFile = Opts.SourceFile;
   if (!SourceFile.empty()) {
     llvm::SmallString<64> AbsSourceFile;
@@ -380,6 +394,10 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
     sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestProtocolVersion);
     break;
 
+  case SourceKitRequest::DemangleNames:
+    prepareDemangleRequest(Req, Opts);
+    break;
+
   case SourceKitRequest::Index:
     sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestIndex);
     break;
@@ -642,6 +660,10 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
       KeepResponseAlive = true;
       break;
 
+    case SourceKitRequest::DemangleNames:
+      printDemangleResults(sourcekitd_response_get_value(Resp), outs());
+      break;
+
     case SourceKitRequest::ProtocolVersion:
     case SourceKitRequest::Index:
     case SourceKitRequest::CodeComplete:
@@ -1107,6 +1129,59 @@ static void printRelatedIdents(sourcekitd_variant_t Info, StringRef Filename,
   OS << "END RANGES\n";
 }
 
+static void prepareDemangleRequest(sourcekitd_object_t Req,
+                                   const TestOptions &Opts) {
+  sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestDemangle);
+  if (Opts.SimplifiedDemangling)
+    sourcekitd_request_dictionary_set_int64(Req, KeySimplified, 1);
+  sourcekitd_object_t arr = sourcekitd_request_array_create(nullptr, 0);
+  sourcekitd_request_dictionary_set_value(Req, KeyNames, arr);
+
+  auto addName = [&](StringRef MangledName) {
+    sourcekitd_request_array_set_stringbuf(arr, SOURCEKITD_ARRAY_APPEND,
+                                        MangledName.data(), MangledName.size());
+  };
+
+  if (Opts.Inputs.empty()) {
+    auto input = llvm::MemoryBuffer::getSTDIN();
+    if (!input) {
+      llvm::errs() << input.getError().message() << '\n';
+      ::exit(1);
+    }
+    llvm::StringRef inputContents = input.get()->getBuffer();
+
+    // This doesn't handle Unicode symbols, but maybe that's okay.
+    llvm::Regex maybeSymbol("_T[_a-zA-Z0-9$]+");
+    llvm::SmallVector<llvm::StringRef, 1> matches;
+    while (maybeSymbol.match(inputContents, &matches)) {
+      addName(matches.front());
+    }
+
+  } else {
+    for (llvm::StringRef name : Opts.Inputs) {
+      addName(name);
+    }
+  }
+
+  sourcekitd_request_release(arr);
+}
+
+static void printDemangleResults(sourcekitd_variant_t Info, raw_ostream &OS) {
+  OS << "START DEMANGLE\n";
+  sourcekitd_variant_t results =
+    sourcekitd_variant_dictionary_get_value(Info, KeyResults);
+  sourcekitd_variant_array_apply(results, ^bool(size_t index, sourcekitd_variant_t value) {
+    StringRef name = sourcekitd_variant_dictionary_get_string(value, KeyName);
+    if (name.empty())
+      OS << "<empty>";
+    else
+      OS << name;
+    OS << '\n';
+    return true;
+  });
+  OS << "END DEMANGLE\n";
+}
+
 static void initializeRewriteBuffer(StringRef Input,
                                     clang::RewriteBuffer &RewriteBuf) {
   RewriteBuf.Initialize(Input);

tools/SourceKit/tools/sourcekitd/lib/API/DictionaryKeys.h

@@ -101,6 +101,7 @@ extern SourceKit::UIdent KeyNextRequestStart;
 extern SourceKit::UIdent KeyPopular;
 extern SourceKit::UIdent KeyUnpopular;
 extern SourceKit::UIdent KeyHide;
+extern SourceKit::UIdent KeySimplified;
 
 extern SourceKit::UIdent KeyIsUnavailable;
 extern SourceKit::UIdent KeyIsDeprecated;

tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp

@@ -22,6 +22,8 @@
 #include "SourceKit/Support/Logging.h"
 #include "SourceKit/Support/UIdent.h"
 
+#include "swift/Basic/DemangleWrappers.h"
+
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
@@ -58,6 +60,9 @@ class LazySKDUID {
 
 static LazySKDUID RequestProtocolVersion("source.request.protocol_version");
 
+static LazySKDUID RequestDemangle("source.request.demangle");
+static LazySKDUID RequestMangleSimpleClass("source.request.mangle_simple_class");
+
 static LazySKDUID RequestIndex("source.request.indexsource");
 static LazySKDUID RequestDocInfo("source.request.docinfo");
 static LazySKDUID RequestCodeComplete("source.request.codecomplete");
@@ -139,6 +144,9 @@ static SourceKit::Context &getGlobalContext() {
   return *GlobalCtx;
 }
 
+static sourcekitd_response_t demangleNames(ArrayRef<const char *> MangledNames,
+                                           bool Simplified);
+
 static sourcekitd_response_t indexSource(StringRef Filename,
                                          ArrayRef<const char *> Args,
                                          StringRef KnownHash);
@@ -300,6 +308,18 @@ void handleRequestImpl(sourcekitd_object_t ReqObj, ResponseReceiver Rec) {
     return Rec(RB.createResponse());
   }
 
+  if (ReqUID == RequestDemangle) {
+    SmallVector<const char *, 8> MangledNames;
+    bool Failed = Req.getStringArray(KeyNames, MangledNames, /*isOptional=*/true);
+    if (Failed) {
+      return Rec(createErrorRequestInvalid(
+                                        "'key.names' not an array of strings"));
+    }
+    int64_t Simplified = false;
+    Req.getInt64(KeySimplified, Simplified, /*isOptional=*/true);
+    return Rec(demangleNames(MangledNames, Simplified));
+  }
+
   // Just accept 'source.request.buildsettings.register' for now, don't do
   // anything else.
   // FIXME: Heavy WIP here.
@@ -927,6 +947,44 @@ class SKDocConsumer : public DocInfoConsumer {
 };
 }
 
+static bool isSwiftPrefixed(StringRef MangledName) {
+  if (MangledName.size() < 2)
+    return false;
+  return (MangledName[0] == '_' && MangledName[1] == 'T');
+}
+
+static sourcekitd_response_t demangleNames(ArrayRef<const char *> MangledNames,
+                                           bool Simplified) {
+  swift::Demangle::DemangleOptions DemangleOptions;
+  if (Simplified) {
+    DemangleOptions =
+      swift::Demangle::DemangleOptions::SimplifiedUIDemangleOptions();
+  }
+
+  auto getDemangledName = [&](StringRef MangledName) -> std::string {
+    if (!isSwiftPrefixed(MangledName))
+      return std::string(); // Not a mangled name
+
+    std::string Result = swift::demangle_wrappers::demangleSymbolAsString(
+        MangledName, DemangleOptions);
+
+    if (Result == MangledName)
+      return std::string(); // Not a mangled name
+
+    return Result;
+  };
+
+  ResponseBuilder RespBuilder;
+  auto Arr = RespBuilder.getDictionary().setArray(KeyResults);
+  for (auto MangledName : MangledNames) {
+    std::string Result = getDemangledName(MangledName);
+    auto Entry = Arr.appendDictionary();
+    Entry.set(KeyName, Result.c_str());
+  }
+
+  return RespBuilder.createResponse();
+}
+
 static sourcekitd_response_t reportDocInfo(llvm::MemoryBuffer *InputBuf,
                                            StringRef ModuleName,
                                            ArrayRef<const char *> Args) {

tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-Common.cpp

@@ -114,6 +114,7 @@ UIdent sourcekitd::KeyNextRequestStart("key.nextrequeststart");
 UIdent sourcekitd::KeyPopular("key.popular");
 UIdent sourcekitd::KeyUnpopular("key.unpopular");
 UIdent sourcekitd::KeyHide("key.hide");
+UIdent sourcekitd::KeySimplified("key.simplified");
 
 UIdent sourcekitd::KeyIsDeprecated("key.is_deprecated");
 UIdent sourcekitd::KeyIsUnavailable("key.is_unavailable");