Initial implementation of faster component function usage search

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

@@ -38,6 +38,7 @@ 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.findDestructuringDeclarationUsages
 import org.jetbrains.kotlin.idea.search.usagesSearch.getClassNameForCompanionObject
 import org.jetbrains.kotlin.idea.search.usagesSearch.getSpecialNamesToSearch
 import org.jetbrains.kotlin.idea.stubindex.KotlinSourceFilterScope
@@ -45,6 +46,7 @@ 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
 
 data class KotlinReferencesSearchOptions(val acceptCallableOverrides: Boolean = false,
                                          val acceptOverloads: Boolean = false,
@@ -96,8 +98,8 @@ class KotlinReferencesSearcher : QueryExecutorBase<PsiReference, ReferencesSearc
 
         val resultProcessor = KotlinRequestResultProcessor(unwrappedElement, filter = refFilter, options = kotlinOptions)
 
+        val name = runReadAction { unwrappedElement.name }
         if (kotlinOptions.anyEnabled()) {
-            val name = runReadAction { unwrappedElement.name }
             if (name != null) {
                 queryParameters.optimizer.searchWord(name, effectiveSearchScope, UsageSearchContext.IN_CODE, true, unwrappedElement,
                                                      resultProcessor)
@@ -114,7 +116,32 @@ class KotlinReferencesSearcher : QueryExecutorBase<PsiReference, ReferencesSearc
         }
 
         if (!(unwrappedElement is KtElement && isOnlyKotlinSearch(effectiveSearchScope))) {
-            searchLightElements(queryParameters, element)
+            searchLightElements(queryParameters, element, consumer)
+        }
+
+        if (kotlinOptions.searchForComponentConventions) {
+            when (element) {
+                is KtFunction -> {
+                    if (name != null && isComponentLike(name)) {
+                        findDestructuringDeclarationUsages(element, effectiveSearchScope, consumer, queryParameters.optimizer)
+                    }
+                }
+
+                is KtParameter -> {
+                    val componentFunctionDescriptor = element.dataClassComponentFunction()
+                    if (componentFunctionDescriptor != null) {
+                        val containingClass = element.getStrictParentOfType<KtClassOrObject>()?.toLightClass()
+                        searchDataClassComponentUsages(queryParameters, containingClass, componentFunctionDescriptor, consumer)
+                    }
+                }
+
+                is KtLightParameter -> {
+                    val componentFunctionDescriptor = element.kotlinOrigin?.dataClassComponentFunction()
+                    if (componentFunctionDescriptor != null) {
+                        searchDataClassComponentUsages(queryParameters, element.method.containingClass, componentFunctionDescriptor, consumer)
+                    }
+                }
+            }
         }
     }
 
@@ -209,16 +236,19 @@ class KotlinReferencesSearcher : QueryExecutorBase<PsiReference, ReferencesSearc
 
         private fun searchDataClassComponentUsages(queryParameters: ReferencesSearch.SearchParameters,
                                                    containingClass: PsiClass?,
-                                                   componentFunctionDescriptor: FunctionDescriptor) {
+                                                   componentFunctionDescriptor: FunctionDescriptor,
+                                                   consumer: Processor<PsiReference>
+        ) {
             val componentFunction = containingClass?.methods?.find {
                 it.name == componentFunctionDescriptor.name.asString() && it.parameterList.parametersCount == 0
             }
             if (componentFunction != null) {
                 searchNamedElement(queryParameters, componentFunction)
+                findDestructuringDeclarationUsages(componentFunction, queryParameters.effectiveSearchScope, consumer, queryParameters.optimizer)
             }
         }
 
-        private fun searchLightElements(queryParameters: ReferencesSearch.SearchParameters, element: PsiElement) {
+        private fun searchLightElements(queryParameters: ReferencesSearch.SearchParameters, element: PsiElement, consumer: Processor<PsiReference>) {
             when (element) {
                 is KtClassOrObject -> processKtClassOrObject(element, queryParameters)
                 is KtNamedFunction, is KtSecondaryConstructor -> {
@@ -242,14 +272,6 @@ class KotlinReferencesSearcher : QueryExecutorBase<PsiReference, ReferencesSearc
 
                 is KtParameter -> {
                     searchPropertyMethods(queryParameters, element)
-                    runReadAction {
-
-                        val componentFunctionDescriptor = element.dataClassComponentFunction()
-                        if (componentFunctionDescriptor != null) {
-                            val containingClass = element.getStrictParentOfType<KtClassOrObject>()?.toLightClass()
-                            searchDataClassComponentUsages(queryParameters, containingClass, componentFunctionDescriptor)
-                        }
-                    }
                 }
 
                 is KtLightMethod -> {
@@ -269,12 +291,6 @@ class KotlinReferencesSearcher : QueryExecutorBase<PsiReference, ReferencesSearc
 
                 is KtLightParameter -> {
                     val origin = element.kotlinOrigin ?: return
-                    runReadAction {
-                        val componentFunctionDescriptor = origin.dataClassComponentFunction()
-                        if (componentFunctionDescriptor != null) {
-                            searchDataClassComponentUsages(queryParameters, element.method.containingClass, componentFunctionDescriptor)
-                        }
-                    }
                     searchPropertyMethods(queryParameters, origin)
                 }
             }

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

@@ -17,12 +17,14 @@
 package org.jetbrains.kotlin.idea.search.usagesSearch
 
 import com.google.common.collect.ImmutableSet
-import org.jetbrains.kotlin.idea.references.*
+import org.jetbrains.kotlin.idea.references.KtArrayAccessReference
+import org.jetbrains.kotlin.idea.references.KtForLoopInReference
+import org.jetbrains.kotlin.idea.references.KtPropertyDelegationMethodsReference
+import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DelegatedPropertyResolver
-import org.jetbrains.kotlin.resolve.dataClassUtils.isComponentLike
 import org.jetbrains.kotlin.types.expressions.OperatorConventions.*
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
@@ -57,8 +59,6 @@ fun Name.getOperationSymbolsToSearch(): Pair<Set<KtToken>, Class<*>?> {
         DelegatedPropertyResolver.PROPERTY_DELEGATED_FUNCTION_NAME -> return setOf(KtTokens.BY_KEYWORD) to KtPropertyDelegationMethodsReference::class.java
     }
 
-    if (isComponentLike(this)) return setOf(KtTokens.LPAR) to KtDestructuringDeclarationReference::class.java
-
     val unaryOp = UNARY_OPERATION_NAMES_WITH_DEPRECATED_INVERTED[this]
     if (unaryOp != null) return setOf(unaryOp) to KtSimpleNameReference::class.java
 

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

@@ -0,0 +1,332 @@
+/*
+ * 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.lang.java.JavaLanguage
+import com.intellij.psi.*
+import com.intellij.psi.search.*
+import com.intellij.psi.search.searches.ClassInheritorsSearch
+import com.intellij.psi.search.searches.ReferencesSearch
+import com.intellij.util.Processor
+import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.asJava.elements.KtLightMethod
+import org.jetbrains.kotlin.asJava.namedUnwrappedElement
+import org.jetbrains.kotlin.asJava.toLightClass
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.diagnostics.DiagnosticUtils
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
+import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
+import org.jetbrains.kotlin.idea.references.KtDestructuringDeclarationReference
+import org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinRequestResultProcessor
+import org.jetbrains.kotlin.idea.search.restrictToKotlinSources
+import org.jetbrains.kotlin.idea.util.FuzzyType
+import org.jetbrains.kotlin.idea.util.fuzzyExtensionReceiverType
+import org.jetbrains.kotlin.idea.util.toFuzzyType
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.*
+import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.util.isValidOperator
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
+import java.util.*
+
+//TODO: compiled code
+//TODO: check if it's too expensive
+
+fun findDestructuringDeclarationUsages(
+        componentFunction: PsiMethod,
+        scope: SearchScope,
+        consumer: Processor<PsiReference>,
+        optimizer: SearchRequestCollector
+) {
+    if (componentFunction !is KtLightMethod) return //TODO
+    val ktDeclarationTarget = componentFunction.kotlinOrigin as? KtDeclaration ?: return //TODO?
+    findDestructuringDeclarationUsages(ktDeclarationTarget, scope, consumer, optimizer)
+}
+
+fun findDestructuringDeclarationUsages(
+        ktDeclaration: KtDeclaration,
+        scope: SearchScope,
+        consumer: Processor<PsiReference>,
+        optimizer: SearchRequestCollector
+) {
+    // for local scope it's faster to use plain search
+    if (scope is LocalSearchScope) {
+        val unwrappedElement = ktDeclaration.namedUnwrappedElement ?: return
+        val resultProcessor = KotlinRequestResultProcessor(unwrappedElement,
+                                                           filter = { ref -> ref is KtDestructuringDeclarationReference })
+        optimizer.searchWord("(", scope.restrictToKotlinSources(), UsageSearchContext.IN_CODE, true, unwrappedElement, resultProcessor)
+        return
+    }
+
+    val descriptor = ktDeclaration.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)
+    }
+
+    Processor(dataType, ktDeclaration, scope, consumer).run()
+}
+
+private class Processor(
+        private val dataType: FuzzyType,
+        private val target: KtDeclaration,
+        private val searchScope: SearchScope,
+        private val consumer: Processor<PsiReference>
+) {
+    private val project = target.project
+    private val declarationsToSearch = ArrayList<PsiElement>()
+    private val declarationsToSearchSet = HashSet<PsiElement>()
+
+    fun run() {
+        val dataClassDescriptor = dataType.type.constructor.declarationDescriptor ?: return
+        val dataClassDeclaration = DescriptorToSourceUtilsIde.getAnyDeclaration(project, dataClassDescriptor)
+        val psiClass = when (dataClassDeclaration) {
+            is PsiClass -> dataClassDeclaration
+            is KtClassOrObject -> dataClassDeclaration.toLightClass() ?: return
+            else -> return
+        }
+
+        val parameters = ClassInheritorsSearch.SearchParameters(psiClass, GlobalSearchScope.allScope(project), true, true, false)
+        val classesToSearch = listOf(psiClass) + ClassInheritorsSearch.search(parameters).findAll()
+
+        for (classToSearch in classesToSearch) {
+            ReferencesSearch.search(classToSearch).forEach(Processor { reference ->
+                if (!processDataClassUsage(reference)) {
+                    //TODO
+                    val element = reference.element
+                    val document = PsiDocumentManager.getInstance(project).getDocument(element.containingFile)
+                    val lineAndCol = DiagnosticUtils.offsetToLineAndColumn(document, element.startOffset)
+                    TODO("Unsupported reference: '${element.text}' in ${element.containingFile.name} line ${lineAndCol.line} column ${lineAndCol.column}")
+                }
+                else {
+                    true
+                }
+            })
+        }
+
+        // we use index instead of iterator because elements are added during iteration
+        var index = 0
+        while (index < declarationsToSearch.size) {
+            val declaration = declarationsToSearch[index++]
+            processDeclarationOfTypeWithDataClass(declaration)
+        }
+    }
+
+    //TODO: check if it's operator (too expensive)
+    private fun addDeclarationToProcess(declaration: PsiElement) {
+        if (declarationsToSearchSet.add(declaration)) {
+            declarationsToSearch.add(declaration)
+        }
+    }
+
+    /**
+     * Process usage of our data class or one of its inheritors
+     */
+    private fun processDataClassUsage(reference: PsiReference): Boolean {
+        val element = reference.element
+        return when (element.language) {
+            KotlinLanguage.INSTANCE -> processKotlinDataClassUsage(element)
+
+            JavaLanguage.INSTANCE -> processJavaDataClassUsage(element)
+
+            else -> false // we don't know anything about usages in other languages - so we downgrade to slow algorithm in this case
+        }
+    }
+
+    private fun processKotlinDataClassUsage(element: PsiElement): Boolean {
+        //TODO: type aliases
+        //TODO: doc-comments
+
+        when (element) {
+            is KtReferenceExpression -> {
+                val parent = element.parent
+                when (parent) {
+                    is KtUserType -> {
+                        val typeRef = parent.parents.lastOrNull { it is KtTypeReference }
+                        val typeRefParent = typeRef?.parent
+                        when (typeRefParent) {
+                            is KtCallableDeclaration -> {
+                                when (typeRef) {
+                                    typeRefParent.typeReference -> {
+                                        addDeclarationToProcess(typeRefParent)
+                                        return true
+                                    }
+
+                                    typeRefParent.receiverTypeReference -> {
+                                        //TODO: search usages inside extensions and member functions of our class & derived
+                                        return true
+                                    }
+                                }
+                            }
+
+                            is KtTypeProjection -> {
+                                val callExpression = (typeRefParent.parent as? KtTypeArgumentList)?.parent as? KtCallExpression
+                                if (callExpression != null) {
+                                    processSuspiciousExpression(callExpression)
+                                    return true
+                                }
+                            }
+
+                            is KtConstructorCalleeExpression -> {
+                                if (typeRefParent.parent is KtSuperTypeCallEntry) {
+                                    // usage in super type list - just ignore, inheritors are processed above
+                                    return true
+                                }
+                            }
+                        }
+                    }
+
+                    is KtCallExpression -> {
+                        if (element == parent.calleeExpression) {
+                            processSuspiciousExpression(parent)
+                            return true
+                        }
+                    }
+                }
+
+                if (element.getStrictParentOfType<KtImportDirective>() != null) return true // ignore usage in import
+            }
+        }
+
+        return false // unsupported type of reference
+    }
+
+    private fun processJavaDataClassUsage(element: PsiElement): Boolean {
+        if (element !is PsiJavaCodeReferenceElement) return true // meaningless reference from Java
+
+        ParentsLoop@
+        for (parent in element.parents) {
+            when (parent) {
+                is PsiCodeBlock,
+                is PsiExpression,
+                is PsiImportStatement,
+                is PsiParameter ->
+                    break@ParentsLoop // ignore local usages, usages in imports and in parameter types
+
+                is PsiMethod, is PsiField -> {
+                    if (!(parent as PsiModifierListOwner).isPrivateOrLocal()) {
+                        addDeclarationToProcess(parent)
+                    }
+                    break@ParentsLoop
+                }
+            }
+        }
+
+        return true
+    }
+
+    /**
+     * Process declaration whose type is our data class (or data class used anywhere inside that type)
+     */
+    private fun processDeclarationOfTypeWithDataClass(declaration: PsiElement) {
+        // we don't need to search usages of declarations in Java because Java doesn't have implicitly typed declarations so such usages cannot affect Kotlin code
+        //TODO: what about Scala and other JVM-languages?
+        val declarationUsageScope = GlobalSearchScope.projectScope(declaration.project).restrictToKotlinSources()
+        ReferencesSearch.search(declaration, declarationUsageScope).forEach { reference ->
+            (reference.element as? KtReferenceExpression)?.let { processSuspiciousExpression(it) }
+        }
+    }
+
+    /**
+     * Process expression which may have type of our data class (or data class used anywhere inside that type)
+     */
+    private fun processSuspiciousExpression(expression: KtExpression) {
+        ParentsLoop@
+        for (container in expression.parentsWithSelf) {
+            if (container !is KtExpression) continue
+            val parent = container.parent
+
+            when (parent) {
+                is KtDestructuringDeclaration -> {
+                    processSuspiciousDeclaration(parent)
+                    break@ParentsLoop
+                }
+
+                is KtWithExpressionInitializer -> {
+                    if (container != parent.initializer) continue@ParentsLoop
+
+                    processSuspiciousDeclaration(parent)
+
+                    break@ParentsLoop
+                }
+
+                is KtContainerNode -> {
+                    if (parent.node.elementType == KtNodeTypes.LOOP_RANGE) {
+                        val forExpression = parent.parent as KtForExpression
+                        (forExpression.destructuringParameter ?: forExpression.loopParameter as KtDeclaration?)?.let {
+                            processSuspiciousDeclaration(it)
+                        }
+                    }
+                }
+
+            //TODO: entry of data type
+            }
+        }
+    }
+
+    /**
+     * Process declaration which may have implicit type of our data class (or data class used anywhere inside that type)
+     */
+    private fun processSuspiciousDeclaration(declaration: KtDeclaration) {
+        if (declaration is KtDestructuringDeclaration) {
+            if (searchScope.contains(declaration)) {
+                val declarationReference = declaration.references.firstIsInstance<KtDestructuringDeclarationReference>()
+                if (declarationReference.isReferenceTo(target)) {
+                    consumer.process(declarationReference)
+                }
+            }
+        }
+        else {
+            if (!isImplicitlyTyped(declaration)) return
+
+            val descriptor = declaration.resolveToDescriptorIfAny() as? CallableDescriptor ?: return
+            val type = descriptor.returnType
+            if (type != null && type.containsTypeOrDerivedInside(dataType)) {
+                addDeclarationToProcess(declaration)
+            }
+        }
+    }
+
+
+    private fun PsiModifierListOwner.isPrivateOrLocal(): Boolean {
+        return hasModifierProperty(PsiModifier.PRIVATE) || parents.any { it is PsiCodeBlock }
+    }
+
+    private fun KotlinType.containsTypeOrDerivedInside(type: FuzzyType): Boolean {
+        return type.checkIsSuperTypeOf(this) != null || arguments.any { it.type.containsTypeOrDerivedInside(type) }
+    }
+
+    private fun isImplicitlyTyped(declaration: KtDeclaration): Boolean {
+        return when (declaration) {
+            is KtFunction -> !declaration.hasDeclaredReturnType()
+            is KtVariableDeclaration -> declaration.typeReference == null
+            is KtParameter -> declaration.typeReference == null
+            else -> false
+        }
+    }
+}

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

@@ -22,10 +22,8 @@ import org.jetbrains.kotlin.asJava.LightClassUtil.PropertyAccessorsPsiMethods
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.idea.references.KtDestructuringDeclarationReference
 import org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinReferencesSearchOptions
 import org.jetbrains.kotlin.lexer.KtSingleValueToken
-import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
@@ -69,14 +67,7 @@ fun PsiNamedElement.getSpecialNamesToSearch(options: KotlinReferencesSearchOptio
     val name = name
     return when {
         name == null || !Name.isValidIdentifier(name) -> Collections.emptyList<String>() to null
-        this is KtParameter -> {
-            if (!options.searchForComponentConventions) return Collections.emptyList<String>() to null
 
-            val componentFunctionName = this.dataClassComponentFunction()?.name
-            if (componentFunctionName == null) return Collections.emptyList<String>() to null
-
-            return listOf(componentFunctionName.asString(), KtTokens.LPAR.value) to KtDestructuringDeclarationReference::class.java
-        }
         else -> {
             val operationSymbolsToSearch = Name.identifier(name).getOperationSymbolsToSearch()
             operationSymbolsToSearch.first.map { (it as KtSingleValueToken).value } to operationSymbolsToSearch.second