Semantic highlight enum variants as enumMember token

Author: GrayJack

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

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.psi.KtVariableDeclaration
 import org.jetbrains.kotlin.psi.KtNamedDeclaration
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtParameter
+import org.jetbrains.kotlin.psi.KtEnumEntry
 import org.jetbrains.kotlin.psi.KtStringTemplateEntry
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameStringTemplateEntry
@@ -135,11 +136,13 @@ private fun elementToken(element: PsiElement, bindingContext: BindingContext): S
                 is PropertyDescriptor -> SemanticTokenType.PROPERTY
                 is VariableDescriptor -> SemanticTokenType.VARIABLE
                 is ConstructorDescriptor -> when (target.constructedClass.kind) {
+                    ClassKind.ENUM_ENTRY -> SemanticTokenType.ENUM_MEMBER
                     ClassKind.ANNOTATION_CLASS -> SemanticTokenType.TYPE // annotations look nicer this way
                     else -> SemanticTokenType.FUNCTION
                 }
                 is FunctionDescriptor -> SemanticTokenType.FUNCTION
                 is ClassDescriptor -> when (target.kind) {
+                    ClassKind.ENUM_ENTRY -> SemanticTokenType.ENUM_MEMBER
                     ClassKind.CLASS -> SemanticTokenType.CLASS
                     ClassKind.OBJECT -> SemanticTokenType.CLASS
                     ClassKind.INTERFACE -> SemanticTokenType.INTERFACE
@@ -160,6 +163,7 @@ private fun elementToken(element: PsiElement, bindingContext: BindingContext): S
             val tokenType = when (element) {
                 is KtParameter -> SemanticTokenType.PARAMETER
                 is KtProperty -> SemanticTokenType.PROPERTY
+                is KtEnumEntry -> SemanticTokenType.ENUM_MEMBER
                 is KtVariableDeclaration -> SemanticTokenType.VARIABLE
                 is KtClassOrObject -> SemanticTokenType.CLASS
                 is KtFunction -> SemanticTokenType.FUNCTION

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

@@ -14,6 +14,7 @@ class SemanticTokensTest : SingleFileTestFixture("semantictokens", "SemanticToke
         val constLine = 2
         val classLine = 4
         val funLine = 6
+        val enumLine = 8
 
         val expectedVar = sequenceOf(
             SemanticToken(range(varLine, 5, varLine, 13), SemanticTokenType.PROPERTY, setOf(SemanticTokenModifier.DECLARATION)), // variable
@@ -38,12 +39,16 @@ class SemanticTokensTest : SingleFileTestFixture("semantictokens", "SemanticToke
             SemanticToken(range(funLine, 30, funLine, 31), SemanticTokenType.FUNCTION), // f
             SemanticToken(range(funLine, 32, funLine, 33), SemanticTokenType.VARIABLE, setOf(SemanticTokenModifier.READONLY)), // x
         )
+        val expectedEnum = sequenceOf(
+            SemanticToken(range(enumLine, 12, enumLine, 16), SemanticTokenType.CLASS, setOf(SemanticTokenModifier.DECLARATION)), // Enum
+            SemanticToken(range(enumLine, 19, enumLine, 27), SemanticTokenType.ENUM_MEMBER, setOf(SemanticTokenModifier.DECLARATION)) // Variant1
+        )
 
         val partialExpected = encodeTokens(expectedConst + expectedClass)
         val partialResponse = languageServer.textDocumentService.semanticTokensRange(semanticTokensRangeParams(file, range(constLine, 0, classLine + 1, 0))).get()!!
         assertThat(partialResponse.data, contains(*partialExpected.toTypedArray()))
 
-        val fullExpected = encodeTokens(expectedVar + expectedConst + expectedClass + expectedFun)
+        val fullExpected = encodeTokens(expectedVar + expectedConst + expectedClass + expectedFun + expectedEnum)
         val fullResponse = languageServer.textDocumentService.semanticTokensFull(semanticTokensParams(file)).get()!!
         assertThat(fullResponse.data, contains(*fullExpected.toTypedArray()))
     }

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

@@ -4,3 +4,5 @@ val constant: String = "test $variable"
 data class Type(val property: Int)
 
 fun f(x: Int? = null): Int = f(x)
+
+enum class Enum { Variant1 }