From cec299ac0ac287d95c76da79b3507fbab2f39000 Mon Sep 17 00:00:00 2001 From: Jinseong Jeon Date: Fri, 8 Sep 2023 00:24:19 -0700 Subject: [PATCH] K2 UAST: Improve performance of PSI declaration provider When it comes to finding regular class (of `PsiClass`) from the given `ClassId`, we've reinvented the wheel in a worse way. `JavaFileManager` is there to find `PsiClass`, and we already build/populate index for that while initializing AA standalone mode's own project environment. We can simply load that manager and send a query, that's it. However, our efforts so far on binary PSI declaration providers are not entirely useless. `JavaFileManager` can't handle top-level declarations and search for all callables/classes in a given package in general. So, we can still use existing logic for non-ClassId queries. We just need to cache the intermediate results: fq name to `VirtualFile`s. To fully utilize the aforementioned cache, we also need to cache an instance of `KotlinPsiDeclarationProvider`. Combing these three ideas, now running K2 UAST on a very large module is "on par", e.g., K1: //tools/adt/idea/android:intellij.android.core_lint_test PASSED in 112.8s Stats over 7 runs: max = 112.8s, min = 104.3s, avg = 107.9s, dev = 3.6s K2 (before): //tools/adt/idea/android:intellij.android.core_lint_test PASSED in 303.0s Stats over 7 runs: max = 303.0s, min = 286.8s, avg = 294.4s, dev = 5.0s K2 (after): //tools/adt/idea/android:intellij.android.core_lint_test PASSED in 112.8s Stats over 7 runs: max = 112.8s, min = 105.7s, avg = 108.9s, dev = 2.4s --- ...cPsiDeclarationFromBinaryModuleProvider.kt | 53 ++++++++++++++----- 1 file changed, 40 insertions(+), 13 deletions(-) diff --git a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticPsiDeclarationFromBinaryModuleProvider.kt b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticPsiDeclarationFromBinaryModuleProvider.kt index 2c26dc5b49a..deac927790e 100644 --- a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticPsiDeclarationFromBinaryModuleProvider.kt +++ b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticPsiDeclarationFromBinaryModuleProvider.kt @@ -11,7 +11,9 @@ import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem import com.intellij.psi.* import com.intellij.psi.impl.compiled.ClsClassImpl import com.intellij.psi.impl.compiled.ClsFileImpl +import com.intellij.psi.impl.file.impl.JavaFileManager import com.intellij.psi.search.GlobalSearchScope +import com.intellij.util.containers.ContainerUtil import org.jetbrains.kotlin.analysis.decompiled.light.classes.ClsJavaStubByVirtualFileCache import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule import org.jetbrains.kotlin.analysis.providers.KotlinPsiDeclarationProvider @@ -24,6 +26,8 @@ import org.jetbrains.kotlin.name.CallableId import org.jetbrains.kotlin.name.ClassId import org.jetbrains.kotlin.name.FqName import org.jetbrains.kotlin.util.capitalizeDecapitalize.decapitalizeSmart +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.ConcurrentMap private class KotlinStaticPsiDeclarationFromBinaryModuleProvider( private val project: Project, @@ -34,17 +38,23 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider( ) : KotlinPsiDeclarationProvider(), AbstractDeclarationFromBinaryModuleProvider { private val psiManager by lazyPub { PsiManager.getInstance(project) } - private fun clsClassImplsByFqName( + private val javaFileManager by lazyPub { project.getService(JavaFileManager::class.java) } + + private val virtualFileCache = ContainerUtil.createConcurrentSoftMap>>() + + private fun clsClassImplsInPackage( fqName: FqName, - isPackageName: Boolean = true, ): Collection { return binaryModules - .flatMap { - val virtualFilesFromKotlinModule = if (isPackageName) virtualFilesFromKotlinModule(it, fqName) else emptySet() - // NB: this assumes Kotlin module has a valid `kotlin_module` info, - // i.e., package part info for the given `fqName` points to exact class paths we're looking for, - // and thus it's redundant to walk through the folders in an exhaustive way. - virtualFilesFromKotlinModule.ifEmpty { virtualFilesFromModule(it, fqName, isPackageName) } + .flatMap { binaryModule -> + val mapPerModule = virtualFileCache.getOrPut(binaryModule) { ConcurrentHashMap() } + mapPerModule.getOrPut(fqName) { + val virtualFilesFromKotlinModule = virtualFilesFromKotlinModule(binaryModule, fqName) + // NB: this assumes Kotlin module has a valid `kotlin_module` info, + // i.e., package part info for the given `fqName` points to exact class paths we're looking for, + // and thus it's redundant to walk through the folders in an exhaustive way. + virtualFilesFromKotlinModule.ifEmpty { virtualFilesFromModule(binaryModule, fqName, isPackageName = true) } + } } .mapNotNull { createClsJavaClassFromVirtualFile(it) @@ -72,14 +82,14 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider( parentClsClass.innerClasses.find { it.name == innerClassName } } } - return clsClassImplsByFqName(classId.asSingleFqName(), isPackageName = false) + return listOfNotNull(javaFileManager.findClass(classId.asFqNameString(), scope)) } // TODO(dimonchik0036): support 'is' accessor override fun getProperties(callableId: CallableId): Collection { val classes = callableId.classId?.let { classId -> getClassesByClassId(classId) - } ?: clsClassImplsByFqName(callableId.packageName) + } ?: clsClassImplsInPackage(callableId.packageName) return classes.flatMap { psiClass -> psiClass.children .filterIsInstance() @@ -100,7 +110,7 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider( override fun getFunctions(callableId: CallableId): Collection { val classes = callableId.classId?.let { classId -> getClassesByClassId(classId) - } ?: clsClassImplsByFqName(callableId.packageName) + } ?: clsClassImplsInPackage(callableId.packageName) return classes.flatMap { psiClass -> psiClass.methods.filter { psiMethod -> psiMethod.name == callableId.callableName.identifier @@ -116,8 +126,11 @@ class KotlinStaticPsiDeclarationProviderFactory( private val binaryModules: Collection, private val jarFileSystem: CoreJarFileSystem, ) : KotlinPsiDeclarationProviderFactory() { - override fun createPsiDeclarationProvider(searchScope: GlobalSearchScope): KotlinPsiDeclarationProvider { - return KotlinStaticPsiDeclarationFromBinaryModuleProvider( + // TODO: For now, [createPsiDeclarationProvider] is always called with the project scope, hence singleton. + // If we come up with a better / optimal search scope, we may need a different way to cache scope-to-provider mapping. + private val provider: KotlinStaticPsiDeclarationFromBinaryModuleProvider by lazy { + val searchScope = GlobalSearchScope.allScope(project) + KotlinStaticPsiDeclarationFromBinaryModuleProvider( project, searchScope, project.createPackagePartProvider(searchScope), @@ -125,4 +138,18 @@ class KotlinStaticPsiDeclarationProviderFactory( jarFileSystem, ) } + + override fun createPsiDeclarationProvider(searchScope: GlobalSearchScope): KotlinPsiDeclarationProvider { + return if (searchScope == provider.scope) { + provider + } else { + KotlinStaticPsiDeclarationFromBinaryModuleProvider( + project, + searchScope, + project.createPackagePartProvider(searchScope), + binaryModules, + jarFileSystem, + ) + } + } }