FIR IDE: optimize search for FirDeclaration by KtDeclaration

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt

@@ -8,11 +8,13 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.FirElementFinder
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
@@ -36,10 +38,7 @@ internal class FirElementBuilder(
         val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
 
         val containerFir = when (val container = element.getNonLocalContainingDeclarationWithFqName()) {
-            is KtDeclaration -> {
-                FirElementFinder.findElementByPsiIn<FirDeclaration>(firFile, container)
-                    ?: error("Declaration was not found in FIR file which was build by declaration KtFile, declaration is\n${container.getElementTextInContext()}")
-            }
+            is KtDeclaration -> container.findNonLocalFirDeclaration(firFileBuilder, firFile.session.firProvider, moduleFileCache)
             null -> firFile
             else -> error("Unsupported: ${container.text}")
         }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -19,11 +19,9 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
-import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
 internal class FirLazyDeclarationResolver(
     private val firFileBuilder: FirFileBuilder
@@ -139,31 +137,9 @@ internal class FirLazyDeclarationResolver(
         if (this is FirFile) return this
         val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration")
         if (!KtPsiUtil.isLocal(ktDeclaration)) return this
-        return when (val nonLocalPsi = ktDeclaration.getNonLocalContainingDeclarationWithFqName()) {
-            is KtClassOrObject -> provider.getFirClassifierByFqName(
-                nonLocalPsi.classIdIfNonLocal()
-                    ?: error("Container classId should not be null for non-local declaration")
-            ) ?: error("Could not find class ${nonLocalPsi.classIdIfNonLocal()}")
-            is KtProperty, is KtNamedFunction -> {
-                val containerClass = nonLocalPsi.containingClassOrObject
-                if (containerClass != null) {
-                    val containerClassId = containerClass.classIdIfNonLocal()
-                        ?: error("Container classId should not be null for non-local declaration")
-                    val containingFir = provider.getFirClassifierByFqName(containerClassId) as? FirRegularClass
-                        ?: error("Could not find class $containerClassId")
-                    containingFir.declarations.first { it.psi === nonLocalPsi }
-                } else {
-                    val packageFqName = ktDeclaration.containingKtFile.packageFqName
-                    provider.symbolProvider.getTopLevelCallableSymbols(packageFqName, nonLocalPsi.nameAsSafeName)
-                        .firstOrNull { it.fir.psi === nonLocalPsi }
-                    val ktFile = ktDeclaration.containingKtFile
-                    val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
-                    firFile.declarations.firstOrNull { it.psi === nonLocalPsi }
-                        ?: error("Cannot find FIR for\n${nonLocalPsi.getElementTextInContext()}")
-                }
-            }
-            else -> error("Invalid container ${nonLocalPsi?.let { it::class } ?: "null"}")
-        }
+        val nonLocalPsi = ktDeclaration.getNonLocalContainingDeclarationWithFqName()
+            ?: error("Container for local declaration cannot be null")
+        return nonLocalPsi.findNonLocalFirDeclaration(firFileBuilder, provider, moduleFileCache)
     }
 }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/declarationUtils.kt

@@ -0,0 +1,51 @@
+/*
+ * 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.fir.low.level.api.util
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+
+internal fun KtDeclaration.findNonLocalFirDeclaration(
+    firFileBuilder: FirFileBuilder,
+    provider: FirProvider,
+    moduleFileCache: ModuleFileCache
+): FirDeclaration {
+    require(!KtPsiUtil.isLocal(this))
+    return when {
+        this is KtClassOrObject -> findFir(provider)
+        this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> {
+            val containerClass = containingClassOrObject
+            if (containerClass != null) {
+                val containerClassFir = containerClass.findFir(provider)
+                containerClassFir.declarations.first { it.psi === this }
+            } else {
+                val ktFile = containingKtFile
+                val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
+                firFile.declarations.first { it.psi === this }
+            }
+        }
+        this is KtConstructor<*> -> {
+            val containerClassFir = containingClassOrObject?.findFir(provider)
+                ?: error("Container class should be not null for KtConstructor")
+            containerClassFir.declarations.first { it.psi === this }
+        }
+        else -> error("Invalid container $this")
+    }
+}
+
+private fun KtClassOrObject.findFir(provider: FirProvider): FirRegularClass {
+    val containerClassId = classIdIfNonLocal()
+        ?: error("Container classId should not be null for non-local declaration")
+    return provider.getFirClassifierByFqName(containerClassId) as? FirRegularClass
+        ?: error("Could not find class $containerClassId")
+}