Skip to content

Commit 4bc03f8

Browse files
ahoppenahoppen
committed
[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.
1 parent 1560e08 commit 4bc03f8

File tree

11 files changed

+206
-33
lines changed

11 files changed

+206
-33
lines changed

‎test/SourceKit/SemanticTokens/semantic_tokens.swift

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// RUN: %sourcekitd-test -req=semantic-tokens %s -- %s | %FileCheck %s
2+
3+
class MyClass {}
4+
5+
let x: String = "test"
6+
var y = MyClass()
7+
8+
9+
// String in line 5
10+
// CHECK: key.kind: source.lang.swift.ref.struct
11+
// CHECK: key.length: 6,
12+
// CHECK key.is_system: 1
13+
// CHECK-LABEL: },
14+
15+
// MyClass in line 6
16+
// CHECK: key.kind: source.lang.swift.ref.class,
17+
// CHECK: key.length: 7

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include "SourceKit/Support/CancellationToken.h"
1818
#include "SourceKit/Support/UIdent.h"
1919
#include "swift/AST/Type.h"
20+
#include "swift/IDE/CodeCompletionResult.h"
2021
#include "llvm/ADT/ArrayRef.h"
2122
#include "llvm/ADT/IntrusiveRefCntPtr.h"
2223
#include "llvm/ADT/Optional.h"
@@ -330,6 +331,35 @@ struct DiagnosticEntryInfo : DiagnosticEntryInfoBase {
330331
SmallVector<DiagnosticEntryInfoBase, 1> Notes;
331332
};
332333

334+
struct SwiftSemanticToken {
335+
unsigned ByteOffset;
336+
unsigned Length : 24;
337+
// The code-completion kinds are a good match for the semantic kinds we want.
338+
// FIXME: Maybe rename CodeCompletionDeclKind to a more general concept ?
339+
swift::ide::CodeCompletionDeclKind Kind : 6;
340+
unsigned IsRef : 1;
341+
unsigned IsSystem : 1;
342+
343+
SwiftSemanticToken(swift::ide::CodeCompletionDeclKind Kind,
344+
unsigned ByteOffset, unsigned Length, bool IsRef,
345+
bool IsSystem)
346+
: ByteOffset(ByteOffset), Length(Length), Kind(Kind), IsRef(IsRef),
347+
IsSystem(IsSystem) {}
348+
349+
bool getIsRef() const { return static_cast<bool>(IsRef); }
350+
351+
bool getIsSystem() const { return static_cast<bool>(IsSystem); }
352+
353+
UIdent getUIdentForKind();
354+
};
355+
356+
#if !defined(_MSC_VER)
357+
static_assert(sizeof(SwiftSemanticToken) == 8, "Too big");
358+
// FIXME: MSVC doesn't pack bitfields with different underlying types.
359+
// Giving up to check this in MSVC for now, because static_assert is only for
360+
// keeping low memory usage.
361+
#endif
362+
333363
struct SourceFileRange {
334364
/// The byte offset at which the range begins
335365
uintptr_t Start;
@@ -709,6 +739,9 @@ struct CursorInfoData {
709739
/// The result type of `LangSupport::getDiagnostics`
710740
typedef ArrayRef<DiagnosticEntryInfo> DiagnosticsResult;
711741

742+
/// The result of `LangSupport::getSemanticTokens`.
743+
typedef std::vector<SwiftSemanticToken> SemanticTokensResult;
744+
712745
struct RangeInfo {
713746
UIdent RangeKind;
714747
StringRef ExprType;
@@ -1103,6 +1136,13 @@ class LangSupport {
11031136
std::function<void(const RequestResult<DiagnosticsResult> &)>
11041137
Receiver) = 0;
11051138

1139+
virtual void getSemanticTokens(
1140+
StringRef PrimaryFilePath, StringRef InputBufferName,
1141+
ArrayRef<const char *> Args, llvm::Optional<VFSOptions> VfsOptions,
1142+
SourceKitCancellationToken CancellationToken,
1143+
std::function<void(const RequestResult<SemanticTokensResult> &)>
1144+
Receiver) = 0;
1145+
11061146
virtual void
11071147
getNameInfo(StringRef PrimaryFilePath, StringRef InputBufferName,
11081148
unsigned Offset, NameTranslatingInfo &Input,

‎tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

Lines changed: 65 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -643,36 +643,6 @@ struct EditorConsumerSyntaxMapEntry {
643643
:Offset(Offset), Length(Length), Kind(Kind) { }
644644
};
645645

646-
struct SwiftSemanticToken {
647-
unsigned ByteOffset;
648-
unsigned Length : 24;
649-
// The code-completion kinds are a good match for the semantic kinds we want.
650-
// FIXME: Maybe rename CodeCompletionDeclKind to a more general concept ?
651-
CodeCompletionDeclKind Kind : 6;
652-
unsigned IsRef : 1;
653-
unsigned IsSystem : 1;
654-
655-
SwiftSemanticToken(CodeCompletionDeclKind Kind,
656-
unsigned ByteOffset, unsigned Length,
657-
bool IsRef, bool IsSystem)
658-
: ByteOffset(ByteOffset), Length(Length), Kind(Kind),
659-
IsRef(IsRef), IsSystem(IsSystem) { }
660-
661-
bool getIsRef() const { return static_cast<bool>(IsRef); }
662-
663-
bool getIsSystem() const { return static_cast<bool>(IsSystem); }
664-
665-
UIdent getUIdentForKind() const {
666-
return SwiftLangSupport::getUIDForCodeCompletionDeclKind(Kind, getIsRef());
667-
}
668-
};
669-
#if !defined(_MSC_VER)
670-
static_assert(sizeof(SwiftSemanticToken) == 8, "Too big");
671-
// FIXME: MSVC doesn't pack bitfields with different underlying types.
672-
// Giving up to check this in MSVC for now, because static_assert is only for
673-
// keeping low memory usage.
674-
#endif
675-
676646
class SwiftDocumentSemanticInfo :
677647
public ThreadSafeRefCountedBase<SwiftDocumentSemanticInfo> {
678648

@@ -2535,3 +2505,68 @@ void SwiftLangSupport::editorExpandPlaceholder(StringRef Name, unsigned Offset,
25352505
}
25362506
EditorDoc->expandPlaceholder(Offset, Length, Consumer);
25372507
}
2508+
2509+
//===----------------------------------------------------------------------===//
2510+
// Semantic Tokens
2511+
//===----------------------------------------------------------------------===//
2512+
2513+
void SwiftLangSupport::getSemanticTokens(
2514+
StringRef PrimaryFilePath, StringRef InputBufferName,
2515+
ArrayRef<const char *> Args, llvm::Optional<VFSOptions> VfsOptions,
2516+
SourceKitCancellationToken CancellationToken,
2517+
std::function<void(const RequestResult<SemanticTokensResult> &)> Receiver) {
2518+
std::string FileSystemError;
2519+
auto FileSystem = getFileSystem(VfsOptions, PrimaryFilePath, FileSystemError);
2520+
if (!FileSystem) {
2521+
Receiver(RequestResult<SemanticTokensResult>::fromError(FileSystemError));
2522+
return;
2523+
}
2524+
2525+
std::string InvocationError;
2526+
SwiftInvocationRef Invok = ASTMgr->getTypecheckInvocation(
2527+
Args, PrimaryFilePath, FileSystem, InvocationError);
2528+
if (!InvocationError.empty()) {
2529+
LOG_WARN_FUNC("error creating ASTInvocation: " << InvocationError);
2530+
}
2531+
if (!Invok) {
2532+
Receiver(RequestResult<SemanticTokensResult>::fromError(InvocationError));
2533+
return;
2534+
}
2535+
2536+
class SemanticTokensConsumer : public SwiftASTConsumer {
2537+
StringRef InputBufferName;
2538+
std::function<void(const RequestResult<SemanticTokensResult> &)> Receiver;
2539+
2540+
public:
2541+
SemanticTokensConsumer(
2542+
StringRef InputBufferName,
2543+
std::function<void(const RequestResult<SemanticTokensResult> &)>
2544+
Receiver)
2545+
: InputBufferName(InputBufferName), Receiver(Receiver) {}
2546+
2547+
void handlePrimaryAST(ASTUnitRef AstUnit) override {
2548+
auto &CompIns = AstUnit->getCompilerInstance();
2549+
SourceFile *SF = retrieveInputFile(InputBufferName, CompIns);
2550+
if (!SF) {
2551+
Receiver(RequestResult<SemanticTokensResult>::fromError(
2552+
"Unable to find input file"));
2553+
return;
2554+
}
2555+
SemanticAnnotator Annotator(CompIns.getSourceMgr(), *SF->getBufferID());
2556+
Annotator.walk(SF);
2557+
Receiver(
2558+
RequestResult<SemanticTokensResult>::fromResult(Annotator.SemaToks));
2559+
}
2560+
2561+
void cancelled() override {
2562+
Receiver(RequestResult<SemanticTokensResult>::cancelled());
2563+
}
2564+
};
2565+
2566+
auto Consumer = std::make_shared<SemanticTokensConsumer>(InputBufferName,
2567+
std::move(Receiver));
2568+
2569+
getASTManager()->processASTAsync(Invok, std::move(Consumer),
2570+
/*OncePerASTToken=*/nullptr,
2571+
CancellationToken, FileSystem);
2572+
}

‎tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -397,6 +397,10 @@ UIdent SwiftLangSupport::getUIDForRefactoringKind(ide::RefactoringKind Kind){
397397
}
398398
}
399399

400+
UIdent SwiftSemanticToken::getUIdentForKind() {
401+
return SwiftLangSupport::getUIDForCodeCompletionDeclKind(Kind, IsRef);
402+
}
403+
400404
UIdent SwiftLangSupport::getUIDForCodeCompletionDeclKind(
401405
ide::CodeCompletionDeclKind Kind, bool IsRef) {
402406
if (IsRef) {

‎tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -665,6 +665,13 @@ class SwiftLangSupport : public LangSupport {
665665
std::function<void(const RequestResult<DiagnosticsResult> &)>
666666
Receiver) override;
667667

668+
void getSemanticTokens(
669+
StringRef PrimaryFilePath, StringRef InputBufferName,
670+
ArrayRef<const char *> Args, llvm::Optional<VFSOptions> VfsOptions,
671+
SourceKitCancellationToken CancellationToken,
672+
std::function<void(const RequestResult<SemanticTokensResult> &)> Receiver)
673+
override;
674+
668675
void getNameInfo(
669676
StringRef PrimaryFilePath, StringRef InputBufferName, unsigned Offset,
670677
NameTranslatingInfo &Input, ArrayRef<const char *> Args,
@@ -812,6 +819,10 @@ class CloseClangModuleFiles {
812819
/// Disable expensive SIL options which do not affect indexing or diagnostics.
813820
void disableExpensiveSILOptions(swift::SILOptions &Opts);
814821

822+
swift::SourceFile *retrieveInputFile(StringRef inputBufferName,
823+
const swift::CompilerInstance &CI,
824+
bool haveRealPath = false);
825+
815826
} // namespace SourceKit
816827

817828
#endif

‎tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1537,9 +1537,9 @@ class CursorRangeInfoConsumer : public SwiftASTConsumer {
15371537
}
15381538
};
15391539

1540-
static SourceFile *retrieveInputFile(StringRef inputBufferName,
1541-
const CompilerInstance &CI,
1542-
bool haveRealPath = false) {
1540+
SourceFile *SourceKit::retrieveInputFile(StringRef inputBufferName,
1541+
const CompilerInstance &CI,
1542+
bool haveRealPath) {
15431543
// Don't bother looking up if we have the same file as the primary file or
15441544
// we weren't given a separate input file
15451545
if (inputBufferName.empty() ||

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
151151
.Case("global-config", SourceKitRequest::GlobalConfiguration)
152152
.Case("dependency-updated", SourceKitRequest::DependencyUpdated)
153153
.Case("diags", SourceKitRequest::Diagnostics)
154+
.Case("semantic-tokens", SourceKitRequest::SemanticTokens)
154155
.Case("compile", SourceKitRequest::Compile)
155156
.Case("compile.close", SourceKitRequest::CompileClose)
156157
.Case("syntactic-expandmacro", SourceKitRequest::SyntacticMacroExpansion)

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ enum class SourceKitRequest {
6868
GlobalConfiguration,
6969
DependencyUpdated,
7070
Diagnostics,
71+
SemanticTokens,
7172
Compile,
7273
CompileClose,
7374
SyntacticMacroExpansion,

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1137,6 +1137,10 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
11371137
case SourceKitRequest::Diagnostics:
11381138
sourcekitd_request_dictionary_set_uid(Req, KeyRequest, RequestDiagnostics);
11391139
break;
1140+
case SourceKitRequest::SemanticTokens:
1141+
sourcekitd_request_dictionary_set_uid(Req, KeyRequest,
1142+
RequestSemanticTokens);
1143+
break;
11401144

11411145
case SourceKitRequest::Compile:
11421146
sourcekitd_request_dictionary_set_string(Req, KeyName, SemaName.c_str());
@@ -1427,6 +1431,7 @@ static bool handleResponse(sourcekitd_response_t Resp, const TestOptions &Opts,
14271431
case SourceKitRequest::ConformingMethodList:
14281432
case SourceKitRequest::DependencyUpdated:
14291433
case SourceKitRequest::Diagnostics:
1434+
case SourceKitRequest::SemanticTokens:
14301435
printRawResponse(Resp);
14311436
break;
14321437
case SourceKitRequest::Compile:

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

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,10 @@ static void reportCursorInfo(const RequestResult<CursorInfoData> &Result, Respon
194194
static void reportDiagnostics(const RequestResult<DiagnosticsResult> &Result,
195195
ResponseReceiver Rec);
196196
197+
static void
198+
reportSemanticTokens(const RequestResult<SemanticTokensResult> &Result,
199+
ResponseReceiver Rec);
200+
197201
static void reportExpressionTypeInfo(const RequestResult<ExpressionTypesInFile> &Result,
198202
ResponseReceiver Rec);
199203
@@ -1866,6 +1870,32 @@ handleRequestDiagnostics(const RequestDict &Req,
18661870
});
18671871
}
18681872
1873+
static void
1874+
handleRequestSemanticTokens(const RequestDict &Req,
1875+
SourceKitCancellationToken CancellationToken,
1876+
ResponseReceiver Rec) {
1877+
handleSemanticRequest(Req, Rec, [Req, CancellationToken, Rec]() {
1878+
Optional<VFSOptions> vfsOptions = getVFSOptions(Req);
1879+
auto PrimaryFilePath = getPrimaryFilePathForRequestOrEmitError(Req, Rec);
1880+
if (!PrimaryFilePath)
1881+
return;
1882+
StringRef InputBufferName = getInputBufferNameForRequest(Req, Rec);
1883+
1884+
SmallVector<const char *, 8> Args;
1885+
if (getCompilerArgumentsForRequestOrEmitError(Req, Args, Rec))
1886+
return;
1887+
1888+
LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
1889+
Lang.getSemanticTokens(
1890+
*PrimaryFilePath, InputBufferName, Args, std::move(vfsOptions),
1891+
CancellationToken,
1892+
[Rec](const RequestResult<SemanticTokensResult> &Result) {
1893+
reportSemanticTokens(Result, Rec);
1894+
});
1895+
return;
1896+
});
1897+
}
1898+
18691899
/// Expand macros in the specified source file syntactically.
18701900
///
18711901
/// Request would look like:
@@ -2100,6 +2130,7 @@ void handleRequestImpl(sourcekitd_object_t ReqObj,
21002130
HANDLE_REQUEST(RequestRelatedIdents, handleRequestRelatedIdents)
21012131
HANDLE_REQUEST(RequestActiveRegions, handleRequestActiveRegions)
21022132
HANDLE_REQUEST(RequestDiagnostics, handleRequestDiagnostics)
2133+
HANDLE_REQUEST(RequestSemanticTokens, handleRequestSemanticTokens)
21032134
HANDLE_REQUEST(RequestSyntacticMacroExpansion,
21042135
handleRequestSyntacticMacroExpansion)
21052136

@@ -2724,6 +2755,32 @@ static void reportDiagnostics(const RequestResult<DiagnosticsResult> &Result,
27242755
Rec(RespBuilder.createResponse());
27252756
}
27262757

2758+
//===----------------------------------------------------------------------===//
2759+
// ReportSemanticTokens
2760+
//===----------------------------------------------------------------------===//
2761+
2762+
static void
2763+
reportSemanticTokens(const RequestResult<SemanticTokensResult> &Result,
2764+
ResponseReceiver Rec) {
2765+
if (Result.isCancelled())
2766+
return Rec(createErrorRequestCancelled());
2767+
if (Result.isError())
2768+
return Rec(createErrorRequestFailed(Result.getError()));
2769+
2770+
ResponseBuilder RespBuilder;
2771+
auto Dict = RespBuilder.getDictionary();
2772+
TokenAnnotationsArrayBuilder SemanticAnnotations;
2773+
for (auto SemaTok : Result.value()) {
2774+
UIdent Kind = SemaTok.getUIdentForKind();
2775+
if (Kind.isValid()) {
2776+
SemanticAnnotations.add(Kind, SemaTok.ByteOffset, SemaTok.Length,
2777+
SemaTok.getIsSystem());
2778+
}
2779+
}
2780+
Dict.setCustomBuffer(KeySemanticTokens, SemanticAnnotations.createBuffer());
2781+
Rec(RespBuilder.createResponse());
2782+
}
2783+
27272784
//===----------------------------------------------------------------------===//
27282785
// ReportRangeInfo
27292786
//===----------------------------------------------------------------------===//

0 commit comments

Comments
 (0)