Optimization for faster search of references to top-level members

2020-04-21 22:52:13 +03:00
parent ec057a6c00
commit 0c3c8b64ce
3 changed files with 99 additions and 0 deletions
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.search.ideaExtensions
+
+import com.intellij.ide.highlighter.JavaFileType
+import com.intellij.psi.PsiElement
+import com.intellij.psi.impl.cache.CacheManager
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.ScopeOptimizer
+import com.intellij.psi.search.SearchScope
+import com.intellij.psi.search.UsageSearchContext
+import org.jetbrains.kotlin.fileClasses.javaFileFacadeFqName
+import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.idea.search.excludeFileTypes
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
+import org.jetbrains.kotlin.psi.KtFile
+
+class KotlinReferenceScopeOptimizer : ScopeOptimizer {
+    override fun getRestrictedUseScope(element: PsiElement): SearchScope? {
+        if (element is KtCallableDeclaration && element.parent is KtFile) {
+            return getRestrictedScopeForTopLevelCallable(element)
+        }
+        return null
+    }
+
+    private fun getRestrictedScopeForTopLevelCallable(callable: KtCallableDeclaration): GlobalSearchScope? {
+        val useScope = callable.useScope as? GlobalSearchScope ?: return null
+        val file = callable.parent as KtFile
+        val packageName = file.packageFqName.takeUnless { it.isRoot } ?: return null
+        val project = file.project
+        val cacheManager = CacheManager.SERVICE.getInstance(project)
+
+        val kotlinScope = GlobalSearchScope.getScopeRestrictedByFileTypes(useScope, KotlinFileType.INSTANCE)
+        val javaScope = GlobalSearchScope.getScopeRestrictedByFileTypes(useScope, JavaFileType.INSTANCE)
+        val restScope = useScope.excludeFileTypes(KotlinFileType.INSTANCE, JavaFileType.INSTANCE) as GlobalSearchScope
+
+        //TODO: use all components of package name?
+        val shortPackageName = packageName.shortName().identifier
+        val kotlinFiles = cacheManager.getVirtualFilesWithWord(shortPackageName, UsageSearchContext.IN_CODE, kotlinScope, true)
+
+        val javaFacadeName = file.javaFileFacadeFqName.shortName().identifier
+        val javaFiles = cacheManager.getVirtualFilesWithWord(javaFacadeName, UsageSearchContext.IN_CODE, javaScope, true)
+
+        return GlobalSearchScope.filesScope(project, (kotlinFiles + javaFiles).asList()).uniteWith(restScope)
+    }
+}
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.search.ideaExtensions
+
+import com.intellij.ide.highlighter.JavaFileType
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
+import com.intellij.psi.search.*
+import org.jetbrains.kotlin.fileClasses.javaFileFacadeFqName
+import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.idea.search.excludeFileTypes
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
+import org.jetbrains.kotlin.psi.KtFile
+
+class KotlinReferenceScopeOptimizer : ScopeOptimizer {
+    override fun getRestrictedUseScope(element: PsiElement): SearchScope? {
+        if (element is KtCallableDeclaration && element.parent is KtFile) {
+            return getRestrictedScopeForTopLevelCallable(element)
+        }
+        return null
+    }
+
+    private fun getRestrictedScopeForTopLevelCallable(callable: KtCallableDeclaration): GlobalSearchScope? {
+        val useScope = callable.useScope as? GlobalSearchScope ?: return null
+        val file = callable.parent as KtFile
+        val packageName = file.packageFqName.takeUnless { it.isRoot } ?: return null
+        val project = file.project
+        val searchHelper = PsiSearchHelper.getInstance(project)
+
+        val kotlinScope = GlobalSearchScope.getScopeRestrictedByFileTypes(useScope, KotlinFileType.INSTANCE)
+        val javaScope = GlobalSearchScope.getScopeRestrictedByFileTypes(useScope, JavaFileType.INSTANCE)
+        val restScope = useScope.excludeFileTypes(KotlinFileType.INSTANCE, JavaFileType.INSTANCE) as GlobalSearchScope
+
+        val kotlinFiles = mutableListOf<VirtualFile>()
+        searchHelper.processCandidateFilesForText(kotlinScope, UsageSearchContext.IN_CODE, true, packageName.asString()) {
+            kotlinFiles.add(it)
+        }
+
+        val javaFiles = mutableListOf<VirtualFile>()
+        searchHelper.processCandidateFilesForText(javaScope, UsageSearchContext.IN_CODE, true, file.javaFileFacadeFqName.asString()) {
+            javaFiles.add(it)
+        }
+
+        return GlobalSearchScope.filesScope(project, kotlinFiles + javaFiles).uniteWith(restScope)
+    }
+}
@@ -704,6 +704,7 @@
    <gotoDeclarationHandler implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinImportAliasGotoDeclarationHandler"/>
    <referencesSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinReferencesSearcher"/>
    <referencesSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinAliasedImportedElementSearcher"/>
+    <codeUsageScopeOptimizer implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinReferenceScopeOptimizer"/>
    <directClassInheritorsSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinDirectInheritorsSearcher"/>
    <overridingMethodsSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinOverridingMethodsWithGenericsSearcher"/>
    <definitionsScopedSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinDefinitionsSearcher"/>