Implement reference contributor for directives

Author: cypressious

resources/META-INF/plugin.xml

@@ -42,6 +42,8 @@
             <className>org.jetbrains.kotlin.test.helper.intentions.CreateContextualOverloadIntention</className>
             <category>Kotlin Compiler DevKit</category>
         </intentionAction>
+
+        <psi.referenceContributor implementation="org.jetbrains.kotlin.test.helper.reference.DirectiveReferenceContributor"/>
     </extensions>
 
     <extensions defaultExtensionNs="org.jetbrains.kotlin">

src/org/jetbrains/kotlin/test/helper/reference/DirectiveReferenceContributor.kt

@@ -0,0 +1,57 @@
+package org.jetbrains.kotlin.test.helper.reference
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.patterns.PlatformPatterns
+import com.intellij.psi.*
+import com.intellij.util.ProcessingContext
+import org.jetbrains.kotlin.test.helper.getTestDataType
+
+class DirectiveReferenceContributor : PsiReferenceContributor() {
+    val regex = Regex("// ([A-Z_]+)(:.*)?")
+    val regexLanguageFeature = Regex("[+-]([A-Za-z0-9_]+)")
+
+    override fun registerReferenceProviders(registrar: PsiReferenceRegistrar) {
+        registrar.registerReferenceProvider(
+            PlatformPatterns.psiElement(PsiComment::class.java),
+            object : PsiReferenceProvider() {
+                override fun getReferencesByElement(
+                    element: PsiElement,
+                    context: ProcessingContext
+                ): Array<PsiReference> {
+                    val file =
+                        element.containingFile?.virtualFile ?: return PsiReference.EMPTY_ARRAY
+                    if (file.getTestDataType(element.project) == null) return PsiReference.EMPTY_ARRAY
+
+                    val text = element.text
+                    val refs = mutableListOf<PsiReference>()
+
+                    val regex = regex
+                    regex.matchAt(text, 0)?.let { match ->
+                        val (name, matchRange) = match.groups[1] ?: return@let
+                        val range = TextRange(matchRange.first, matchRange.last + 1)
+                        refs.add(TestDirectiveReference(element, range, name))
+
+                        if (name == "LANGUAGE") {
+                            val (features, featuresRange) = match.groups[2] ?: return@let
+                            regexLanguageFeature.findAll(features).forEach {
+                                val (feature, featureRange) = it.groups[1] ?: return@forEach
+                                refs.add(
+                                    LanguageFeatureReference(
+                                        element,
+                                        TextRange(
+                                            featuresRange.first + featureRange.first,
+                                            featuresRange.first + featureRange.last + 1,
+                                        ),
+                                        feature
+                                    )
+                                )
+                            }
+                        }
+                    }
+
+                    return refs.toTypedArray()
+                }
+            }
+        )
+    }
+}

src/org/jetbrains/kotlin/test/helper/reference/LanguageFeatureReference.kt

@@ -0,0 +1,43 @@
+package org.jetbrains.kotlin.test.helper.reference
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.JavaPsiFacade
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiElementResolveResult
+import com.intellij.psi.PsiEnumConstant
+import com.intellij.psi.PsiPolyVariantReference
+import com.intellij.psi.PsiReferenceBase
+import com.intellij.psi.ResolveResult
+import com.intellij.psi.search.GlobalSearchScope
+
+private const val LANGUAGE_FEATURE_FQ_NAME = "org.jetbrains.kotlin.config.LanguageFeature"
+
+class LanguageFeatureReference(
+    element: PsiElement,
+    range: TextRange,
+    private val key: String
+) : PsiReferenceBase<PsiElement>(element, range), PsiPolyVariantReference {
+
+    override fun multiResolve(incompleteCode: Boolean): Array<ResolveResult> {
+        val project = myElement.project
+        val psiFacade = JavaPsiFacade.getInstance(project)
+        val classes = psiFacade
+            .findClasses(
+                LANGUAGE_FEATURE_FQ_NAME,
+                GlobalSearchScope.allScope(project)
+            )
+
+        return classes.mapNotNull { clazz ->
+            clazz.fields
+                .filterIsInstance<PsiEnumConstant>()
+                .firstOrNull { it.name == key }
+                ?.let { PsiElementResolveResult(it) }
+        }.toTypedArray()
+    }
+
+    override fun resolve(): PsiElement? {
+        return multiResolve(false).singleOrNull()?.element
+    }
+
+    override fun getVariants(): Array<Any> = emptyArray()
+}

src/org/jetbrains/kotlin/test/helper/reference/TestDirectiveReference.kt

@@ -0,0 +1,44 @@
+package org.jetbrains.kotlin.test.helper.reference
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiElementResolveResult
+import com.intellij.psi.PsiPolyVariantReference
+import com.intellij.psi.PsiReferenceBase
+import com.intellij.psi.ResolveResult
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.idea.stubindex.KotlinPropertyShortNameIndex
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+
+private val DIRECTIVE_CLASS_ID =
+    ClassId.topLevel(FqName("org.jetbrains.kotlin.test.directives.model.Directive"))
+
+class TestDirectiveReference(
+    element: PsiElement,
+    range: TextRange,
+    private val key: String
+) : PsiReferenceBase<PsiElement>(element, range), PsiPolyVariantReference {
+
+    override fun multiResolve(incompleteCode: Boolean): Array<ResolveResult> {
+        val project = myElement.project
+
+        val declarations =
+            KotlinPropertyShortNameIndex.Helper[key, project, GlobalSearchScope.allScope(project)]
+
+        return declarations
+            .filter {
+                analyze(it) {
+                    it.returnType.isSubtypeOf(DIRECTIVE_CLASS_ID)
+                }
+            }
+            .map { PsiElementResolveResult(it) }.toTypedArray()
+    }
+
+    override fun resolve(): PsiElement? {
+        return multiResolve(false).singleOrNull()?.element
+    }
+
+    override fun getVariants(): Array<Any> = emptyArray()
+}