[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.

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);
+}