[lsp] implement textDocument/semanticTokens/range request

^LSP-89 fixed

GitOrigin-RevId: 9eb57a3200039e92f51fa1fec1cc08626d534896

Author: darthorimar

api.core/BUILD.bazel

@@ -39,6 +39,8 @@ jvm_library(
     "@community//fleet/lsp.protocol",
     "@lib//:junit5",
     "@lib//:junit5Jupiter",
+    "@community//platform/testFramework",
+    "@community//platform/testFramework:testFramework_test_lib",
   ]
 )
 

api.core/language-server.api.core.iml

@@ -32,5 +32,6 @@
     <orderEntry type="module" module-name="fleet.lsp.protocol" />
     <orderEntry type="library" scope="TEST" name="JUnit5" level="project" />
     <orderEntry type="library" scope="TEST" name="JUnit5Jupiter" level="project" />
+    <orderEntry type="module" module-name="intellij.platform.testFramework" scope="TEST" />
   </component>
 </module>

api.core/src/com/jetbrains/ls/api/core/util/LSUtil.kt

@@ -10,8 +10,29 @@ import com.jetbrains.lsp.protocol.*
 val VirtualFile.uri: URI
     get() = url.intellijUriToLspUri()
 
-fun Document.offsetByPosition(position: Position): Int = 
-    getLineStartOffset(position.line) + position.character
+/**
+ * Calculates the absolute offset in the document text based on the given position (line and character offset).
+ *
+ * @param position The position in the document, represented by a line number (zero-based) and
+ *                 character offset in that line (zero-based).
+ * @return The absolute offset in the document, which represents the character index corresponding
+ *         to the given position. If the position line is outside the document bounds, returns the document
+ *         end offset. If the character offset is outside the line bounds, returns the line end offset.
+ */
+fun Document.offsetByPosition(position: Position): Int {
+    val textLength = textLength
+    if (position.line >= lineCount) {
+        // lsp position may be outside the document, which means the end of the document
+        return textLength
+    }
+    val lineStart = getLineStartOffset(position.line)
+    val lineEnd = getLineEndOffset(position.line)
+    if (position.character > lineEnd - lineStart) {
+        // lsp position may be outside the line range, which means the end of the line
+        return lineEnd
+    }
+    return lineStart + position.character
+}
 
 fun Document.positionByOffset(offset: Int): Position {
     val line = getLineNumber(offset) 

api.core/test/com/jetbrains/ls/api/core/util/PositionsTest.kt

@@ -0,0 +1,90 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.ls.api.core.util
+
+import com.intellij.mock.MockMultiLineImmutableDocument
+import com.intellij.openapi.editor.Document
+import com.jetbrains.lsp.protocol.Position
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+
+class PositionsTest {
+    @Test
+    fun `should convert offsets on empty document`() {
+        val document = createDocument()
+
+        assertEquals(0, document.offsetByPosition(Position(0, 0)))
+        assertEquals(0, document.offsetByPosition(Position(10, 20)))
+    }
+
+    @Test
+    fun `should convert offsets outside the range of the document`() {
+        val document = createDocument(5, 6)
+
+        assertEquals(document.textLength, document.offsetByPosition(Position(1, 6)))
+        assertEquals(document.textLength, document.offsetByPosition(Position(1, 7)))
+        assertEquals(document.textLength, document.offsetByPosition(Position(2, 0)))
+        assertEquals(document.textLength, document.offsetByPosition(Position(2, 10)))
+    }
+
+
+    @Test
+    fun `should convert offsets  outside line ranges`() {
+        val document = createDocument(5, 6)
+
+        assertEquals(5, document.offsetByPosition(Position(0, 6)))
+        assertEquals(5, document.offsetByPosition(Position(0, 7)))
+        assertEquals(5, document.offsetByPosition(Position(0, 100)))
+        assertEquals(12, document.offsetByPosition(Position(1, 7)))
+        assertEquals(12, document.offsetByPosition(Position(1, 8)))
+        assertEquals(12, document.offsetByPosition(Position(1, 100)))
+    }
+
+
+    @Test
+    fun `should convert offsets inside the range of the document`() {
+        val document = createDocument(5, 6)
+
+        testInsideOffset(0, Position(0, 0), document)
+        testInsideOffset(1, Position(0, 1), document)
+        testInsideOffset(4, Position(0, 4), document)
+        testInsideOffset(6, Position(1, 0), document)
+        testInsideOffset(10, Position(1, 4), document)
+        testInsideOffset(11, Position(1, 5), document)
+    }
+
+    private fun testInsideOffset(
+        expectedOffset: Int,
+        position: Position,
+        document: Document,
+    ) {
+        val actualOffset = document.offsetByPosition(position)
+        assertEquals(expectedOffset, actualOffset)
+
+        val positionByOffset = document.positionByOffset(actualOffset)
+        assertEquals(position, positionByOffset)
+    }
+
+
+    private fun createDocument(
+        vararg lineLengths: Int
+    ): Document {
+        return createDocument(lineLengths.toList())
+    }
+
+    private fun createDocument(
+        lineLengths: List<Int>
+    ): Document {
+        val char = 'a'
+        val text = buildString {
+            for (length in lineLengths) {
+                repeat(length) {
+                    append(char)
+                }
+                if (lineLengths.last() != length) {
+                    append('\n')
+                }
+            }
+        }
+        return MockMultiLineImmutableDocument(text)
+    }
+}

api.features/src/com/jetbrains/ls/api/features/semanticTokens/LSSemanticTokens.kt

@@ -6,18 +6,28 @@ import com.jetbrains.ls.api.features.LSConfiguration
 import com.jetbrains.ls.api.features.semanticTokens.encoding.SemanticTokensEncoder
 import com.jetbrains.lsp.protocol.SemanticTokens
 import com.jetbrains.lsp.protocol.SemanticTokensParams
+import com.jetbrains.lsp.protocol.SemanticTokensRangeParams
 
+// todo send partial results here
 object LSSemanticTokens {
     context(LSServer, LSConfiguration)
     suspend fun semanticTokensFull(params: SemanticTokensParams): SemanticTokens {
-        // todo send partial results here
         val providers = entriesFor<LSSemanticTokensProvider>(params.textDocument)
         val result = providers.flatMap { it.full(params) }
         val registry = createRegistry()
         val encoded = SemanticTokensEncoder.encode(result, registry)
         return SemanticTokens(resultId = null, data = encoded)
     }
 
+    context(LSServer, LSConfiguration)
+    suspend fun semanticTokensRange(params: SemanticTokensRangeParams): SemanticTokens {
+        val providers = entriesFor<LSSemanticTokensProvider>(params.textDocument)
+        val result = providers.flatMap { it.range(params) }
+        val registry = createRegistry()
+        val encoded = SemanticTokensEncoder.encode(result, registry)
+        return SemanticTokens(resultId = null, data = encoded)
+    }
+
     context(LSConfiguration)
     fun createRegistry(): LSSemanticTokenRegistry {
         val registries = entries<LSSemanticTokensProvider>().map { it.createRegistry() }

api.features/src/com/jetbrains/ls/api/features/semanticTokens/LSSemanticTokensProvider.kt

@@ -2,9 +2,9 @@
 package com.jetbrains.ls.api.features.semanticTokens
 
 import com.jetbrains.ls.api.core.LSServer
-import com.jetbrains.ls.api.features.LSConfigurationEntry
 import com.jetbrains.ls.api.features.LSLanguageSpecificConfigurationEntry
 import com.jetbrains.lsp.protocol.SemanticTokensParams
+import com.jetbrains.lsp.protocol.SemanticTokensRangeParams
 
 interface LSSemanticTokensProvider : LSLanguageSpecificConfigurationEntry {
     fun createRegistry(): LSSemanticTokenRegistry
@@ -14,4 +14,10 @@ interface LSSemanticTokensProvider : LSLanguageSpecificConfigurationEntry {
      */
     context(LSServer)
     suspend fun full(params: SemanticTokensParams): List<LSSemanticTokenWithRange>
+
+    /**
+     * LSP method `textDocument/semanticTokens/range`
+     */
+    context(LSServer)
+    suspend fun range(params: SemanticTokensRangeParams): List<LSSemanticTokenWithRange>
 }

api.features/src/com/jetbrains/ls/api/features/utils/ranges.kt

@@ -0,0 +1,27 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.ls.api.features.utils
+
+import com.intellij.openapi.editor.Document
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiWhiteSpace
+import com.intellij.psi.util.descendantsOfType
+import com.jetbrains.ls.api.core.util.toTextRange
+import com.jetbrains.lsp.protocol.Range
+
+/**
+ * Retrieves all non-whitespace child elements from the given `PsiFile`.
+ * Optionally, filters these elements based on the provided range.
+ *
+ * @param range an optional range that defines a subset of the document; only elements within this range are included
+ * @return a list of non-whitespace child elements, filtered by the specified range if provided
+ */
+fun PsiFile.allNonWhitespaceChildren(
+    document: Document,
+    range: Range?,
+): List<PsiElement> {
+    val all = descendantsOfType<PsiElement>().filter { it !is PsiWhiteSpace }
+    if (range == null) return all.toList()
+    val textRange = range.toTextRange(document)
+    return all.filter { textRange.intersects(it.textRange) }.toList()
+}

features-impl/kotlin/src/com/jetbrains/ls/api/features/impl/common/kotlin/semanticTokens/LSSemanticTokensProviderKotlinImpl.kt

@@ -8,13 +8,20 @@ import com.intellij.openapi.vfs.findDocument
 import com.intellij.openapi.vfs.findPsiFile
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
+import com.intellij.psi.util.descendantsOfType
 import com.jetbrains.ls.api.core.LSServer
 import com.jetbrains.ls.api.core.util.findVirtualFile
 import com.jetbrains.ls.api.core.util.toLspRange
+import com.jetbrains.ls.api.core.util.toTextRange
 import com.jetbrains.ls.api.features.impl.common.kotlin.language.LSKotlinLanguage
 import com.jetbrains.ls.api.features.language.LSLanguage
 import com.jetbrains.ls.api.features.semanticTokens.*
+import com.jetbrains.ls.api.features.utils.allNonWhitespaceChildren
+import com.jetbrains.lsp.protocol.Range
 import com.jetbrains.lsp.protocol.SemanticTokensParams
+import com.jetbrains.lsp.protocol.SemanticTokensRangeParams
+import com.jetbrains.lsp.protocol.TextDocumentIdentifier
+import com.jetbrains.lsp.protocol.intersects
 import kotlinx.coroutines.CancellationException
 import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.kotlin.analysis.api.KaSession
@@ -41,12 +48,26 @@ object LSSemanticTokensProviderKotlinImpl : LSSemanticTokensProvider {
 
     context(LSServer)
     override suspend fun full(params: SemanticTokensParams): List<LSSemanticTokenWithRange> {
+        return getTokens(params.textDocument, range = null)
+
+    }
+
+    context(LSServer)
+    override suspend fun range(params: SemanticTokensRangeParams): List<LSSemanticTokenWithRange> {
+        return getTokens(params.textDocument, params.range)
+    }
+
+    /**
+     * @param range `null` means tokens from the whole file`
+     */
+    context(LSServer)
+    private suspend fun getTokens(textDocument: TextDocumentIdentifier, range:Range?): List<LSSemanticTokenWithRange> {
         return withAnalysisContext {
             runReadAction {
-                val file = params.textDocument.findVirtualFile() ?: return@runReadAction emptyList()
+                val file = textDocument.findVirtualFile() ?: return@runReadAction emptyList()
                 val ktFile = file.findPsiFile(project) as? KtFile ?: return@runReadAction emptyList()
                 val document = file.findDocument() ?: return@runReadAction emptyList()
-                val leafs = ktFile.allNonWhitespaceChildren()
+                val leafs = ktFile.allNonWhitespaceChildren(document, range)
                 if (leafs.isEmpty()) return@runReadAction emptyList()
                 analyze(ktFile) {
                     leafs.mapNotNull { element ->
@@ -181,11 +202,6 @@ object LSSemanticTokensProviderKotlinImpl : LSSemanticTokensProvider {
         }
     }
 
-
-    private fun KtFile.allNonWhitespaceChildren(): List<PsiElement> {
-        return collectDescendantsOfType<PsiElement>().filter { it !is PsiWhiteSpace }
-    }
-
     private fun FqName.isFromKotlinStdlib(): Boolean {
         return startsWith(StandardNames.BUILT_INS_PACKAGE_NAME)
     }

kotlin-lsp/src/com/jetbrains/ls/kotlinLsp/requests/core/initialize.kt

@@ -74,7 +74,7 @@ internal fun LspHandlersBuilder.initializeRequest() {
                             tokenModifiers = registry.modifiers.map { it.name },
                         )
                     },
-                    range = false,
+                    range = true,
                     full = true,
                 ),
                 completionProvider = CompletionRegistrationOptionsImpl(

kotlin-lsp/src/com/jetbrains/ls/kotlinLsp/requests/features.kt

@@ -19,6 +19,7 @@ import com.jetbrains.lsp.protocol.CodeActions.CodeActionRequest
 import com.jetbrains.lsp.protocol.Commands.ExecuteCommand
 import com.jetbrains.lsp.protocol.Diagnostics.DocumentDiagnosticRequestType
 import com.jetbrains.lsp.protocol.SemanticTokensRequests.SemanticTokensFullRequest
+import com.jetbrains.lsp.protocol.SemanticTokensRequests.SemanticTokensRangeRequest
 import com.jetbrains.lsp.protocol.WorkspaceSymbolRequests.WorkspaceSymbolRequest
 
 context(LSServer, LSConfiguration)
@@ -32,6 +33,7 @@ internal fun LspHandlersBuilder.features() {
     request(HoverRequestType) { LSHover.getHover(it) }
     request(ReferenceRequestType) { LSReferences.getReferences(it) }
     request(SemanticTokensFullRequest) { LSSemanticTokens.semanticTokensFull(it) }
+    request(SemanticTokensRangeRequest) { LSSemanticTokens.semanticTokensRange(it) }
     request(WorkspaceSymbolRequest) { LSWorkspaceSymbols.getSymbols(it) }
     request(DocumentSymbolRequest) { LSDocumentSymbols.getSymbols(it) }
 }