From 69827577001e917e3c6f55aff0e11124abda6d9e Mon Sep 17 00:00:00 2001
From: Evgeny Gerashchenko <Evgeny.Gerashchenko@jetbrains.com>
Date: Tue, 3 Mar 2015 19:07:23 +0300
Subject: [PATCH] Optimized method. Moved logic about package view to
 JetReference.

---
 .../DescriptorToDeclarationUtil.kt            | 30 ++++++++-----------
 .../kotlin/idea/references/JetReference.kt    | 24 ++++++++++-----
 2 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/codeInsight/DescriptorToDeclarationUtil.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/codeInsight/DescriptorToDeclarationUtil.kt
index 171c936b41a..a81fc01f894 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/codeInsight/DescriptorToDeclarationUtil.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/codeInsight/DescriptorToDeclarationUtil.kt
@@ -32,26 +32,22 @@ import org.jetbrains.kotlin.descriptors.*
 // TODO rename it it DescriptorToSourceIde
 public object DescriptorToDeclarationUtil {
     public fun getDeclaration(project: Project, descriptor: DeclarationDescriptor): PsiElement? {
-        // TODO optimize it!
-        return resolveToPsiElements(project, descriptor).firstOrNull()
+        return getPsiElementsStream(project, descriptor).firstOrNull()
+    }
+
+    private fun getPsiElementsStream(project: Project, targetDescriptor: DeclarationDescriptor): Stream<PsiElement> {
+        val effectiveReferencedDescriptors = DescriptorToSourceUtils.getEffectiveReferencedDescriptors(targetDescriptor).stream()
+        return effectiveReferencedDescriptors.flatMap {
+            streamOf(
+                    DescriptorToSourceUtils.getSourceFromDescriptor(it),
+                    findBuiltinDeclaration(project, it),
+                    DecompiledNavigationUtils.getDeclarationFromDecompiledClassFile(project, it)
+            )
+        }.filterNotNull()
     }
 
     public fun resolveToPsiElements(project: Project, targetDescriptor: DeclarationDescriptor): Collection<PsiElement> {
-        val result = HashSet<PsiElement>()
-
-        val effectiveReferencedDescriptors: Collection<DeclarationDescriptor> = DescriptorToSourceUtils.getEffectiveReferencedDescriptors(targetDescriptor)
-
-        for (descriptor in effectiveReferencedDescriptors) {
-            result.addIfNotNull(DescriptorToSourceUtils.getSourceFromDescriptor(descriptor))
-            result.addIfNotNull(findBuiltinDeclaration(project, descriptor))
-            result.addIfNotNull(DecompiledNavigationUtils.getDeclarationFromDecompiledClassFile(project, descriptor))
-        }
-
-        if (targetDescriptor is PackageViewDescriptor) {
-            val psiFacade = JavaPsiFacade.getInstance(project)
-            val fqName = targetDescriptor.getFqName().asString()
-            result.addIfNotNull(psiFacade.findPackage(fqName))
-        }
+        val result = getPsiElementsStream(project, targetDescriptor).toHashSet()
         // filter out elements which are navigate to some other element of the result
         // this is needed to avoid duplicated results for references to declaration in same library source file
         return result.filter { element -> result.none { element != it && it.getNavigationElement() == element } }
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetReference.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetReference.kt
index de16df05bd9..15ecb5bbcd0 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetReference.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetReference.kt
@@ -16,14 +16,10 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.psi.PsiPolyVariantReference
-import com.intellij.psi.PsiPolyVariantReferenceBase
-import com.intellij.psi.PsiElement
-import com.intellij.psi.ResolveResult
-import com.intellij.psi.PsiElementResolveResult
+import com.intellij.psi.*
 import com.intellij.util.IncorrectOperationException
-import com.intellij.psi.PsiReference
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.PackageViewDescriptor
 import org.jetbrains.kotlin.resolve.BindingContext
 import java.util.Collections
 import org.jetbrains.kotlin.psi.JetReferenceExpression
@@ -31,6 +27,7 @@ import org.jetbrains.kotlin.idea.codeInsight.DescriptorToDeclarationUtil
 import org.jetbrains.kotlin.psi.JetElement
 import org.jetbrains.kotlin.utils.keysToMap
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.utils.addToStdlib.singletonOrEmptyList
 
 public trait JetReference : PsiPolyVariantReference {
     public fun resolveToDescriptors(): Collection<DeclarationDescriptor>
@@ -81,12 +78,12 @@ public abstract class AbstractJetReference<T : JetElement>(element: T)
     }
 
     override fun resolveMap(): Map<DeclarationDescriptor, Collection<PsiElement>> {
-        return getTargetDescriptors(expression.analyze()) keysToMap { DescriptorToDeclarationUtil.resolveToPsiElements(expression.getProject(), it) }
+        return getTargetDescriptors(expression.analyze()) keysToMap { resolveToPsiElements(it) }
     }
 
     private fun resolveToPsiElements(context: BindingContext, targetDescriptors: Collection<DeclarationDescriptor>): Collection<PsiElement> {
         if (targetDescriptors.isNotEmpty()) {
-            return targetDescriptors flatMap { target -> DescriptorToDeclarationUtil.resolveToPsiElements(expression.getProject(), target) }
+            return targetDescriptors flatMap { target -> resolveToPsiElements(target) }
         }
 
         val labelTargets = getLabelTargets(context)
@@ -97,6 +94,17 @@ public abstract class AbstractJetReference<T : JetElement>(element: T)
         return Collections.emptySet()
     }
 
+    private fun resolveToPsiElements(targetDescriptor: DeclarationDescriptor): Collection<PsiElement> {
+        if (targetDescriptor is PackageViewDescriptor) {
+            val psiFacade = JavaPsiFacade.getInstance(expression.getProject())
+            val fqName = targetDescriptor.getFqName().asString()
+            return psiFacade.findPackage(fqName).singletonOrEmptyList()
+        }
+        else {
+            return DescriptorToDeclarationUtil.resolveToPsiElements(expression.getProject(), targetDescriptor)
+        }
+    }
+
     protected abstract fun getTargetDescriptors(context: BindingContext): Collection<DeclarationDescriptor>
 
     private fun getLabelTargets(context: BindingContext): Collection<PsiElement>? {