[SourceKit] Add a request to get the semantic tokens of a file

We need this request for semantic highlighting in LSP. Previously, we were getting the semantic tokens using a 0,0 edit after a document update notification but that will no longer be possible if we open the documents in syntactic only mode.

test/SourceKit/SemanticTokens/semantic_tokens.swift

@@ -0,0 +1,17 @@
+// RUN: %sourcekitd-test -req=semantic-tokens %s -- %s | %FileCheck %s
+
+class MyClass {}
+
+let x: String = "test"
+var y = MyClass()
+
+
+// String in line 5
+// CHECK: key.kind: source.lang.swift.ref.struct
+// CHECK: key.length: 6,
+// CHECK key.is_system: 1
+// CHECK-LABEL: },
+    
+// MyClass in line 6
+// CHECK: key.kind: source.lang.swift.ref.class,
+// CHECK: key.length: 7

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

@@ -17,6 +17,7 @@
 #include "SourceKit/Support/CancellationToken.h"
 #include "SourceKit/Support/UIdent.h"
 #include "swift/AST/Type.h"
+#include "swift/IDE/CodeCompletionResult.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/Optional.h"
@@ -330,6 +331,35 @@ struct DiagnosticEntryInfo : DiagnosticEntryInfoBase {
   SmallVector<DiagnosticEntryInfoBase, 1> Notes;
 };
 
+struct SwiftSemanticToken {
+  unsigned ByteOffset;
+  unsigned Length : 24;
+  // The code-completion kinds are a good match for the semantic kinds we want.
+  // FIXME: Maybe rename CodeCompletionDeclKind to a more general concept ?
+  swift::ide::CodeCompletionDeclKind Kind : 6;
+  unsigned IsRef : 1;
+  unsigned IsSystem : 1;
+
+  SwiftSemanticToken(swift::ide::CodeCompletionDeclKind Kind,
+                     unsigned ByteOffset, unsigned Length, bool IsRef,
+                     bool IsSystem)
+      : ByteOffset(ByteOffset), Length(Length), Kind(Kind), IsRef(IsRef),
+        IsSystem(IsSystem) {}
+
+  bool getIsRef() const { return static_cast<bool>(IsRef); }
+
+  bool getIsSystem() const { return static_cast<bool>(IsSystem); }
+
+  UIdent getUIdentForKind();
+};
+
+#if !defined(_MSC_VER)
+static_assert(sizeof(SwiftSemanticToken) == 8, "Too big");
+// FIXME: MSVC doesn't pack bitfields with different underlying types.
+// Giving up to check this in MSVC for now, because static_assert is only for
+// keeping low memory usage.
+#endif
+
 struct SourceFileRange {
   /// The byte offset at which the range begins
   uintptr_t Start;
@@ -709,6 +739,9 @@ struct CursorInfoData {
 /// The result type of `LangSupport::getDiagnostics`
 typedef ArrayRef<DiagnosticEntryInfo> DiagnosticsResult;
 
+/// The result of `LangSupport::getSemanticTokens`.
+typedef std::vector<SwiftSemanticToken> SemanticTokensResult;
+
 struct RangeInfo {
   UIdent RangeKind;
   StringRef ExprType;
@@ -1103,6 +1136,13 @@ public:
                  std::function<void(const RequestResult<DiagnosticsResult> &)>
                      Receiver) = 0;
 
+  virtual void getSemanticTokens(
+      StringRef PrimaryFilePath, StringRef InputBufferName,
+      ArrayRef<const char *> Args, llvm::Optional<VFSOptions> VfsOptions,
+      SourceKitCancellationToken CancellationToken,
+      std::function<void(const RequestResult<SemanticTokensResult> &)>
+          Receiver) = 0;
+
   virtual void
   getNameInfo(StringRef PrimaryFilePath, StringRef InputBufferName,
               unsigned Offset, NameTranslatingInfo &Input,

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -643,36 +643,6 @@ struct EditorConsumerSyntaxMapEntry {
     :Offset(Offset), Length(Length), Kind(Kind) { }
 };
 
-struct SwiftSemanticToken {
-  unsigned ByteOffset;
-  unsigned Length : 24;
-  // The code-completion kinds are a good match for the semantic kinds we want.
-  // FIXME: Maybe rename CodeCompletionDeclKind to a more general concept ?
-  CodeCompletionDeclKind Kind : 6;
-  unsigned IsRef : 1;
-  unsigned IsSystem : 1;
-
-  SwiftSemanticToken(CodeCompletionDeclKind Kind,
-                     unsigned ByteOffset, unsigned Length,
-                     bool IsRef, bool IsSystem)
-    : ByteOffset(ByteOffset), Length(Length), Kind(Kind),
-      IsRef(IsRef), IsSystem(IsSystem) { }
-
-  bool getIsRef() const { return static_cast<bool>(IsRef); }
-
-  bool getIsSystem() const { return static_cast<bool>(IsSystem); }
-
-  UIdent getUIdentForKind() const {
-    return SwiftLangSupport::getUIDForCodeCompletionDeclKind(Kind, getIsRef());
-  }
-};
-#if !defined(_MSC_VER)
-static_assert(sizeof(SwiftSemanticToken) == 8, "Too big");
-// FIXME: MSVC doesn't pack bitfields with different underlying types.
-// Giving up to check this in MSVC for now, because static_assert is only for
-// keeping low memory usage.
-#endif
-
 class SwiftDocumentSemanticInfo :
     public ThreadSafeRefCountedBase<SwiftDocumentSemanticInfo> {
 
@@ -2535,3 +2505,68 @@ void SwiftLangSupport::editorExpandPlaceholder(StringRef Name, unsigned Offset,
   }
   EditorDoc->expandPlaceholder(Offset, Length, Consumer);
 }
+
+//===----------------------------------------------------------------------===//
+// Semantic Tokens
+//===----------------------------------------------------------------------===//
+
+void SwiftLangSupport::getSemanticTokens(
+    StringRef PrimaryFilePath, StringRef InputBufferName,
+    ArrayRef<const char *> Args, llvm::Optional<VFSOptions> VfsOptions,
+    SourceKitCancellationToken CancellationToken,
+    std::function<void(const RequestResult<SemanticTokensResult> &)> Receiver) {
+  std::string FileSystemError;
+  auto FileSystem = getFileSystem(VfsOptions, PrimaryFilePath, FileSystemError);
+  if (!FileSystem) {
+    Receiver(RequestResult<SemanticTokensResult>::fromError(FileSystemError));
+    return;
+  }
+
+  std::string InvocationError;
+  SwiftInvocationRef Invok = ASTMgr->getTypecheckInvocation(
+      Args, PrimaryFilePath, FileSystem, InvocationError);
+  if (!InvocationError.empty()) {
+    LOG_WARN_FUNC("error creating ASTInvocation: " << InvocationError);
+  }
+  if (!Invok) {
+    Receiver(RequestResult<SemanticTokensResult>::fromError(InvocationError));
+    return;
+  }
+
+  class SemanticTokensConsumer : public SwiftASTConsumer {
+    StringRef InputBufferName;
+    std::function<void(const RequestResult<SemanticTokensResult> &)> Receiver;
+
+  public:
+    SemanticTokensConsumer(
+        StringRef InputBufferName,
+        std::function<void(const RequestResult<SemanticTokensResult> &)>
+            Receiver)
+        : InputBufferName(InputBufferName), Receiver(Receiver) {}
+
+    void handlePrimaryAST(ASTUnitRef AstUnit) override {
+      auto &CompIns = AstUnit->getCompilerInstance();
+      SourceFile *SF = retrieveInputFile(InputBufferName, CompIns);
+      if (!SF) {
+        Receiver(RequestResult<SemanticTokensResult>::fromError(
+            "Unable to find input file"));
+        return;
+      }
+      SemanticAnnotator Annotator(CompIns.getSourceMgr(), *SF->getBufferID());
+      Annotator.walk(SF);
+      Receiver(
+          RequestResult<SemanticTokensResult>::fromResult(Annotator.SemaToks));
+    }
+
+    void cancelled() override {
+      Receiver(RequestResult<SemanticTokensResult>::cancelled());
+    }
+  };
+
+  auto Consumer = std::make_shared<SemanticTokensConsumer>(InputBufferName,
+                                                           std::move(Receiver));
+
+  getASTManager()->processASTAsync(Invok, std::move(Consumer),
+                                   /*OncePerASTToken=*/nullptr,
+                                   CancellationToken, FileSystem);
+}

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp

@@ -397,6 +397,10 @@ UIdent SwiftLangSupport::getUIDForRefactoringKind(ide::RefactoringKind Kind){
   }
 }
 
+UIdent SwiftSemanticToken::getUIdentForKind() {
+  return SwiftLangSupport::getUIDForCodeCompletionDeclKind(Kind, IsRef);
+}
+
 UIdent SwiftLangSupport::getUIDForCodeCompletionDeclKind(
     ide::CodeCompletionDeclKind Kind, bool IsRef) {
   if (IsRef) {

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -665,6 +665,13 @@ public:
                  std::function<void(const RequestResult<DiagnosticsResult> &)>
                      Receiver) override;
 
+  void getSemanticTokens(
+      StringRef PrimaryFilePath, StringRef InputBufferName,
+      ArrayRef<const char *> Args, llvm::Optional<VFSOptions> VfsOptions,
+      SourceKitCancellationToken CancellationToken,
+      std::function<void(const RequestResult<SemanticTokensResult> &)> Receiver)
+      override;
+
   void getNameInfo(
       StringRef PrimaryFilePath, StringRef InputBufferName, unsigned Offset,
       NameTranslatingInfo &Input, ArrayRef<const char *> Args,
@@ -812,6 +819,10 @@ public:
 /// Disable expensive SIL options which do not affect indexing or diagnostics.
 void disableExpensiveSILOptions(swift::SILOptions &Opts);
 
+swift::SourceFile *retrieveInputFile(StringRef inputBufferName,
+                                     const swift::CompilerInstance &CI,
+                                     bool haveRealPath = false);
+
 } // namespace SourceKit
 
 #endif

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -1537,9 +1537,9 @@ public:
   }
 };
 
-static SourceFile *retrieveInputFile(StringRef inputBufferName,
+SourceFile *SourceKit::retrieveInputFile(StringRef inputBufferName,
                                          const CompilerInstance &CI,
-                                     bool haveRealPath = false) {
+                                         bool haveRealPath) {
   // Don't bother looking up if we have the same file as the primary file or
   // we weren't given a separate input file
   if (inputBufferName.empty() ||

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

@@ -151,6 +151,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
         .Case("global-config", SourceKitRequest::GlobalConfiguration)
         .Case("dependency-updated", SourceKitRequest::DependencyUpdated)
         .Case("diags", SourceKitRequest::Diagnostics)
+        .Case("semantic-tokens", SourceKitRequest::SemanticTokens)
         .Case("compile", SourceKitRequest::Compile)
         .Case("compile.close", SourceKitRequest::CompileClose)
         .Case("syntactic-expandmacro", SourceKitRequest::SyntacticMacroExpansion)

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

@@ -68,6 +68,7 @@ enum class SourceKitRequest {
   GlobalConfiguration,
   DependencyUpdated,
   Diagnostics,
+  SemanticTokens,
   Compile,
   CompileClose,
   SyntacticMacroExpansion,

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

@@ -1137,6 +1137,10 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
   case SourceKitRequest::Diagnostics:
     sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestDiagnostics);
     break;
+  case SourceKitRequest::SemanticTokens:
+    sourcekitd_request_dictionary_set_uid(Req, KeyRequest,
+                                          RequestSemanticTokens);
+    break;
 
   case SourceKitRequest::Compile:
     sourcekitd_request_dictionary_set_string(Req, KeyName, SemaName.c_str());
@@ -1427,6 +1431,7 @@ static bool handleResponse(sourcekitd_response_t Resp, const TestOptions &Opts,
     case SourceKitRequest::ConformingMethodList:
     case SourceKitRequest::DependencyUpdated:
     case SourceKitRequest::Diagnostics:
+    case SourceKitRequest::SemanticTokens:
       printRawResponse(Resp);
       break;
     case SourceKitRequest::Compile:

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

@@ -194,6 +194,10 @@ static void reportCursorInfo(const RequestResult<CursorInfoData> &Result, Respon
 static void reportDiagnostics(const RequestResult<DiagnosticsResult> &Result,
                               ResponseReceiver Rec);
 
+static void
+reportSemanticTokens(const RequestResult<SemanticTokensResult> &Result,
+                     ResponseReceiver Rec);
+
 static void reportExpressionTypeInfo(const RequestResult<ExpressionTypesInFile> &Result,
                                      ResponseReceiver Rec);
 
@@ -1866,6 +1870,32 @@ handleRequestDiagnostics(const RequestDict &Req,
   });
 }
 
+static void
+handleRequestSemanticTokens(const RequestDict &Req,
+                            SourceKitCancellationToken CancellationToken,
+                            ResponseReceiver Rec) {
+  handleSemanticRequest(Req, Rec, [Req, CancellationToken, Rec]() {
+    Optional<VFSOptions> vfsOptions = getVFSOptions(Req);
+    auto PrimaryFilePath = getPrimaryFilePathForRequestOrEmitError(Req, Rec);
+    if (!PrimaryFilePath)
+      return;
+    StringRef InputBufferName = getInputBufferNameForRequest(Req, Rec);
+
+    SmallVector<const char *, 8> Args;
+    if (getCompilerArgumentsForRequestOrEmitError(Req, Args, Rec))
+      return;
+
+    LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
+    Lang.getSemanticTokens(
+        *PrimaryFilePath, InputBufferName, Args, std::move(vfsOptions),
+        CancellationToken,
+        [Rec](const RequestResult<SemanticTokensResult> &Result) {
+          reportSemanticTokens(Result, Rec);
+        });
+    return;
+  });
+}
+
 /// Expand macros in the specified source file syntactically.
 ///
 /// Request would look like:
@@ -2100,6 +2130,7 @@ void handleRequestImpl(sourcekitd_object_t ReqObj,
   HANDLE_REQUEST(RequestRelatedIdents, handleRequestRelatedIdents)
   HANDLE_REQUEST(RequestActiveRegions, handleRequestActiveRegions)
   HANDLE_REQUEST(RequestDiagnostics, handleRequestDiagnostics)
+  HANDLE_REQUEST(RequestSemanticTokens, handleRequestSemanticTokens)
   HANDLE_REQUEST(RequestSyntacticMacroExpansion,
                  handleRequestSyntacticMacroExpansion)
 
@@ -2724,6 +2755,32 @@ static void reportDiagnostics(const RequestResult<DiagnosticsResult> &Result,
   Rec(RespBuilder.createResponse());
 }
 
+//===----------------------------------------------------------------------===//
+// ReportSemanticTokens
+//===----------------------------------------------------------------------===//
+
+static void
+reportSemanticTokens(const RequestResult<SemanticTokensResult> &Result,
+                     ResponseReceiver Rec) {
+  if (Result.isCancelled())
+    return Rec(createErrorRequestCancelled());
+  if (Result.isError())
+    return Rec(createErrorRequestFailed(Result.getError()));
+
+  ResponseBuilder RespBuilder;
+  auto Dict = RespBuilder.getDictionary();
+  TokenAnnotationsArrayBuilder SemanticAnnotations;
+  for (auto SemaTok : Result.value()) {
+    UIdent Kind = SemaTok.getUIdentForKind();
+    if (Kind.isValid()) {
+      SemanticAnnotations.add(Kind, SemaTok.ByteOffset, SemaTok.Length,
+                              SemaTok.getIsSystem());
+    }
+  }
+  Dict.setCustomBuffer(KeySemanticTokens, SemanticAnnotations.createBuffer());
+  Rec(RespBuilder.createResponse());
+}
+
 //===----------------------------------------------------------------------===//
 // ReportRangeInfo
 //===----------------------------------------------------------------------===//

utils/gyb_sourcekit_support/UIDs.py

@@ -69,6 +69,7 @@ UID_KEYS = [
     KEY('NumBytesToErase', 'key.num_bytes_to_erase'),
     KEY('NotRecommended', 'key.not_recommended'),
     KEY('Annotations', 'key.annotations'),
+    KEY('SemanticTokens', 'key.semantic_tokens'),
     KEY('DiagnosticStage', 'key.diagnostic_stage'),
     KEY('SyntaxMap', 'key.syntaxmap'),
     KEY('IsSystem', 'key.is_system'),
@@ -276,6 +277,7 @@ UID_REQUESTS = [
     REQUEST('GlobalConfiguration', 'source.request.configuration.global'),
     REQUEST('DependencyUpdated', 'source.request.dependency_updated'),
     REQUEST('Diagnostics', 'source.request.diagnostics'),
+    REQUEST('SemanticTokens', 'source.request.semantic_tokens'),
     REQUEST('Compile', 'source.request.compile'),
     REQUEST('CompileClose', 'source.request.compile.close'),
     REQUEST('EnableRequestBarriers', 'source.request.enable_request_barriers'),