Use partial resolve multiple in operator reference search

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/BinaryOperatorReferenceSearcher.kt

@@ -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>()

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/ContainsOperatorReferenceSearcher.kt

@@ -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.*
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtOperationReferenceExpression
-import org.jetbrains.kotlin.psi.KtWhenConditionInRange
 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>()

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/DestructuringDeclarationReferenceSearcher.kt

@@ -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

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/IndexingOperatorReferenceSearcher.kt

@@ -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>()

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/InvokeOperatorReferenceSearcher.kt

@@ -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>()
     }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/OperatorReferenceSearcher.kt

@@ -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() {
+                        val refs = ArrayList<PsiReference>()
-                            override fun visitElement(element: PsiElement) {
+                        val elements = element.collectDescendantsOfType<KtElement> {
-                                //TODO: resolve of multiple references at once
+                            val ref = extractReference(it) ?: return@collectDescendantsOfType false
-                                val reference = extractReference(element)
+                            refs.add(ref)
-                                if (reference != null && reference.isReferenceTo(targetDeclaration)) {
+                            true
-                                    consumer.process(reference)
+                        }
-                                }
 
-                                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) }
                     }
                 }
             }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/usagesSearch/UnaryOperatorReferenceSearcher.kt

@@ -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>()