Spring Support: Replace light elements bound to line markers with their originals

#KT-12091 Fixed

ChangeLog.md

@@ -51,6 +51,7 @@ Issues fixed:
 - [KT-12018](https://youtrack.jetbrains.com/issue/KT-12018) Auto-format spaces between function name and arguments in infix calls
 - [KT-12067](https://youtrack.jetbrains.com/issue/KT-12067) Deadlock in Kotlin debugger is fixed
 - [KT-12070](https://youtrack.jetbrains.com/issue/KT-12070) Add empty line in error message of Maven and Gradle configuration
+- [KT-12091](https://youtrack.jetbrains.com/issue/KT-12091) Fixed unstable behavior of Spring line markers
 
 #### Debugger
 

compiler/light-classes/src/org/jetbrains/kotlin/asJava/KtLightIdentifier.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.asJava
 
 import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNameIdentifierOwner
 import com.intellij.psi.impl.light.LightIdentifier
 import org.jetbrains.kotlin.psi.KtNamedDeclaration
@@ -25,6 +26,9 @@ class KtLightIdentifier(
         private val lightOwner: PsiNameIdentifierOwner,
         private val ktDeclaration: KtNamedDeclaration?
 ) : LightIdentifier(lightOwner.manager, ktDeclaration?.name ?: "") {
+    val origin: PsiElement?
+        get() = ktDeclaration?.nameIdentifier
+
     override fun isPhysical() = true
     override fun getParent() = lightOwner
     override fun getTextRange() = ktDeclaration?.nameIdentifier?.textRange ?: TextRange.EMPTY_RANGE

ultimate/src/org/jetbrains/kotlin/idea/spring/lineMarking/KotlinSpringClassAnnotator.kt

@@ -16,16 +16,23 @@
 
 package org.jetbrains.kotlin.idea.spring.lineMarking
 
+import com.intellij.codeInsight.daemon.LineMarkerInfo
 import com.intellij.codeInsight.daemon.RelatedItemLineMarkerInfo
+import com.intellij.navigation.GotoRelatedItem
+import com.intellij.openapi.editor.markup.GutterIconRenderer
+import com.intellij.openapi.util.NotNullLazyValue
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNameIdentifierOwner
 import com.intellij.spring.gutter.SpringClassAnnotator
+import com.intellij.util.Function
+import com.intellij.util.SmartList
 import org.jetbrains.kotlin.asJava.*
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
 import org.jetbrains.kotlin.resolve.annotations.hasJvmStaticAnnotation
+import javax.swing.Icon
 
 class KotlinSpringClassAnnotator : SpringClassAnnotator() {
     override fun getElementToProcess(psiElement: PsiElement): PsiElement? {
@@ -50,7 +57,7 @@ class KotlinSpringClassAnnotator : SpringClassAnnotator() {
         return null
     }
 
-    override fun collectNavigationMarkers(psiElement: PsiElement, result: MutableCollection<in RelatedItemLineMarkerInfo<PsiElement>>?) {
+    private fun doCollectMarkers(psiElement: PsiElement, result: MutableCollection<in RelatedItemLineMarkerInfo<PsiElement>>?) {
         if (psiElement is KtProperty) {
             for (it in psiElement.toLightElements()) {
                 val nameIdentifier = (it as? PsiNameIdentifierOwner)?.nameIdentifier ?: continue
@@ -64,4 +71,41 @@ class KotlinSpringClassAnnotator : SpringClassAnnotator() {
 
         super.collectNavigationMarkers(psiElement, result)
     }
+
+    // TODO
+    // Weak references to light elements may be reclaimed by GC after original file is modified causing line markers to misbehave
+    // This workaround allows reuse of SpringClassAnnotator logic and avoids binding of line markers to light elements
+
+    private val toolTipProviderField by lazy { LineMarkerInfo::class.java.getDeclaredField("myTooltipProvider").apply { isAccessible = true } }
+    private val iconField by lazy { LineMarkerInfo::class.java.getDeclaredField("myIcon").apply { isAccessible = true } }
+    private val iconAlignmentField by lazy { LineMarkerInfo::class.java.getDeclaredField("myIconAlignment").apply { isAccessible = true } }
+    private val targetsField by lazy { RelatedItemLineMarkerInfo::class.java.getDeclaredField("myTargets").apply { isAccessible = true } }
+
+    override fun collectNavigationMarkers(psiElement: PsiElement, result: MutableCollection<in RelatedItemLineMarkerInfo<PsiElement>>) {
+        val newItems = SmartList<RelatedItemLineMarkerInfo<PsiElement>>()
+
+        doCollectMarkers(psiElement, newItems)
+
+        newItems.mapNotNullTo(result) { item ->
+            val itemElement = item.element
+            val elementToAnnotate = when (itemElement) {
+                is KtLightIdentifier -> itemElement.origin
+                is KtLightElement<*, *> -> itemElement.kotlinOrigin
+                else -> return@mapNotNullTo item
+            }
+            if (elementToAnnotate == null) return@mapNotNullTo null
+
+            @Suppress("UNCHECKED_CAST")
+            RelatedItemLineMarkerInfo<PsiElement>(
+                    elementToAnnotate,
+                    elementToAnnotate.textRange,
+                    iconField.get(item) as Icon?,
+                    item.updatePass,
+                    toolTipProviderField.get(item) as Function<PsiElement, String>,
+                    item.navigationHandler,
+                    iconAlignmentField.get(item) as GutterIconRenderer.Alignment,
+                    targetsField.get(item) as NotNullLazyValue<Collection<GotoRelatedItem>>
+            )
+        }
+    }
 }