Add caching for reference resolution in idea ResolveCache

idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/KtReference.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.*
+import com.intellij.psi.impl.source.resolve.ResolveCache
 import com.intellij.util.IncorrectOperationException
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.PackageViewDescriptor
@@ -38,18 +39,14 @@ interface KtReference : PsiPolyVariantReference {
     val resolvesByNames: Collection<Name>
 }
 
-abstract class AbstractKtReference<T : KtElement>(element: T)
+abstract class AbstractKtReference<T : KtElement>(element: T) : PsiPolyVariantReferenceBase<T>(element), KtReference {
-: PsiPolyVariantReferenceBase<T>(element), KtReference {
-
     val expression: T
         get() = element
 
     override fun multiResolve(incompleteCode: Boolean): Array<ResolveResult> {
-        return PsiElementResolveResult.createResults(resolveToPsiElements())
+        @Suppress("UNCHECKED_CAST")
-    }
+        val kotlinResolver = KOTLIN_RESOLVER as ResolveCache.PolyVariantResolver<AbstractKtReference<T>>
-
+        return ResolveCache.getInstance(expression.project).resolveWithCaching(this, kotlinResolver, false, incompleteCode)
-    override fun resolve(): PsiElement? {
-        return resolveToPsiElements().singleOrNull()
     }
 
     override fun isReferenceTo(element: PsiElement?): Boolean {
@@ -68,51 +65,64 @@ abstract class AbstractKtReference<T : KtElement>(element: T)
 
     override fun isSoft(): Boolean = false
 
-    private fun resolveToPsiElements(): Collection<PsiElement> {
-        val bindingContext = expression.analyze(BodyResolveMode.PARTIAL_WITH_DIAGNOSTICS)
-        return resolveToPsiElements(bindingContext, getTargetDescriptors(bindingContext))
-    }
-
     override fun resolveToDescriptors(bindingContext: BindingContext): Collection<DeclarationDescriptor> {
         return getTargetDescriptors(bindingContext)
     }
 
-    private fun resolveToPsiElements(context: BindingContext, targetDescriptors: Collection<DeclarationDescriptor>): Collection<PsiElement> {
-        if (targetDescriptors.isNotEmpty()) {
-            return targetDescriptors.flatMap { target -> resolveToPsiElements(target) }.toSet()
-        }
-
-        val labelTargets = getLabelTargets(context)
-        if (labelTargets != null) {
-            return labelTargets
-        }
-
-        return Collections.emptySet()
-    }
-
-    private fun resolveToPsiElements(targetDescriptor: DeclarationDescriptor): Collection<PsiElement> {
-        if (targetDescriptor is PackageViewDescriptor) {
-            val psiFacade = JavaPsiFacade.getInstance(expression.project)
-            val fqName = targetDescriptor.fqName.asString()
-            return listOfNotNull(psiFacade.findPackage(fqName))
-        }
-        else {
-            return DescriptorToSourceUtilsIde.getAllDeclarations(expression.project, targetDescriptor, expression.resolveScope)
-        }
-    }
-
     protected abstract fun getTargetDescriptors(context: BindingContext): Collection<DeclarationDescriptor>
 
-    private fun getLabelTargets(context: BindingContext): Collection<PsiElement>? {
-        val reference = expression as? KtReferenceExpression ?: return null
-        val labelTarget = context[BindingContext.LABEL_TARGET, reference]
-        if (labelTarget != null) {
-            return listOf(labelTarget)
-        }
-        return context[BindingContext.AMBIGUOUS_LABEL_TARGET, reference]
-    }
-
     override fun toString() = this::class.java.simpleName + ": " + expression.text
+
+    companion object {
+        class KotlinReferenceResolver : ResolveCache.PolyVariantResolver<AbstractKtReference<KtElement>> {
+            class KotlinResolveResult(element: PsiElement) : PsiElementResolveResult(element)
+
+            private fun resolveToPsiElements(ref: AbstractKtReference<KtElement>): Collection<PsiElement> {
+                val bindingContext = ref.expression.analyze(BodyResolveMode.PARTIAL_WITH_DIAGNOSTICS)
+                return resolveToPsiElements(ref, bindingContext, ref.getTargetDescriptors(bindingContext))
+            }
+
+            private fun resolveToPsiElements(ref: AbstractKtReference<KtElement>, context: BindingContext, targetDescriptors: Collection<DeclarationDescriptor>): Collection<PsiElement> {
+                if (targetDescriptors.isNotEmpty()) {
+                    return targetDescriptors.flatMap { target -> resolveToPsiElements(ref, target) }.toSet()
+                }
+
+                val labelTargets = getLabelTargets(ref, context)
+                if (labelTargets != null) {
+                    return labelTargets
+                }
+
+                return Collections.emptySet()
+            }
+
+            private fun resolveToPsiElements(ref: AbstractKtReference<KtElement>, targetDescriptor: DeclarationDescriptor): Collection<PsiElement> {
+                if (targetDescriptor is PackageViewDescriptor) {
+                    val psiFacade = JavaPsiFacade.getInstance(ref.expression.project)
+                    val fqName = targetDescriptor.fqName.asString()
+                    return listOfNotNull(psiFacade.findPackage(fqName))
+                }
+                else {
+                    return DescriptorToSourceUtilsIde.getAllDeclarations(ref.expression.project, targetDescriptor, ref.expression.resolveScope)
+                }
+            }
+
+            private fun getLabelTargets(ref: AbstractKtReference<KtElement>, context: BindingContext): Collection<PsiElement>? {
+                val reference = ref.expression as? KtReferenceExpression ?: return null
+                val labelTarget = context[BindingContext.LABEL_TARGET, reference]
+                if (labelTarget != null) {
+                    return listOf(labelTarget)
+                }
+                return context[BindingContext.AMBIGUOUS_LABEL_TARGET, reference]
+            }
+
+            override fun resolve(ref: AbstractKtReference<KtElement>, incompleteCode: Boolean): Array<ResolveResult> {
+                val resolveToPsiElements = resolveToPsiElements(ref)
+                return resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
+            }
+        }
+
+        val KOTLIN_RESOLVER = KotlinReferenceResolver()
+    }
 }
 
 abstract class KtSimpleReference<T : KtReferenceExpression>(expression: T) : AbstractKtReference<T>(expression) {