[SourceKit] Allow cursorinfo to take a USR instead of an offset

This eventually calls the code from ReconstructType to try to find the
Decl for a USR. For now, only works in a file, not a generated
interface.

rdar://problem/25017817

test/SourceKit/SourceDocInfo/cursor_usr.swift

@@ -0,0 +1,57 @@
+// The RUN lines are at the bottom in case we ever need to rely on line:col info.
+import Foo
+import FooSwiftModule
+
+var global: Int
+
+struct S1 {}
+
+func foo(x: FooStruct1) -> S1 {}
+
+// RUN: rm -rf %t
+// RUN: mkdir %t
+// RUN: %swiftc_driver -emit-module -o %t/FooSwiftModule.swiftmodule %S/Inputs/FooSwiftModule.swift
+
+// Sanity check that we have identical responses when things work.
+// RUN: %sourcekitd-test -req=cursor -pos=5:5 %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s > %t.from_offset.txt
+// RUN: %sourcekitd-test -req=cursor -usr "s:v10cursor_usr6globalSi" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s > %t.from_usr.txt
+// RUN: FileCheck %s -check-prefix=CHECK_SANITY1 < %t.from_offset.txt
+// RUN: FileCheck %s -check-prefix=CHECK_SANITY1 < %t.from_usr.txt
+// RUN: diff -u %t.from_usr.txt %t.from_offset.txt
+// CHECK_SANITY1: source.lang.swift.decl.var.global (5:5-5:11)
+// CHECK_SANITY1-NEXT: global
+// CHECK_SANITY1-NEXT: s:v10cursor_usr6globalSi
+// CHECK_SANITY1-NEXT: Int
+// CHECK_SANITY1-NEXT: <Declaration>var global: <Type usr="s:Si">Int</Type></Declaration>
+// CHECK_SANITY1-NEXT: <decl.var.global><syntaxtype.keyword>var</syntaxtype.keyword> <decl.name>global</decl.name>: <decl.var.type><ref.struct usr="s:Si">Int</ref.struct></decl.var.type></decl.var.global>
+
+// Bogus USR.
+// RUN: %sourcekitd-test -req=cursor -usr "s:blahblahblah" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s | FileCheck %s -check-prefix=EMPTY
+// Missing s: prefix.
+// RUN: %sourcekitd-test -req=cursor -usr "v10cursor_usr6globalSi" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s | FileCheck %s -check-prefix=EMPTY
+// FIXME: no support for clang USRs.
+// RUN: %sourcekitd-test -req=cursor -usr "c:@S@FooStruct1" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s | FileCheck %s -check-prefix=EMPTY
+// EMPTY: <empty cursor info>
+
+// FIXME: missing symbol shows up as some other part of the USR (the type here).
+// RUN: %sourcekitd-test -req=cursor -usr "s:v10cursor_usr11global_noneSi" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s | FileCheck %s -check-prefix=SHOULD_BE_EMPTY
+// SHOULD_BE_EMPTY: source.lang.swift.decl.struct ()
+// SHOULD_BE_EMPTY: Int
+
+// RUN: %sourcekitd-test -req=cursor -usr "s:V10cursor_usr2S1" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s | FileCheck %s -check-prefix=CHECK1
+// CHECK1: source.lang.swift.decl.struct (7:8-7:10)
+// CHECK1: S1
+// CHECK1: <decl.struct><syntaxtype.keyword>struct</syntaxtype.keyword> <decl.name>S1</decl.name></decl.struct>
+
+// RUN: %sourcekitd-test -req=cursor -usr "s:F14FooSwiftModule12fooSwiftFuncFT_Si" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s | FileCheck %s -check-prefix=CHECK2
+// CHECK2: source.lang.swift.decl.function.free ()
+// CHECK2: fooSwiftFunc()
+// CHECK2: () -> Int
+// CHECK2: FooSwiftModule
+// CHECK2: <decl.function.free><syntaxtype.keyword>func</syntaxtype.keyword> <decl.name>fooSwiftFunc</decl.name>() -&gt; <decl.function.returntype><ref.struct usr="s:Si">Int</ref.struct></decl.function.returntype></decl.function.free>
+
+// RUN: %sourcekitd-test -req=cursor -usr "s:F10cursor_usr3fooFVSC10FooStruct1VS_2S1" %s -- -I %t -F %S/../Inputs/libIDE-mock-sdk %mcp_opt %s | FileCheck %s -check-prefix=CHECK3
+// CHECK3: source.lang.swift.decl.function.free (9:6-9:24)
+// CHECK3: foo(_:)
+// CHECK3: (FooStruct1) -> S1
+// CHECK3: <decl.function.free><syntaxtype.keyword>func</syntaxtype.keyword> <decl.name>foo</decl.name>(<decl.var.parameter><decl.var.parameter.name>x</decl.var.parameter.name>: <decl.var.parameter.type><ref.struct usr="c:@S@FooStruct1">FooStruct1</ref.struct></decl.var.parameter.type></decl.var.parameter>) -&gt; <decl.function.returntype><ref.struct usr="s:V10cursor_usr2S1">S1</ref.struct></decl.function.returntype></decl.function.free>

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -458,6 +458,11 @@ public:
                              ArrayRef<const char *> Args,
                           std::function<void(const CursorInfo &)> Receiver) = 0;
 
+  virtual void
+  getCursorInfoFromUSR(StringRef Filename, StringRef USR,
+                       ArrayRef<const char *> Args,
+                       std::function<void(const CursorInfo &)> Receiver) = 0;
+
   virtual void findRelatedIdentifiersInFile(StringRef Filename,
                                             unsigned Offset,
                                             ArrayRef<const char *> Args,

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -354,6 +354,10 @@ public:
                      ArrayRef<const char *> Args,
                      std::function<void(const CursorInfo &)> Receiver) override;
 
+  void getCursorInfoFromUSR(
+      StringRef Filename, StringRef USR, ArrayRef<const char *> Args,
+      std::function<void(const CursorInfo &)> Receiver) override;
+
   void findRelatedIdentifiersInFile(StringRef Filename, unsigned Offset,
                                     ArrayRef<const char *> Args,
               std::function<void(const RelatedIdentsInfo &)> Receiver) override;

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -994,6 +994,151 @@ void SwiftLangSupport::getCursorInfo(
                 Receiver);
 }
 
+static void
+resolveCursorFromUSR(SwiftLangSupport &Lang, StringRef InputFile, StringRef USR,
+                     SwiftInvocationRef Invok, bool TryExistingAST,
+                     std::function<void(const CursorInfo &)> Receiver) {
+  assert(Invok);
+
+  class CursorInfoConsumer : public SwiftASTConsumer {
+    std::string InputFile;
+    StringRef USR;
+    SwiftLangSupport &Lang;
+    SwiftInvocationRef ASTInvok;
+    const bool TryExistingAST;
+    std::function<void(const CursorInfo &)> Receiver;
+    SmallVector<ImmutableTextSnapshotRef, 4> PreviousASTSnaps;
+
+  public:
+    CursorInfoConsumer(StringRef InputFile, StringRef USR,
+                       SwiftLangSupport &Lang, SwiftInvocationRef ASTInvok,
+                       bool TryExistingAST,
+                       std::function<void(const CursorInfo &)> Receiver)
+        : InputFile(InputFile), USR(USR), Lang(Lang),
+          ASTInvok(std::move(ASTInvok)), TryExistingAST(TryExistingAST),
+          Receiver(std::move(Receiver)) {}
+
+    bool canUseASTWithSnapshots(
+        ArrayRef<ImmutableTextSnapshotRef> Snapshots) override {
+      if (!TryExistingAST) {
+        LOG_INFO_FUNC(High, "will resolve using up-to-date AST");
+        return false;
+      }
+
+      if (!Snapshots.empty()) {
+        PreviousASTSnaps.append(Snapshots.begin(), Snapshots.end());
+        LOG_INFO_FUNC(High, "will try existing AST");
+        return true;
+      }
+
+      LOG_INFO_FUNC(High, "will resolve using up-to-date AST");
+      return false;
+    }
+
+    void handlePrimaryAST(ASTUnitRef AstUnit) override {
+      auto &CompIns = AstUnit->getCompilerInstance();
+      Module *MainModule = CompIns.getMainModule();
+
+      unsigned BufferID =
+          AstUnit->getPrimarySourceFile().getBufferID().getValue();
+
+      trace::TracedOperation TracedOp;
+      if (trace::enabled()) {
+        trace::SwiftInvocation SwiftArgs;
+        ASTInvok->raw(SwiftArgs.Args.Args, SwiftArgs.Args.PrimaryFile);
+        trace::initTraceFiles(SwiftArgs, CompIns);
+        TracedOp.start(trace::OperationKind::CursorInfoForSource, SwiftArgs,
+                       {std::make_pair("USR", USR)});
+      }
+
+      std::string mangledName(USR);
+      if (USR.startswith("s:")) {
+        mangledName.replace(0, 2, "_T");
+      } else 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);
+
+      if (!D) {
+        Receiver({});
+        return;
+      }
+
+      CompilerInvocation CompInvok;
+      ASTInvok->applyTo(CompInvok);
+
+      if (auto *M = dyn_cast<ModuleDecl>(D)) {
+        passCursorInfoForModule(M, Lang.getIFaceGenContexts(), CompInvok,
+                                Receiver);
+      } else if (auto *VD = dyn_cast<ValueDecl>(D)) {
+        bool Failed =
+            passCursorInfoForDecl(VD, MainModule, VD->getType(),
+                                  /*isRef=*/false, BufferID, Lang, CompInvok,
+                                  PreviousASTSnaps, Receiver);
+        if (Failed) {
+          if (!PreviousASTSnaps.empty()) {
+            // Attempt again using the up-to-date AST.
+            resolveCursorFromUSR(Lang, InputFile, USR, ASTInvok,
+                                 /*TryExistingAST=*/false, Receiver);
+          } else {
+            Receiver({});
+          }
+        }
+      }
+    }
+
+    void cancelled() override {
+      CursorInfo Info;
+      Info.IsCancelled = true;
+      Receiver(Info);
+    }
+
+    void failed(StringRef Error) override {
+      LOG_WARN_FUNC("cursor info failed: " << Error);
+      Receiver({});
+    }
+  };
+
+  auto Consumer = std::make_shared<CursorInfoConsumer>(
+      InputFile, USR, Lang, Invok, TryExistingAST, Receiver);
+  /// FIXME: When request cancellation is implemented and Xcode adopts it,
+  /// don't use 'OncePerASTToken'.
+  static const char OncePerASTToken = 0;
+  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer),
+                                       &OncePerASTToken);
+}
+
+void SwiftLangSupport::getCursorInfoFromUSR(
+    StringRef filename, StringRef USR, ArrayRef<const char *> args,
+    std::function<void(const CursorInfo &)> receiver) {
+  if (auto IFaceGenRef = IFaceGenContexts.get(filename)) {
+    LOG_WARN_FUNC("info from usr for generated interface not implemented yet");
+    receiver({});
+    return;
+  }
+
+  std::string error;
+  SwiftInvocationRef invok = ASTMgr->getInvocation(args, filename, error);
+  if (!invok) {
+    // FIXME: Report it as failed request.
+    LOG_WARN_FUNC("failed to create an ASTInvocation: " << error);
+    receiver({});
+    return;
+  }
+
+  resolveCursorFromUSR(*this, filename, USR, invok, /*TryExistingAST=*/true,
+                       receiver);
+}
+
 //===----------------------------------------------------------------------===//
 // SwiftLangSupport::findUSRRange
 //===----------------------------------------------------------------------===//

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

@@ -384,6 +384,7 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
           getBufferForFilename(SourceFile)->getBuffer(), SourceFile);
   }
 
+  // FIXME: we should detect if offset is required but not set.
   unsigned ByteOffset = Opts.Offset;
   if (Opts.Line != 0) {
     ByteOffset = resolveFromLineCol(Opts.Line, Opts.Col, SourceBuf.get());
@@ -486,7 +487,11 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
 
   case SourceKitRequest::CursorInfo:
     sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestCursorInfo);
-    sourcekitd_request_dictionary_set_int64(Req, KeyOffset, ByteOffset);
+    if (!Opts.USR.empty()) {
+      sourcekitd_request_dictionary_set_string(Req, KeyUSR, Opts.USR.c_str());
+    } else {
+      sourcekitd_request_dictionary_set_int64(Req, KeyOffset, ByteOffset);
+    }
     break;
 
   case SourceKitRequest::RelatedIdents:

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

@@ -701,13 +701,22 @@ handleSemanticRequest(RequestDict Req,
       return Rec(createErrorRequestFailed("semantic editor is disabled"));
 
   if (ReqUID == RequestCursorInfo) {
-    int64_t Offset;
-    if (Req.getInt64(KeyOffset, Offset, /*isOptional=*/false))
-      return Rec(createErrorRequestInvalid("missing 'key.offset'"));
     LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
-    return Lang.getCursorInfo(
+
-        *SourceFile, Offset, Args,
+    int64_t Offset;
-        [Rec](const CursorInfo &Info) { reportCursorInfo(Info, Rec); });
+    if (!Req.getInt64(KeyOffset, Offset, /*isOptional=*/false)) {
+      return Lang.getCursorInfo(
+          *SourceFile, Offset, Args,
+          [Rec](const CursorInfo &Info) { reportCursorInfo(Info, Rec); });
+    }
+    if (auto USR = Req.getString(KeyUSR)) {
+      return Lang.getCursorInfoFromUSR(
+          *SourceFile, *USR, Args,
+          [Rec](const CursorInfo &Info) { reportCursorInfo(Info, Rec); });
+    }
+
+    return Rec(createErrorRequestInvalid(
+        "either 'key.offset' or 'key.usr' is required"));
   }
 
   if (ReqUID == RequestRelatedIdents) {