Use partial resolve multiple in operator reference search

2016-09-16 21:31:55 +03:00
parent f181458a58
commit 4cc16e42b2
7 changed files with 27 additions and 21 deletions
@@ -24,6 +24,7 @@ import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.lexer.KtSingleValueToken
 import org.jetbrains.kotlin.psi.KtBinaryExpression
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance

@@ -49,7 +50,7 @@ class BinaryOperatorReferenceSearcher(
        return element.getReferencedNameElementType() in operationTokens
    }

-    override fun extractReference(element: PsiElement): PsiReference? {
+    override fun extractReference(element: KtElement): PsiReference? {
        val binaryExpression = element as? KtBinaryExpression ?: return null
        if (binaryExpression.operationToken !in operationTokens) return null
        return binaryExpression.operationReference.references.firstIsInstance<KtSimpleNameReference>()
@@ -23,10 +23,7 @@ import com.intellij.psi.search.SearchScope
 import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.KtBinaryExpression
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtOperationReferenceExpression
-import org.jetbrains.kotlin.psi.KtWhenConditionInRange
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance

 class ContainsOperatorReferenceSearcher(
@@ -59,7 +56,7 @@ class ContainsOperatorReferenceSearcher(
        return element.getReferencedNameElementType() in OPERATION_TOKENS
    }

-    override fun extractReference(element: PsiElement): PsiReference? {
+    override fun extractReference(element: KtElement): PsiReference? {
        val referenceExpression = element as? KtOperationReferenceExpression ?: return null
        if (referenceExpression.getReferencedNameElementType() !in OPERATION_TOKENS) return null
        return referenceExpression.references.firstIsInstance<KtSimpleNameReference>()
@@ -44,7 +44,7 @@ class DestructuringDeclarationReferenceSearcher(
        }
    }

-    override fun extractReference(element: PsiElement): PsiReference? {
+    override fun extractReference(element: KtElement): PsiReference? {
        val destructuringDeclaration = element as? KtDestructuringDeclaration ?: return null
        val entries = destructuringDeclaration.entries
        if (entries.size < componentIndex) return null
@@ -24,6 +24,7 @@ import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.references.KtArrayAccessReference
 import org.jetbrains.kotlin.idea.references.readWriteAccess
 import org.jetbrains.kotlin.psi.KtArrayAccessExpression
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
@@ -45,7 +46,7 @@ class IndexingOperatorReferenceSearcher(

    override fun isReferenceToCheck(ref: PsiReference) = ref is KtArrayAccessReference && checkAccessExpression(ref.element as KtArrayAccessExpression)

-    override fun extractReference(element: PsiElement): PsiReference? {
+    override fun extractReference(element: KtElement): PsiReference? {
        val accessExpression = element as? KtArrayAccessExpression ?: return null
        if (!checkAccessExpression(accessExpression)) return null
        return accessExpression.references.firstIsInstance<KtArrayAccessReference>()
@@ -23,6 +23,7 @@ import com.intellij.psi.search.SearchScope
 import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.references.KtInvokeFunctionReference
 import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance

@@ -40,7 +41,7 @@ class InvokeOperatorReferenceSearcher(

    override fun isReferenceToCheck(ref: PsiReference) = ref is KtInvokeFunctionReference

-    override fun extractReference(element: PsiElement): PsiReference? {
+    override fun extractReference(element: KtElement): PsiReference? {
        return (element as? KtCallExpression)?.references?.firstIsInstance<KtInvokeFunctionReference>()
    }
 }
@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.idea.caches.resolve.getJavaOrKotlinMemberDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinReferencesSearchOptions
 import org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinRequestResultProcessor
@@ -39,8 +40,10 @@ import org.jetbrains.kotlin.idea.util.fuzzyExtensionReceiverType
 import org.jetbrains.kotlin.idea.util.toFuzzyType
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.resolve.DataClassDescriptorResolver
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.util.isValidOperator
@@ -64,7 +67,7 @@ abstract class OperatorReferenceSearcher<TReferenceElement : KtElement>(
    /**
     * Extract reference that may resolve to our operator (no actual resolve to be performed)
     */
-    protected abstract fun extractReference(element: PsiElement): PsiReference?
+    protected abstract fun extractReference(element: KtElement): PsiReference?

    /**
     * Check if reference may potentially resolve to our operator (no actual resolve to be performed)
@@ -221,17 +224,19 @@ abstract class OperatorReferenceSearcher<TReferenceElement : KtElement>(
            for (element in scope.scope) {
                runReadAction {
                    if (element.isValid) {
-                        element.accept(object : PsiRecursiveElementWalkingVisitor() {
-                            override fun visitElement(element: PsiElement) {
-                                //TODO: resolve of multiple references at once
-                                val reference = extractReference(element)
-                                if (reference != null && reference.isReferenceTo(targetDeclaration)) {
-                                    consumer.process(reference)
-                                }
+                        val refs = ArrayList<PsiReference>()
+                        val elements = element.collectDescendantsOfType<KtElement> {
+                            val ref = extractReference(it) ?: return@collectDescendantsOfType false
+                            refs.add(ref)
+                            true
+                        }

-                                super.visitElement(element)
-                            }
-                        })
+                        // resolve all references at once
+                        (element.containingFile as KtFile).getResolutionFacade().analyze(elements, BodyResolveMode.PARTIAL)
+
+                        refs
+                                .filter { it.isReferenceTo(targetDeclaration) }
+                                .forEach { consumer.process(it) }
                    }
                }
            }
@@ -23,6 +23,7 @@ import com.intellij.psi.search.SearchScope
 import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.lexer.KtSingleValueToken
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtUnaryExpression
@@ -49,7 +50,7 @@ class UnaryOperatorReferenceSearcher(
        return element.getReferencedNameElementType() == operationToken
    }

-    override fun extractReference(element: PsiElement): PsiReference? {
+    override fun extractReference(element: KtElement): PsiReference? {
        val unaryExpression = element as? KtUnaryExpression ?: return null
        if (unaryExpression.operationToken != operationToken) return null
        return unaryExpression.operationReference.references.firstIsInstance<KtSimpleNameReference>()