KtLightPsiJavaCodeReferenceElement reference getting made lazy (KT-26036, EA-125884)

Collecting references from contributors could be expensive and even recursive (EA-125884)

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

@@ -5,20 +5,22 @@
 
 package org.jetbrains.kotlin.asJava.elements
 
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.*
 import com.intellij.psi.scope.PsiScopeProcessor
+import com.intellij.util.IncorrectOperationException
 import org.jetbrains.kotlin.asJava.classes.lazyPub
 
 class KtLightPsiJavaCodeReferenceElement(
     private val ktElement: PsiElement,
-    private val reference: PsiReference,
+    reference: () -> PsiReference?,
-    private val clsDelegateProvider: () -> PsiJavaCodeReferenceElement?
+    clsDelegateProvider: () -> PsiJavaCodeReferenceElement?
 ) :
     PsiElement by ktElement,
-    PsiReference by reference,
+    PsiReference by LazyPsiReferenceDelegate(ktElement, reference),
     PsiJavaCodeReferenceElement {
 
-    private val delegate by lazyPub { clsDelegateProvider() }
+    private val delegate by lazyPub(clsDelegateProvider)
 
     override fun advancedResolve(incompleteCode: Boolean): JavaResolveResult =
         delegate?.advancedResolve(incompleteCode) ?: JavaResolveResult.EMPTY
@@ -43,3 +45,32 @@ class KtLightPsiJavaCodeReferenceElement(
 
     override fun getParameterList(): PsiReferenceParameterList? = delegate?.parameterList
 }
+
+private class LazyPsiReferenceDelegate(
+    private val psiElement: PsiElement,
+    referenceProvider: () -> PsiReference?
+) : PsiReference {
+
+    private val delegate by lazyPub(referenceProvider)
+
+    override fun getElement(): PsiElement = psiElement
+
+    override fun resolve(): PsiElement? = delegate?.resolve()
+
+    override fun getRangeInElement(): TextRange = delegate?.rangeInElement ?: psiElement.textRange
+
+    override fun getCanonicalText(): String = delegate?.canonicalText ?: "<no-text>"
+
+    @Throws(IncorrectOperationException::class)
+    override fun handleElementRename(newElementName: String): PsiElement = delegate?.handleElementRename(newElementName) ?: element
+
+    @Throws(IncorrectOperationException::class)
+    override fun bindToElement(element: PsiElement): PsiElement =
+        delegate?.bindToElement(element) ?: throw IncorrectOperationException("can't rename LazyPsiReferenceDelegate")
+
+    override fun isSoft(): Boolean = delegate?.isSoft ?: false
+
+    override fun isReferenceTo(element: PsiElement): Boolean = delegate?.isReferenceTo(element) ?: false
+
+    override fun getVariants(): Array<Any> = delegate?.variants ?: emptyArray()
+}

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

@@ -127,15 +127,14 @@ class KtLightAnnotationForSourceEntry(
     }
 
 
-    override fun getNameReferenceElement(): PsiJavaCodeReferenceElement? {
+    override fun getNameReferenceElement(): PsiJavaCodeReferenceElement? = KtLightPsiJavaCodeReferenceElement(
-        val reference = (kotlinOrigin as? KtAnnotationEntry)?.typeReference?.reference
+        kotlinOrigin.navigationElement,
+        {
+            (kotlinOrigin as? KtAnnotationEntry)?.typeReference?.reference
                 ?: (kotlinOrigin.calleeExpression?.nameReference)?.references?.firstOrNull()
-                ?: return null
+        },
-        return KtLightPsiJavaCodeReferenceElement(
+        { super.getNameReferenceElement() }
-            kotlinOrigin.navigationElement,
+    )
-            reference,
-            { super.getNameReferenceElement() })
-    }
 
 
     private val ktLightAnnotationParameterList by lazyPub { KtLightAnnotationParameterList() }

compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt.173

@@ -129,15 +129,14 @@ class KtLightAnnotationForSourceEntry(
     }
 
 
-    override fun getNameReferenceElement(): PsiJavaCodeReferenceElement? {
+    override fun getNameReferenceElement(): PsiJavaCodeReferenceElement? = KtLightPsiJavaCodeReferenceElement(
-        val reference = (kotlinOrigin as? KtAnnotationEntry)?.typeReference?.reference
+        kotlinOrigin.navigationElement,
+        {
+            (kotlinOrigin as? KtAnnotationEntry)?.typeReference?.reference
                 ?: (kotlinOrigin.calleeExpression?.nameReference)?.references?.firstOrNull()
-                ?: return null
+        },
-        return KtLightPsiJavaCodeReferenceElement(
+        { super.getNameReferenceElement() }
-            kotlinOrigin.navigationElement,
+    )
-            reference,
-            { super.getNameReferenceElement() })
-    }
 
 
     private val ktLightAnnotationParameterList by lazyPub { KtLightAnnotationParameterList() }