Fix enum references in Language Server

language-server/src/main/kotlin/tools/samt/ls/SamtDeclarationLookup.kt

@@ -50,6 +50,11 @@ class SamtDeclarationLookup private constructor() : SamtSemanticLookup<Location,
         this[aliasType.declaration.name.location] = aliasType
     }
 
+    override fun markEnumDeclaration(enumType: EnumType) {
+        super.markEnumDeclaration(enumType)
+        this[enumType.declaration.name.location] = enumType
+    }
+
     companion object {
         fun analyze(fileNode: FileNode, samtPackage: Package) =
             SamtDeclarationLookup().also { it.analyze(fileNode, samtPackage) }

language-server/src/test/kotlin/tools/samt/ls/SamtDeclarationLookupTest.kt

@@ -16,10 +16,17 @@ class SamtDeclarationLookupTest {
     fun `correctly find definition in complex model`() {
         val serviceSource = """
             package test
+
+            enum Friendliness {
+                FRIENDLY,
+                NEUTRAL,
+                HOSTILE
+            }
 
             record Person {
                 name: List<String? (size(1..100))>
                 age: Int
+                friendliness: Friendliness
             }
 
             service PersonService {
@@ -52,8 +59,9 @@ class SamtDeclarationLookupTest {
         """.trimIndent()
         parseAndCheck(
             serviceSource to listOf(
-                ExpectedDefinition("8:31" to "8:37") { it is RecordDeclarationNode && it.name.name == "Person" },
-                ExpectedDefinition("8:52" to "8:58") { it is RecordDeclarationNode && it.name.name == "Person" },
+                ExpectedDefinition("11:18" to "11:30") { it is EnumDeclarationNode && it.name.name == "Friendliness" },
+                ExpectedDefinition("15:31" to "15:37") { it is RecordDeclarationNode && it.name.name == "Person" },
+                ExpectedDefinition("15:52" to "15:58") { it is RecordDeclarationNode && it.name.name == "Person" },
             ),
             providerSource to listOf(
                 ExpectedDefinition("3:15" to "3:28") { it is ServiceDeclarationNode && it.name.name == "PersonService" },
@@ -76,6 +84,12 @@ class SamtDeclarationLookupTest {
     fun `finds definition for name of the user defined types themselves`() {
         val serviceSource = """
             package test
+
+            enum Friendliness {
+                FRIENDLY,
+                NEUTRAL,
+                HOSTILE
+            }
 
             record Person {
                 name: List<String? (size(1..100))>
@@ -97,9 +111,10 @@ class SamtDeclarationLookupTest {
         """.trimIndent()
         parseAndCheck(
             serviceSource to listOf(
-                ExpectedDefinition("2:7" to "2:13") { it is RecordDeclarationNode && it.name.name == "Person" },
-                ExpectedDefinition("7:8" to "7:21") { it is ServiceDeclarationNode && it.name.name == "PersonService" },
-                ExpectedDefinition("8:4" to "8:7") { it is OperationNode && it.name.name == "foo" },
+                ExpectedDefinition("2:5" to "2:17") { it is EnumDeclarationNode && it.name.name == "Friendliness" },
+                ExpectedDefinition("8:7" to "8:13") { it is RecordDeclarationNode && it.name.name == "Person" },
+                ExpectedDefinition("13:8" to "13:21") { it is ServiceDeclarationNode && it.name.name == "PersonService" },
+                ExpectedDefinition("14:4" to "14:7") { it is OperationNode && it.name.name == "foo" },
             ),
             providerSource to listOf(
                 ExpectedDefinition("2:8" to "2:22") { it is ProviderDeclarationNode && it.name.name == "PersonEndpoint" },

language-server/src/test/kotlin/tools/samt/ls/SamtReferencesLookupTest.kt

@@ -14,8 +14,16 @@ class SamtReferencesLookupTest {
     fun `correctly find references in complex model`() {
         val serviceSource = """
             package test
+
+            enum Friendliness {
+                FRIENDLY,
+                NEUTRAL,
+                HOSTILE
+            }
 
-            record Person { }
+            record Person {
+                friendliness: Friendliness
+            }
 
             service PersonService {
                 getNeighbors(person: Person): Map<String, Person?>? (*..100)
@@ -48,16 +56,22 @@ class SamtReferencesLookupTest {
         val (samtPackage, referencesLookup) = parse(serviceSource, providerSource, consumerOneSource, consumerTwoSource)
 
         val testPackage = samtPackage.subPackages.single { it.name == "test" }
+        val friendliness = testPackage.enums.single { it.name == "Friendliness" }
         val person = testPackage.records.single { it.name == "Person" }
         val personService = testPackage.services.single { it.name == "PersonService" }
         val personEndpoint = testPackage.providers.single { it.name == "PersonEndpoint" }
         val getNeighbors = personService.operations.single { it.name == "getNeighbors" }
 
+        val friendlinessReferences = referencesLookup[friendliness]
+        assertNotNull(friendlinessReferences)
+        assertEquals(1, friendlinessReferences.size, "Following list had unexpected amount of entries: $friendlinessReferences")
+        assertContains(friendlinessReferences, TestLocation("9:18" to "9:30").getLocation(serviceSource))
+
         val personReferences = referencesLookup[person]
         assertNotNull(personReferences)
         assertEquals(2, personReferences.size, "Following list had unexpected amount of entries: $personReferences")
-        assertContains(personReferences, TestLocation("5:25" to "5:31").getLocation(serviceSource))
-        assertContains(personReferences, TestLocation("5:46" to "5:52").getLocation(serviceSource))
+        assertContains(personReferences, TestLocation("13:25" to "13:31").getLocation(serviceSource))
+        assertContains(personReferences, TestLocation("13:46" to "13:52").getLocation(serviceSource))
 
         val personServiceReferences = referencesLookup[personService]
         assertNotNull(personServiceReferences)