Implemented fast search for invoke operator calls

#KT-13643 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/ideaExtensions/KotlinReferencesSearcher.kt

@@ -37,22 +37,21 @@ import org.jetbrains.kotlin.idea.search.KOTLIN_NAMED_ARGUMENT_SEARCH_CONTEXT
 import org.jetbrains.kotlin.idea.search.allScope
 import org.jetbrains.kotlin.idea.search.effectiveSearchScope
 import org.jetbrains.kotlin.idea.search.unionSafe
-import org.jetbrains.kotlin.idea.search.usagesSearch.dataClassComponentFunction
+import org.jetbrains.kotlin.idea.search.usagesSearch.*
-import org.jetbrains.kotlin.idea.search.usagesSearch.findDestructuringDeclarationUsages
-import org.jetbrains.kotlin.idea.search.usagesSearch.getClassNameForCompanionObject
-import org.jetbrains.kotlin.idea.search.usagesSearch.getSpecialNamesToSearch
 import org.jetbrains.kotlin.idea.stubindex.KotlinSourceFilterScope
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.resolve.dataClassUtils.isComponentLike
+import org.jetbrains.kotlin.util.OperatorNameConventions
 
 data class KotlinReferencesSearchOptions(val acceptCallableOverrides: Boolean = false,
                                          val acceptOverloads: Boolean = false,
                                          val acceptExtensionsOfDeclarationClass: Boolean = false,
                                          val acceptCompanionObjectMembers: Boolean = false,
                                          val searchForComponentConventions: Boolean = true,
+                                         val searchInvokeOperator: Boolean = true,
                                          val searchNamedArguments: Boolean = true) {
     fun anyEnabled(): Boolean = acceptCallableOverrides || acceptOverloads || acceptExtensionsOfDeclarationClass
 
@@ -144,6 +143,11 @@ class KotlinReferencesSearcher : QueryExecutorBase<PsiReference, ReferencesSearc
                 }
             }
         }
+
+        //TODO: Java invoke's
+        if (kotlinOptions.searchInvokeOperator && element is KtFunction && name == OperatorNameConventions.INVOKE.asString()) {
+            findInvokeOperatorUsages(element, effectiveSearchScope, consumer)
+        }
     }
 
     private fun searchNamedArguments(parameter: KtParameter, queryParameters: ReferencesSearch.SearchParameters) {

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

@@ -51,6 +51,7 @@ import org.jetbrains.kotlin.load.java.sam.SingleAbstractMethodUtils
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.util.OperatorNameConventions
 import java.util.*
 
 //TODO: check if smart search is too expensive
@@ -592,7 +593,7 @@ class ExpressionsOfTypeProcessor(
     private fun PsiElement.isOperatorExpensiveToSearch(): Boolean {
         when (this) {
             is KtFunction -> {
-                if (name?.startsWith("component") == true) return false // component functions are not so expensive to search
+                if (name?.startsWith("component") == true || name == OperatorNameConventions.INVOKE.asString()) return false
                 return  hasModifier(KtTokens.OPERATOR_KEYWORD)
                         || hasModifier(KtTokens.OVERRIDE_KEYWORD) && (resolveToDescriptorIfAny() as? FunctionDescriptor)?.isOperator == true
             }

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

@@ -0,0 +1,136 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.search.usagesSearch
+
+import com.intellij.openapi.roots.ProjectRootManager
+import com.intellij.psi.PsiManager
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiReference
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.LocalSearchScope
+import com.intellij.psi.search.SearchScope
+import com.intellij.util.Processor
+import org.jetbrains.kotlin.asJava.elements.KtLightMethod
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.idea.references.KtInvokeFunctionReference
+import org.jetbrains.kotlin.idea.search.usagesSearch.ExpressionsOfTypeProcessor.Companion.testLog
+import org.jetbrains.kotlin.idea.util.fuzzyExtensionReceiverType
+import org.jetbrains.kotlin.idea.util.toFuzzyType
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
+import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
+import org.jetbrains.kotlin.util.isValidOperator
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
+
+//TODO: problem with infinite recursion not solved
+
+//TODO: code duplication
+
+fun findInvokeOperatorUsages(
+        invokeFunction: PsiMethod,
+        scope: SearchScope,
+        consumer: Processor<PsiReference>
+) {
+    if (invokeFunction !is KtLightMethod) return //TODO
+    val ktDeclarationTarget = invokeFunction.kotlinOrigin as? KtDeclaration ?: return //TODO?
+    findInvokeOperatorUsages(ktDeclarationTarget, scope, consumer)
+}
+
+fun findInvokeOperatorUsages(
+        targetDeclaration: KtDeclaration,
+        searchScope: SearchScope,
+        consumer: Processor<PsiReference>
+) {
+    val usePlainSearch = when (ExpressionsOfTypeProcessor.mode) {
+        ExpressionsOfTypeProcessor.Mode.ALWAYS_SMART -> false
+        ExpressionsOfTypeProcessor.Mode.ALWAYS_PLAIN -> true
+        ExpressionsOfTypeProcessor.Mode.PLAIN_WHEN_NEEDED -> searchScope is LocalSearchScope // for local scope it's faster to use plain search
+    }
+    if (usePlainSearch) {
+        doPlainSearch(targetDeclaration, searchScope, consumer)
+        return
+    }
+
+    val descriptor = targetDeclaration.resolveToDescriptor() as? CallableDescriptor ?: return
+
+    if (descriptor is FunctionDescriptor && !descriptor.isValidOperator()) return
+
+    val dataType = if (descriptor.isExtension) {
+        descriptor.fuzzyExtensionReceiverType()!!
+    }
+    else {
+        val classDescriptor = descriptor.containingDeclaration as? ClassDescriptor ?: return
+        classDescriptor.defaultType.toFuzzyType(classDescriptor.typeConstructor.parameters)
+    }
+
+    ExpressionsOfTypeProcessor(
+            dataType,
+            searchScope,
+            suspiciousExpressionHandler = { expression -> processSuspiciousExpression(expression, targetDeclaration, consumer) },
+            suspiciousScopeHandler = { searchScope -> doPlainSearch(targetDeclaration, searchScope, consumer) },
+            resolutionFacade = targetDeclaration.getResolutionFacade()
+    ).run()
+}
+
+private fun processSuspiciousExpression(expression: KtExpression, targetDeclaration: KtDeclaration, consumer: Processor<PsiReference>) {
+    val callExpression = expression.parent as? KtCallExpression ?: return
+    testLog?.add("Resolving call ${callExpression.text}")
+    processCallExpression(callExpression, targetDeclaration, consumer)
+}
+
+private fun doPlainSearch(ktDeclaration: KtDeclaration, scope: SearchScope, consumer: Processor<PsiReference>) {
+    if (scope is LocalSearchScope) {
+        for (element in scope.scope) {
+            val stop = element.anyDescendantOfType<KtCallExpression> { !processCallExpression(it, ktDeclaration, consumer) }
+            if (stop) break
+        }
+    }
+    else {
+        scope as GlobalSearchScope
+        val project = ktDeclaration.project
+        val psiManager = PsiManager.getInstance(project)
+        ProjectRootManager.getInstance(project).fileIndex.iterateContent { file ->
+            if (file in scope) {
+                val ktFile = psiManager.findFile(file) as? KtFile
+                if (ktFile != null) {
+                    val stop = ktFile.anyDescendantOfType<KtCallExpression> {
+                        !processCallExpression(it, ktDeclaration, consumer)
+                    }
+                    if (stop) return@iterateContent false
+                }
+            }
+            true
+        }
+    }
+}
+
+private fun processCallExpression(callExpression: KtCallExpression, targetDeclaration: KtDeclaration, consumer: Processor<PsiReference>): Boolean {
+    val reference = callExpression.references.firstIsInstance<KtInvokeFunctionReference>()
+    if (reference.isReferenceTo(targetDeclaration)) {
+        return consumer.process(reference)
+    }
+    else {
+        return true
+    }
+}

idea/src/org/jetbrains/kotlin/idea/findUsages/handlers/KotlinFindMemberUsagesHandler.kt

@@ -131,9 +131,9 @@ abstract class KotlinFindMemberUsagesHandler<T : KtNamedDeclaration>
         val uniqueProcessor = CommonProcessors.UniqueProcessor(processor)
 
         if (options.isUsages) {
-            // we disable searchForComponentConventions for ReferencesSearch because they will be searched by MethodReferencesSearch
+            // we disable searchForComponentConventions and searchInvokeOperator for ReferencesSearch because they will be searched by MethodReferencesSearch
             val kotlinSearchOptions = createKotlinReferencesSearchOptions(options)
-                    .copy(searchForComponentConventions = false)
+                    .copy(searchForComponentConventions = false, searchInvokeOperator = false)
             val searchParameters = KotlinReferencesSearchParameters(element, options.searchScope, kotlinOptions = kotlinSearchOptions)
 
             with(applyQueryFilters(element, options, ReferencesSearch.search(searchParameters))) {

idea/src/org/jetbrains/kotlin/idea/search/ideaExtensions/KotlinConventionMethodReferencesSearcher.kt

@@ -23,11 +23,13 @@ import com.intellij.psi.search.searches.MethodReferencesSearch
 import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.search.restrictToKotlinSources
 import org.jetbrains.kotlin.idea.search.usagesSearch.findDestructuringDeclarationUsages
+import org.jetbrains.kotlin.idea.search.usagesSearch.findInvokeOperatorUsages
 import org.jetbrains.kotlin.idea.search.usagesSearch.getOperationSymbolsToSearch
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.lexer.KtSingleValueToken
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.dataClassUtils.isComponentLike
+import org.jetbrains.kotlin.util.OperatorNameConventions
 
 class KotlinConventionMethodReferencesSearcher() : QueryExecutorBase<PsiReference, MethodReferencesSearch.SearchParameters>(true) {
     override fun processQuery(queryParameters: MethodReferencesSearch.SearchParameters, consumer: Processor<PsiReference>) {
@@ -39,6 +41,9 @@ class KotlinConventionMethodReferencesSearcher() : QueryExecutorBase<PsiReferenc
         if (isComponentLike(identifier)) {
             findDestructuringDeclarationUsages(method, queryParameters.effectiveSearchScope, consumer, queryParameters.optimizer)
         }
+        else if (identifier == OperatorNameConventions.INVOKE) {
+            findInvokeOperatorUsages(method, queryParameters.effectiveSearchScope, consumer)
+        }
         else {
             val operationSymbolsToSearch = identifier.getOperationSymbolsToSearch()
             val wordsToSearch = operationSymbolsToSearch.first.map { (it as KtSingleValueToken).value }

idea/testData/findUsages/kotlin/conventions/invoke.0.kt

@@ -2,10 +2,12 @@
 // OPTIONS: usages
 
 class B(val n: Int) {
-    fun <caret>invoke(i: Int): B = B(i)
+    operator fun <caret>invoke(i: Int){}
 }
 
+fun f() = B(1)
+
 fun test() {
-    B(1).invoke(2)
+    f(1).invoke(2)
-    B(1)(2)
+    f(2)(2)
 }

idea/testData/findUsages/kotlin/conventions/invoke.log

@@ -0,0 +1,8 @@
+Checked type of f()
+Resolving call f(1)
+Resolving call f(2)
+Resolving call f(2)(2)
+Searched references to B
+Searched references to f() in Kotlin files
+Used plain search in LocalSearchScope:
+    CLASS:B

idea/testData/findUsages/kotlin/conventions/invoke.results.txt

@@ -1 +1,2 @@
-Function call 9 B(1).invoke(2)
+Function call 11 f(1).invoke(2)
+Implicit 'invoke' 12 f(2)(2)

idea/testData/refactoring/rename/renameInvoke/after/invoke.kt

@@ -4,5 +4,5 @@ class A(val n: Int) {
 
 fun test() {
     A(1).foo(2)
-    A(1)(2)
+    A(1).foo(2)
 }

idea/testData/search/references/testInvokeFun.kt

@@ -0,0 +1,10 @@
+class B(val n: Int) {
+    operator fun <caret>invoke(i: Int){}
+}
+
+fun f() = B(1)
+
+fun test() {
+    f(1).invoke(2)
+    f(2)(2)
+}

idea/tests/org/jetbrains/kotlin/findUsages/AbstractFindUsagesTest.java

@@ -478,6 +478,16 @@ public abstract class AbstractFindUsagesTest extends KotlinLightCodeInsightFixtu
 
         if (log != null) {
             KotlinTestUtils.assertEqualsToFile(new File(rootPath, prefix + "log"), log);
+
+            // if log is empty then compare results with plain search
+            try {
+                ExpressionsOfTypeProcessor.Companion.setMode(ExpressionsOfTypeProcessor.Mode.ALWAYS_PLAIN);
+
+                findUsagesAndCheckResults(mainFileText, prefix, rootPath, caretElement, options);
+            }
+            finally {
+                ExpressionsOfTypeProcessor.Companion.setMode(ExpressionsOfTypeProcessor.Mode.ALWAYS_SMART);
+            }
         }
     }
 

idea/tests/org/jetbrains/kotlin/search/KotlinReferencesSearchTest.kt

@@ -22,6 +22,7 @@ import com.intellij.psi.search.LocalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.testFramework.fixtures.JavaCodeInsightTestFixture
 import org.jetbrains.kotlin.idea.references.KtDestructuringDeclarationReference
+import org.jetbrains.kotlin.idea.references.KtInvokeFunctionReference
 import org.jetbrains.kotlin.idea.search.usagesSearch.ExpressionsOfTypeProcessor
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
 import org.jetbrains.kotlin.psi.KtFunction
@@ -58,6 +59,13 @@ class KotlinReferencesSearchTest(): AbstractSearcherTest() {
         Assert.assertTrue(refs[1] is KtDestructuringDeclarationReference)
     }
 
+    fun testInvokeFun() {
+        val refs = doTest<KtFunction>()
+        Assert.assertEquals(2, refs.size)
+        Assert.assertEquals("invoke", refs[0].canonicalText)
+        Assert.assertTrue(refs[1] is KtInvokeFunctionReference)
+    }
+
     // workaround for KT-9788 AssertionError from backand when we read field from inline function
     private val myFixtureProxy: JavaCodeInsightTestFixture get() = myFixture