Skip to content

Commit 8965020

Browse files
fwcdfwcd
authored
Merge pull request #278 from fwcd/semantic-tokens
Add support for semantic tokens/highlighting
2 parents f820c26 + 38ffd87 commit 8965020

File tree

8 files changed

+302
-5
lines changed

8 files changed

+302
-5
lines changed

‎server/src/main/kotlin/org/javacs/kt/KotlinLanguageServer.kt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import org.javacs.kt.util.TemporaryDirectory
1414
import org.javacs.kt.util.parseURI
1515
import org.javacs.kt.progress.Progress
1616
import org.javacs.kt.progress.LanguageClientProgress
17+
import org.javacs.kt.semantictokens.semanticTokensLegend
1718
import java.net.URI
1819
import java.io.Closeable
1920
import java.nio.file.Paths
@@ -81,6 +82,7 @@ class KotlinLanguageServer : LanguageServer, LanguageClientAware, Closeable {
8182
serverCapabilities.documentSymbolProvider = Either.forLeft(true)
8283
serverCapabilities.workspaceSymbolProvider = Either.forLeft(true)
8384
serverCapabilities.referencesProvider = Either.forLeft(true)
85+
serverCapabilities.semanticTokensProvider = SemanticTokensWithRegistrationOptions(semanticTokensLegend, true, true)
8486
serverCapabilities.codeActionProvider = Either.forLeft(true)
8587
serverCapabilities.documentFormattingProvider = Either.forLeft(true)
8688
serverCapabilities.documentRangeFormattingProvider = Either.forLeft(true)

‎server/src/main/kotlin/org/javacs/kt/KotlinTextDocumentService.kt

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import org.javacs.kt.position.offset
1414
import org.javacs.kt.position.extractRange
1515
import org.javacs.kt.position.position
1616
import org.javacs.kt.references.findReferences
17+
import org.javacs.kt.semantictokens.encodedSemanticTokens
1718
import org.javacs.kt.signaturehelp.fetchSignatureHelpAt
1819
import org.javacs.kt.symbols.documentSymbols
1920
import org.javacs.kt.util.noResult
@@ -223,7 +224,35 @@ class KotlinTextDocumentService(
223224
val offset = offset(content, position.position.line, position.position.character)
224225
findReferences(file, offset, sp)
225226
}
227+
}
228+
229+
override fun semanticTokensFull(params: SemanticTokensParams) = async.compute {
230+
LOG.info("Full semantic tokens in {}", describeURI(params.textDocument.uri))
231+
232+
reportTime {
233+
val uri = parseURI(params.textDocument.uri)
234+
val file = sp.currentVersion(uri)
235+
236+
val tokens = encodedSemanticTokens(file)
237+
LOG.info("Found {} tokens", tokens.size)
238+
239+
SemanticTokens(tokens)
240+
}
241+
}
242+
243+
override fun semanticTokensRange(params: SemanticTokensRangeParams) = async.compute {
244+
LOG.info("Ranged semantic tokens in {}", describeURI(params.textDocument.uri))
245+
246+
reportTime {
247+
val uri = parseURI(params.textDocument.uri)
248+
val file = sp.currentVersion(uri)
249+
250+
val tokens = encodedSemanticTokens(file, params.range)
251+
LOG.info("Found {} tokens", tokens.size)
252+
253+
SemanticTokens(tokens)
226254
}
255+
}
227256

228257
override fun resolveCodeLens(unresolved: CodeLens): CompletableFuture<CodeLens> {
229258
TODO("not implemented")

‎server/src/main/kotlin/org/javacs/kt/semantictokens/SemanticTokens.kt

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
package org.javacs.kt.semantictokens
2+
3+
import org.eclipse.lsp4j.SemanticTokenTypes
4+
import org.eclipse.lsp4j.SemanticTokenModifiers
5+
import org.eclipse.lsp4j.SemanticTokensLegend
6+
import org.eclipse.lsp4j.Range
7+
import org.javacs.kt.CompiledFile
8+
import org.javacs.kt.position.range
9+
import org.javacs.kt.position.offset
10+
import org.javacs.kt.util.preOrderTraversal
11+
import org.jetbrains.kotlin.descriptors.ClassDescriptor
12+
import org.jetbrains.kotlin.descriptors.ClassKind
13+
import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
14+
import org.jetbrains.kotlin.descriptors.FunctionDescriptor
15+
import org.jetbrains.kotlin.descriptors.PropertyDescriptor
16+
import org.jetbrains.kotlin.descriptors.VariableDescriptor
17+
import org.jetbrains.kotlin.lexer.KtTokens
18+
import org.jetbrains.kotlin.psi.KtClassOrObject
19+
import org.jetbrains.kotlin.psi.KtFunction
20+
import org.jetbrains.kotlin.psi.KtModifierListOwner
21+
import org.jetbrains.kotlin.psi.KtNameReferenceExpression
22+
import org.jetbrains.kotlin.psi.KtVariableDeclaration
23+
import org.jetbrains.kotlin.psi.KtNamedDeclaration
24+
import org.jetbrains.kotlin.psi.KtProperty
25+
import org.jetbrains.kotlin.psi.KtParameter
26+
import org.jetbrains.kotlin.psi.KtStringTemplateEntry
27+
import org.jetbrains.kotlin.psi.KtStringTemplateExpression
28+
import org.jetbrains.kotlin.psi.KtSimpleNameStringTemplateEntry
29+
import org.jetbrains.kotlin.psi.KtBlockStringTemplateEntry
30+
import org.jetbrains.kotlin.psi.KtEscapeStringTemplateEntry
31+
import org.jetbrains.kotlin.resolve.BindingContext
32+
import com.intellij.psi.PsiElement
33+
import com.intellij.psi.PsiNameIdentifierOwner
34+
import com.intellij.psi.PsiLiteralExpression
35+
import com.intellij.psi.PsiType
36+
import com.intellij.openapi.util.TextRange
37+
38+
enum class SemanticTokenType(val typeName: String) {
39+
KEYWORD(SemanticTokenTypes.Keyword),
40+
VARIABLE(SemanticTokenTypes.Variable),
41+
FUNCTION(SemanticTokenTypes.Function),
42+
PROPERTY(SemanticTokenTypes.Property),
43+
PARAMETER(SemanticTokenTypes.Parameter),
44+
ENUM_MEMBER(SemanticTokenTypes.EnumMember),
45+
CLASS(SemanticTokenTypes.Class),
46+
INTERFACE(SemanticTokenTypes.Interface),
47+
ENUM(SemanticTokenTypes.Enum),
48+
TYPE(SemanticTokenTypes.Type),
49+
STRING(SemanticTokenTypes.String),
50+
NUMBER(SemanticTokenTypes.Number),
51+
// Since LSP does not provide a token type for string interpolation
52+
// entries, we use Variable as a fallback here for now
53+
INTERPOLATION_ENTRY(SemanticTokenTypes.Variable)
54+
}
55+
56+
enum class SemanticTokenModifier(val modifierName: String) {
57+
DECLARATION(SemanticTokenModifiers.Declaration),
58+
DEFINITION(SemanticTokenModifiers.Definition),
59+
ABSTRACT(SemanticTokenModifiers.Abstract),
60+
READONLY(SemanticTokenModifiers.Readonly)
61+
}
62+
63+
val semanticTokensLegend = SemanticTokensLegend(
64+
SemanticTokenType.values().map { it.typeName },
65+
SemanticTokenModifier.values().map { it.modifierName }
66+
)
67+
68+
data class SemanticToken(val range: Range, val type: SemanticTokenType, val modifiers: Set<SemanticTokenModifier> = setOf())
69+
70+
/**
71+
* Computes LSP-encoded semantic tokens for the given range in the
72+
* document. No range means the entire document.
73+
*/
74+
fun encodedSemanticTokens(file: CompiledFile, range: Range? = null): List<Int> =
75+
encodeTokens(semanticTokens(file, range))
76+
77+
/**
78+
* Computes semantic tokens for the given range in the document.
79+
* No range means the entire document.
80+
*/
81+
fun semanticTokens(file: CompiledFile, range: Range? = null): Sequence<SemanticToken> =
82+
elementTokens(file.parse, file.compile, range)
83+
84+
fun encodeTokens(tokens: Sequence<SemanticToken>): List<Int> {
85+
val encoded = mutableListOf<Int>()
86+
var last: SemanticToken? = null
87+
88+
for (token in tokens) {
89+
// Tokens must be on a single line
90+
if (token.range.start.line == token.range.end.line) {
91+
val length = token.range.end.character - token.range.start.character
92+
val deltaLine = token.range.start.line - (last?.range?.start?.line ?: 0)
93+
val deltaStart = token.range.start.character - (last?.takeIf { deltaLine == 0 }?.range?.start?.character ?: 0)
94+
95+
encoded.add(deltaLine)
96+
encoded.add(deltaStart)
97+
encoded.add(length)
98+
encoded.add(encodeType(token.type))
99+
encoded.add(encodeModifiers(token.modifiers))
100+
101+
last = token
102+
}
103+
}
104+
105+
return encoded
106+
}
107+
108+
private fun encodeType(type: SemanticTokenType): Int = type.ordinal
109+
110+
private fun encodeModifiers(modifiers: Set<SemanticTokenModifier>): Int = modifiers
111+
.map { 1 shl it.ordinal }
112+
.fold(0, Int::or)
113+
114+
private fun elementTokens(element: PsiElement, bindingContext: BindingContext, range: Range? = null): Sequence<SemanticToken> {
115+
val file = element.containingFile
116+
val textRange = range?.let { TextRange(offset(file.text, it.start), offset(file.text, it.end)) }
117+
return element
118+
// TODO: Ideally we would like to cut-off subtrees outside our range, but this doesn't quite seem to work
119+
// .preOrderTraversal { elem -> textRange?.let { it.contains(elem.textRange) } ?: true }
120+
.preOrderTraversal()
121+
.filter { elem -> textRange?.let { it.contains(elem.textRange) } ?: true }
122+
.mapNotNull { elementToken(it, bindingContext) }
123+
}
124+
125+
private fun elementToken(element: PsiElement, bindingContext: BindingContext): SemanticToken? {
126+
val file = element.containingFile
127+
val elementRange = range(file.text, element.textRange)
128+
129+
return when (element) {
130+
// References (variables, types, functions, ...)
131+
132+
is KtNameReferenceExpression -> {
133+
val target = bindingContext[BindingContext.REFERENCE_TARGET, element]
134+
val tokenType = when (target) {
135+
is PropertyDescriptor -> SemanticTokenType.PROPERTY
136+
is VariableDescriptor -> SemanticTokenType.VARIABLE
137+
is ConstructorDescriptor -> when (target.constructedClass.kind) {
138+
ClassKind.ANNOTATION_CLASS -> SemanticTokenType.TYPE // annotations look nicer this way
139+
else -> SemanticTokenType.FUNCTION
140+
}
141+
is FunctionDescriptor -> SemanticTokenType.FUNCTION
142+
is ClassDescriptor -> when (target.kind) {
143+
ClassKind.CLASS -> SemanticTokenType.CLASS
144+
ClassKind.OBJECT -> SemanticTokenType.CLASS
145+
ClassKind.INTERFACE -> SemanticTokenType.INTERFACE
146+
ClassKind.ENUM_CLASS -> SemanticTokenType.ENUM
147+
else -> SemanticTokenType.TYPE
148+
}
149+
else -> return null
150+
}
151+
val isConstant = (target as? VariableDescriptor)?.let { !it.isVar() || it.isConst() } ?: false
152+
val modifiers = if (isConstant) setOf(SemanticTokenModifier.READONLY) else setOf()
153+
154+
SemanticToken(elementRange, tokenType, modifiers)
155+
}
156+
157+
// Declarations (variables, types, functions, ...)
158+
159+
is PsiNameIdentifierOwner -> {
160+
val tokenType = when (element) {
161+
is KtParameter -> SemanticTokenType.PARAMETER
162+
is KtProperty -> SemanticTokenType.PROPERTY
163+
is KtVariableDeclaration -> SemanticTokenType.VARIABLE
164+
is KtClassOrObject -> SemanticTokenType.CLASS
165+
is KtFunction -> SemanticTokenType.FUNCTION
166+
else -> return null
167+
}
168+
val identifierRange = element.nameIdentifier?.let { range(file.text, it.textRange) } ?: return null
169+
val modifiers = mutableSetOf(SemanticTokenModifier.DECLARATION)
170+
171+
if (element is KtVariableDeclaration && (!element.isVar() || element.hasModifier(KtTokens.CONST_KEYWORD)) || element is KtParameter) {
172+
modifiers.add(SemanticTokenModifier.READONLY)
173+
}
174+
175+
if (element is KtModifierListOwner) {
176+
if (element.hasModifier(KtTokens.ABSTRACT_KEYWORD)) {
177+
modifiers.add(SemanticTokenModifier.ABSTRACT)
178+
}
179+
}
180+
181+
SemanticToken(identifierRange, tokenType, modifiers)
182+
}
183+
184+
// Literals and string interpolations
185+
186+
is KtSimpleNameStringTemplateEntry, is KtBlockStringTemplateEntry ->
187+
SemanticToken(elementRange, SemanticTokenType.INTERPOLATION_ENTRY)
188+
is KtStringTemplateExpression -> SemanticToken(elementRange, SemanticTokenType.STRING)
189+
is PsiLiteralExpression -> {
190+
val tokenType = when (element.type) {
191+
PsiType.INT, PsiType.LONG, PsiType.DOUBLE -> SemanticTokenType.NUMBER
192+
PsiType.CHAR -> SemanticTokenType.STRING
193+
PsiType.BOOLEAN, PsiType.NULL -> SemanticTokenType.KEYWORD
194+
else -> return null
195+
}
196+
SemanticToken(elementRange, tokenType)
197+
}
198+
else -> null
199+
}
200+
}

‎server/src/main/kotlin/org/javacs/kt/util/PsiUtils.kt

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,18 @@ import java.nio.file.Path
88
inline fun<reified Find> PsiElement.findParent() =
99
this.parentsWithSelf.filterIsInstance<Find>().firstOrNull()
1010

11-
fun PsiElement.preOrderTraversal(): Sequence<PsiElement> {
11+
fun PsiElement.preOrderTraversal(shouldTraverse: (PsiElement) -> Boolean = { true }): Sequence<PsiElement> {
1212
val root = this
1313

1414
return sequence {
15-
yield(root)
15+
if (shouldTraverse(root)) {
16+
yield(root)
1617

17-
for (child in root.children) {
18-
yieldAll(child.preOrderTraversal())
18+
for (child in root.children) {
19+
if (shouldTraverse(child)) {
20+
yieldAll(child.preOrderTraversal(shouldTraverse))
21+
}
22+
}
1923
}
2024
}
2125
}

‎server/src/test/kotlin/org/javacs/kt/LanguageServerTestFixture.kt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,12 @@ abstract class LanguageServerTestFixture(relativeWorkspaceRoot: String) : Langua
6868
fun hoverParams(relativePath: String, line: Int, column: Int): HoverParams =
6969
textDocumentPosition(relativePath, line, column).run { HoverParams(textDocument, position) }
7070

71+
fun semanticTokensParams(relativePath: String): SemanticTokensParams =
72+
textDocumentPosition(relativePath, 0, 0).run { SemanticTokensParams(textDocument) }
73+
74+
fun semanticTokensRangeParams(relativePath: String, range: Range): SemanticTokensRangeParams =
75+
textDocumentPosition(relativePath, 0, 0).run { SemanticTokensRangeParams(textDocument, range) }
76+
7177
fun signatureHelpParams(relativePath: String, line: Int, column: Int): SignatureHelpParams =
7278
textDocumentPosition(relativePath, line, column).run { SignatureHelpParams(textDocument, position) }
7379

‎server/src/test/kotlin/org/javacs/kt/SemanticTokensTest.kt

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package org.javacs.kt
2+
3+
import org.hamcrest.Matchers.*
4+
import org.junit.Assert.assertThat
5+
import org.junit.Test
6+
import org.javacs.kt.semantictokens.encodeTokens
7+
import org.javacs.kt.semantictokens.SemanticToken
8+
import org.javacs.kt.semantictokens.SemanticTokenType
9+
import org.javacs.kt.semantictokens.SemanticTokenModifier
10+
11+
class SemanticTokensTest : SingleFileTestFixture("semantictokens", "SemanticTokens.kt") {
12+
@Test fun `tokenize file`() {
13+
val varLine = 1
14+
val constLine = 2
15+
val classLine = 4
16+
val funLine = 6
17+
18+
val expectedVar = sequenceOf(
19+
SemanticToken(range(varLine, 5, varLine, 13), SemanticTokenType.PROPERTY, setOf(SemanticTokenModifier.DECLARATION)), // variable
20+
)
21+
val expectedConst = sequenceOf(
22+
SemanticToken(range(constLine, 5, constLine, 13), SemanticTokenType.PROPERTY, setOf(SemanticTokenModifier.DECLARATION, SemanticTokenModifier.READONLY)), // constant
23+
SemanticToken(range(constLine, 15, constLine, 21), SemanticTokenType.CLASS), // String
24+
SemanticToken(range(constLine, 24, constLine, 40), SemanticTokenType.STRING), // "test $variable"
25+
SemanticToken(range(constLine, 30, constLine, 39), SemanticTokenType.INTERPOLATION_ENTRY), // $variable
26+
SemanticToken(range(constLine, 31, constLine, 39), SemanticTokenType.PROPERTY), // variable
27+
)
28+
val expectedClass = sequenceOf(
29+
SemanticToken(range(classLine, 12, classLine, 16), SemanticTokenType.CLASS, setOf(SemanticTokenModifier.DECLARATION)), // Type
30+
SemanticToken(range(classLine, 21, classLine, 29), SemanticTokenType.PARAMETER, setOf(SemanticTokenModifier.DECLARATION, SemanticTokenModifier.READONLY)), // property
31+
SemanticToken(range(classLine, 31, classLine, 34), SemanticTokenType.CLASS), // Int
32+
)
33+
val expectedFun = sequenceOf(
34+
SemanticToken(range(funLine, 5, funLine, 6), SemanticTokenType.FUNCTION, setOf(SemanticTokenModifier.DECLARATION)), // f
35+
SemanticToken(range(funLine, 7, funLine, 8), SemanticTokenType.PARAMETER, setOf(SemanticTokenModifier.DECLARATION, SemanticTokenModifier.READONLY)), // x
36+
SemanticToken(range(funLine, 10, funLine, 13), SemanticTokenType.CLASS), // Int?
37+
SemanticToken(range(funLine, 24, funLine, 27), SemanticTokenType.CLASS), // Int
38+
SemanticToken(range(funLine, 30, funLine, 31), SemanticTokenType.FUNCTION), // f
39+
SemanticToken(range(funLine, 32, funLine, 33), SemanticTokenType.VARIABLE, setOf(SemanticTokenModifier.READONLY)), // x
40+
)
41+
42+
val partialExpected = encodeTokens(expectedConst + expectedClass)
43+
val partialResponse = languageServer.textDocumentService.semanticTokensRange(semanticTokensRangeParams(file, range(constLine, 0, classLine + 1, 0))).get()!!
44+
assertThat(partialResponse.data, contains(*partialExpected.toTypedArray()))
45+
46+
val fullExpected = encodeTokens(expectedVar + expectedConst + expectedClass + expectedFun)
47+
val fullResponse = languageServer.textDocumentService.semanticTokensFull(semanticTokensParams(file)).get()!!
48+
assertThat(fullResponse.data, contains(*fullExpected.toTypedArray()))
49+
}
50+
}

‎server/src/test/resources/semantictokens/SemanticTokens.kt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
var variable = 3
2+
val constant: String = "test $variable"
3+
4+
data class Type(val property: Int)
5+
6+
fun f(x: Int? = null): Int = f(x)

‎shared/src/main/kotlin/org/javacs/kt/classpath/GradleClassPathResolver.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ private fun readDependenciesViaGradleCLI(projectDirectory: Path, gradleScripts:
8080
val dependencies = findGradleCLIDependencies(command, projectDirectory)
8181
?.also { LOG.debug("Classpath for task {}", it) }
8282
.orEmpty()
83-
.filter { it.toString().toLowerCase().endsWith(".jar") || Files.isDirectory(it) } // Some Gradle plugins seem to cause this to output POMs, therefore filter JARs
83+
.filter { it.toString().lowercase().endsWith(".jar") || Files.isDirectory(it) } // Some Gradle plugins seem to cause this to output POMs, therefore filter JARs
8484
.toSet()
8585

8686
tmpScripts.forEach(Files::delete)

0 commit comments

Comments
 (0)