[IDE] Factor code that looks up USRs into getDeclFromUSR

Author: benlangmuir

include/swift/IDE/Utils.h

@@ -125,6 +125,7 @@ void getLocationInfoForClangNode(ClangNode ClangNode,
 
 Optional<std::pair<unsigned, unsigned>> parseLineCol(StringRef LineCol);
 
+Decl *getDeclFromUSR(ASTContext &context, StringRef USR, std::string &error);
 Decl *getDeclFromMangledSymbolName(ASTContext &context, StringRef mangledName,
                                    std::string &error);
 

lib/IDE/ReconstructType.cpp

@@ -2375,6 +2375,22 @@ VisitNode(ASTContext *ast, std::vector<Demangle::NodePointer> &nodes,
   nodes.pop_back();
 }
 
+Decl *ide::getDeclFromUSR(ASTContext &context, StringRef USR, std::string &error) {
+  if (!USR.startswith("s:")) {
+    error = "not a Swift USR";
+    return nullptr;
+  }
+
+  std::string mangledName(USR);
+  // Convert to a symbol name by replacing the USR prefix.
+  // This relies on USR generation being very close to symbol mangling; if we
+  // need to support entities with customized USRs (e.g. extensions), we will
+  // need to do something smarter here.
+  mangledName.replace(0, 2, "_T");
+
+  return getDeclFromMangledSymbolName(context, mangledName, error);
+}
+
 Decl *ide::getDeclFromMangledSymbolName(ASTContext &context,
                                         StringRef mangledName,
                                         std::string &error) {

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -1051,22 +1051,15 @@ resolveCursorFromUSR(SwiftLangSupport &Lang, StringRef InputFile, StringRef USR,
                        {std::make_pair("USR", USR)});
       }
 
-      std::string mangledName(USR);
-      if (USR.startswith("s:")) {
-        mangledName.replace(0, 2, "_T");
-      } else if (USR.startswith("c:")) {
+      if (USR.startswith("c:")) {
         LOG_WARN_FUNC("lookup for C/C++/ObjC USRs not implemented");
         Receiver({});
         return;
-      } else if (!USR.startswith("_T")) {
-        LOG_WARN_FUNC("unknown USR prefix");
-        Receiver({});
-        return;
       }
 
       auto &context = CompIns.getASTContext();
       std::string error;
-      Decl *D = ide::getDeclFromMangledSymbolName(context, mangledName, error);
+      Decl *D = ide::getDeclFromUSR(context, USR, error);
 
       if (!D) {
         Receiver({});

tools/swift-ide-test/swift-ide-test.cpp

@@ -2423,19 +2423,14 @@ class TypeReconstructWalker : public SourceEntityWalker {
 
 private:
   void tryDemangleDecl(ValueDecl *VD, CharSourceRange range) {
-    std::string mangledName;
+    std::string USR;
     {
-      llvm::raw_string_ostream OS(mangledName);
+      llvm::raw_string_ostream OS(USR);
       printDeclUSR(VD, OS);
     }
 
-    // Put the expected symbol _T prefix on the name by replacing the s:.
-    assert(StringRef(mangledName).startswith("s:"));
-    mangledName[0] = '_';
-    mangledName[1] = 'T';
-
     std::string error;
-    if (Decl *reDecl = getDeclFromMangledSymbolName(Ctx, mangledName, error)) {
+    if (Decl *reDecl = getDeclFromUSR(Ctx, USR, error)) {
       Stream << "reconstructed decl from usr for '" << range.str() << "' is '";
       reDecl->print(Stream, PrintOptions());
       Stream << "'\n";