[SourceKit] Add a 'mangle simple classes' request, intended to supersede _swift_mangleSimpleClass.

rdar://23525281

Author: akyrtzi

test/SourceKit/Demangle/demangle.swift

@@ -11,3 +11,9 @@
 // SIMPLIFIED-NEXT: Builtin.Float80
 // SIMPLIFIED-NEXT: bar
 // SIMPLIFIED-NEXT: END DEMANGLE
+
+// RUN: %sourcekitd-test -req=mangle Foo.Baru Swift.Beer | FileCheck %s -check-prefix=MANGLED
+// MANGLED:      START MANGLE
+// MANGLED-NEXT: _TtC3Foo4Baru
+// MANGLED-NEXT: _TtCs4Beer
+// MANGLED-NEXT: END MANGLE

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

@@ -104,6 +104,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
       Request = llvm::StringSwitch<SourceKitRequest>(InputArg->getValue())
         .Case("version", SourceKitRequest::ProtocolVersion)
         .Case("demangle", SourceKitRequest::DemangleNames)
+        .Case("mangle", SourceKitRequest::MangleSimpleClasses)
         .Case("index", SourceKitRequest::Index)
         .Case("complete", SourceKitRequest::CodeComplete)
         .Case("complete.open", SourceKitRequest::CodeCompleteOpen)
@@ -132,7 +133,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/demangle/index/complete/cursor/related-idents/syntax-map/structure/"
+            << "version/demangle/mangle/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";

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

@@ -23,6 +23,7 @@ enum class SourceKitRequest {
   None,
   ProtocolVersion,
   DemangleNames,
+  MangleSimpleClasses,
   Index,
   CodeComplete,
   CodeCompleteOpen,

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

@@ -57,6 +57,9 @@ static void printModuleGroupNames(sourcekitd_variant_t Info,
 static void prepareDemangleRequest(sourcekitd_object_t Req,
                                    const TestOptions &Opts);
 static void printDemangleResults(sourcekitd_variant_t Info, raw_ostream &OS);
+static void prepareMangleRequest(sourcekitd_object_t Req,
+                                 const TestOptions &Opts);
+static void printMangleResults(sourcekitd_variant_t Info, raw_ostream &OS);
 
 static unsigned resolveFromLineCol(unsigned Line, unsigned Col,
                                    StringRef Filename);
@@ -116,6 +119,7 @@ static sourcekitd_uid_t KeySimplified;
 
 static sourcekitd_uid_t RequestProtocolVersion;
 static sourcekitd_uid_t RequestDemangle;
+static sourcekitd_uid_t RequestMangleSimpleClass;
 static sourcekitd_uid_t RequestIndex;
 static sourcekitd_uid_t RequestCodeComplete;
 static sourcekitd_uid_t RequestCodeCompleteOpen;
@@ -216,6 +220,7 @@ static int skt_main(int argc, const char **argv) {
 
   RequestProtocolVersion = sourcekitd_uid_get_from_cstr("source.request.protocol_version");
   RequestDemangle = sourcekitd_uid_get_from_cstr("source.request.demangle");
+  RequestMangleSimpleClass = sourcekitd_uid_get_from_cstr("source.request.mangle_simple_class");
   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");
@@ -354,7 +359,8 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
   if (Optargc < Args.size())
     Opts.CompilerArgs = Args.slice(Optargc+1);
 
-  if (Opts.Request == SourceKitRequest::DemangleNames)
+  if (Opts.Request == SourceKitRequest::DemangleNames ||
+      Opts.Request == SourceKitRequest::MangleSimpleClasses)
     Opts.SourceFile.clear();
 
   std::string SourceFile = Opts.SourceFile;
@@ -400,6 +406,10 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
     prepareDemangleRequest(Req, Opts);
     break;
 
+  case SourceKitRequest::MangleSimpleClasses:
+    prepareMangleRequest(Req, Opts);
+    break;
+
   case SourceKitRequest::Index:
     sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestIndex);
     break;
@@ -668,6 +678,10 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
       printDemangleResults(sourcekitd_response_get_value(Resp), outs());
       break;
 
+    case SourceKitRequest::MangleSimpleClasses:
+      printMangleResults(sourcekitd_response_get_value(Resp), outs());
+      break;
+
     case SourceKitRequest::ProtocolVersion:
     case SourceKitRequest::Index:
     case SourceKitRequest::CodeComplete:
@@ -1170,6 +1184,37 @@ static void prepareDemangleRequest(sourcekitd_object_t Req,
   sourcekitd_request_release(arr);
 }
 
+static void prepareMangleRequest(sourcekitd_object_t Req,
+                                 const TestOptions &Opts) {
+  sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestMangleSimpleClass);
+  sourcekitd_object_t arr = sourcekitd_request_array_create(nullptr, 0);
+  sourcekitd_request_dictionary_set_value(Req, KeyNames, arr);
+
+  auto addPair = [&](StringRef ModuleName, StringRef ClassName) {
+    sourcekitd_object_t pair =
+      sourcekitd_request_dictionary_create(nullptr, nullptr, 0);
+    sourcekitd_request_dictionary_set_stringbuf(pair, KeyModuleName,
+                                          ModuleName.data(), ModuleName.size());
+    sourcekitd_request_dictionary_set_stringbuf(pair, KeyName,
+                                            ClassName.data(), ClassName.size());
+    sourcekitd_request_array_set_value(arr, SOURCEKITD_ARRAY_APPEND, pair);
+    sourcekitd_request_release(pair);
+  };
+
+  for (StringRef pair : Opts.Inputs) {
+    auto Idx = pair.find('.');
+    if (Idx == StringRef::npos) {
+      errs() << "expected pairs with format '<module>.<class name>'\n";
+      ::exit(1);
+    }
+    StringRef moduleName = pair.substr(0, Idx);
+    StringRef className = pair.substr(Idx+1);
+    addPair(moduleName, className);
+  }
+
+  sourcekitd_request_release(arr);
+}
+
 static void printDemangleResults(sourcekitd_variant_t Info, raw_ostream &OS) {
   OS << "START DEMANGLE\n";
   sourcekitd_variant_t results =
@@ -1186,6 +1231,22 @@ static void printDemangleResults(sourcekitd_variant_t Info, raw_ostream &OS) {
   OS << "END DEMANGLE\n";
 }
 
+static void printMangleResults(sourcekitd_variant_t Info, raw_ostream &OS) {
+  OS << "START MANGLE\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 MANGLE\n";
+}
+
 static void initializeRewriteBuffer(StringRef Input,
                                     clang::RewriteBuffer &RewriteBuf) {
   RewriteBuf.Initialize(Input);

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

@@ -61,6 +61,7 @@ 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");
@@ -146,6 +147,9 @@ static SourceKit::Context &getGlobalContext() {
 static sourcekitd_response_t demangleNames(ArrayRef<const char *> MangledNames,
                                            bool Simplified);
 
+static sourcekitd_response_t
+mangleSimpleClassNames(ArrayRef<std::pair<StringRef, StringRef>> ModuleClassPairs);
+
 static sourcekitd_response_t indexSource(StringRef Filename,
                                          ArrayRef<const char *> Args,
                                          StringRef KnownHash);
@@ -321,6 +325,33 @@ void handleRequestImpl(sourcekitd_object_t ReqObj, ResponseReceiver Rec) {
     return Rec(demangleNames(MangledNames, Simplified));
   }
 
+  if (ReqUID == RequestMangleSimpleClass) {
+    SmallVector<std::pair<StringRef, StringRef>, 16> ModuleClassPairs;
+    sourcekitd_response_t err = nullptr;
+    bool failed = Req.dictionaryArrayApply(KeyNames, [&](RequestDict dict) {
+      Optional<StringRef> ModuleName = dict.getString(KeyModuleName);
+      if (!ModuleName.hasValue()) {
+        err = createErrorRequestInvalid("missing 'key.modulename'");
+        return true;
+      }
+      Optional<StringRef> ClassName = dict.getString(KeyName);
+      if (!ClassName.hasValue()) {
+        err = createErrorRequestInvalid("missing 'key.name'");
+        return true;
+      }
+      ModuleClassPairs.push_back(std::make_pair(*ModuleName, *ClassName));
+      return false;
+    });
+
+    if (failed) {
+      if (!err)
+        err = createErrorRequestInvalid("missing 'key.names'");
+      return Rec(err);
+    }
+
+    return Rec(mangleSimpleClassNames(ModuleClassPairs));
+  }
+
   // Just accept 'source.request.buildsettings.register' for now, don't do
   // anything else.
   // FIXME: Heavy WIP here.
@@ -994,6 +1025,37 @@ static sourcekitd_response_t demangleNames(ArrayRef<const char *> MangledNames,
   return RespBuilder.createResponse();
 }
 
+static std::string mangleSimpleClass(StringRef moduleName,
+                                     StringRef className) {
+  using namespace swift::Demangle;
+
+  auto moduleNode = NodeFactory::create(Node::Kind::Module, moduleName);
+  auto IdNode = NodeFactory::create(Node::Kind::Identifier, className);
+  auto classNode = NodeFactory::create(Node::Kind::Class);
+  auto typeNode = NodeFactory::create(Node::Kind::Type);
+  auto typeManglingNode = NodeFactory::create(Node::Kind::TypeMangling);
+  auto globalNode = NodeFactory::create(Node::Kind::Global);
+
+  classNode->addChildren(moduleNode, IdNode);
+  typeNode->addChild(classNode);
+  typeManglingNode->addChild(typeNode);
+  globalNode->addChild(typeManglingNode);
+  return mangleNode(globalNode);
+}
+
+static sourcekitd_response_t
+mangleSimpleClassNames(ArrayRef<std::pair<StringRef, StringRef>> ModuleClassPairs) {
+  ResponseBuilder RespBuilder;
+  auto Arr = RespBuilder.getDictionary().setArray(KeyResults);
+  for (auto &pair : ModuleClassPairs) {
+    std::string Result = mangleSimpleClass(pair.first, pair.second);
+    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) {