[LL FIR] optimize search for fir declaration by kt declaration

Introduce new API to traverse FirFile by known declaration path
instead of unoptimized findSourceByTraversingWholeTree

^KT-59454

analysis/low-level-api-fir/build.gradle.kts

@@ -76,6 +76,7 @@ allprojects {
         kotlinOptions {
             freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.fir.symbols.SymbolInternals"
             freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals"
+            freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals"
         }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt

@@ -10,80 +10,104 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.packageFqName
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
 object FirElementFinder {
     fun findClassifierWithClassId(
         firFile: FirFile,
         classId: ClassId,
-    ): FirClassLikeDeclaration? = findPathToClassifierWithClassId(
+    ): FirClassLikeDeclaration? = collectDesignationPath(
         firFile = firFile,
-        classId = classId,
-        expectedMemberDeclaration = null
-    )?.second
+        containerClassId = classId.outerClassId,
+        expectedDeclarationAcceptor = { it is FirClassLikeDeclaration && it.symbol.name == classId.shortClassName },
+    )?.target?.let { it as FirClassLikeDeclaration }
 
     fun findClassPathToDeclaration(
         firFile: FirFile,
         declarationContainerClassId: ClassId,
         targetMemberDeclaration: FirDeclaration,
-    ): List<FirRegularClass>? = findPathToClassifierWithClassId(
+    ): List<FirRegularClass>? = collectDesignationPath(
         firFile = firFile,
-        classId = declarationContainerClassId,
-        expectedMemberDeclaration = targetMemberDeclaration
-    )?.first
+        containerClassId = declarationContainerClassId,
+        expectedDeclarationAcceptor = { it == targetMemberDeclaration },
+    )?.path
 
-    private fun findPathToClassifierWithClassId(
+    fun findDeclaration(firFile: FirFile, nonLocalDeclaration: KtDeclaration): FirDeclaration? = collectDesignationPath(
+        firFile = firFile,
+        nonLocalDeclaration = nonLocalDeclaration,
+    )?.target
+
+    fun findPathToDeclarationWithTarget(
         firFile: FirFile,
-        classId: ClassId,
-        expectedMemberDeclaration: FirDeclaration?,
-    ): Pair<List<FirRegularClass>, FirClassLikeDeclaration>? {
-        requireWithAttachmentBuilder(!classId.isLocal, { "ClassId should not be local" }) {
-            withEntry("classId", classId) { it.asString() }
-        }
-        requireWithAttachmentBuilder(
-            firFile.packageFqName == classId.packageFqName,
-            { "ClassId should not be local" }
-        ) {
-            withEntry("FirFile.packageName", firFile.packageFqName) { it.asString() }
-            withEntry("ClassId.packageName", classId.packageFqName) { it.asString() }
+        nonLocalDeclaration: KtDeclaration,
+    ): List<FirDeclaration>? = collectDesignationPath(
+        firFile = firFile,
+        nonLocalDeclaration = nonLocalDeclaration,
+    )?.pathWithTarget
+
+    private class FirDeclarationDesignation(
+        val path: List<FirRegularClass>,
+        val target: FirDeclaration,
+    ) {
+        val pathWithTarget: List<FirDeclaration> get() = path + target
+    }
+
+    private fun collectDesignationPath(
+        firFile: FirFile,
+        nonLocalDeclaration: KtDeclaration,
+    ): FirDeclarationDesignation? = collectDesignationPath(
+        firFile = firFile,
+        containerClassId = nonLocalDeclaration.containingClassOrObject?.getClassId(),
+        expectedDeclarationAcceptor = { it.psi == nonLocalDeclaration },
+    )
+
+    private fun collectDesignationPath(
+        firFile: FirFile,
+        containerClassId: ClassId?,
+        expectedDeclarationAcceptor: (FirDeclaration) -> Boolean,
+    ): FirDeclarationDesignation? {
+        if (containerClassId != null) {
+            requireWithAttachmentBuilder(!containerClassId.isLocal, { "ClassId should not be local" }) {
+                withEntry("classId", containerClassId) { it.asString() }
+            }
+
+            requireWithAttachmentBuilder(
+                firFile.packageFqName == containerClassId.packageFqName,
+                { "ClassId should not be local" }
+            ) {
+                withEntry("FirFile.packageName", firFile.packageFqName) { it.asString() }
+                withEntry("ClassId.packageName", containerClassId.packageFqName) { it.asString() }
+            }
         }
 
-        val classIdPathSegment = classId.relativeClassName.pathSegments()
+        val classIdPathSegment = containerClassId?.relativeClassName?.pathSegments().orEmpty()
         val path = ArrayList<FirRegularClass>(classIdPathSegment.size)
-        var result: FirClassLikeDeclaration? = null
+        var result: FirDeclaration? = null
 
         fun find(declarations: Iterable<FirDeclaration>, classIdPathIndex: Int): Boolean {
-            val currentClassSegment = classIdPathSegment[classIdPathIndex]
-
+            val currentClassSegment = classIdPathSegment.getOrNull(classIdPathIndex)
             for (subDeclaration in declarations) {
-                if (subDeclaration is FirScript) {
-                    val scriptDeclarations = subDeclaration.statements.asSequence().filterIsInstance<FirDeclaration>()
-                    if (find(scriptDeclarations.asIterable(), classIdPathIndex)) {
+                when {
+                    currentClassSegment == null && expectedDeclarationAcceptor(subDeclaration) -> {
+                        result = subDeclaration
                         return true
                     }
 
-                    continue
-                }
-
-                if (subDeclaration is FirClassLikeDeclaration && currentClassSegment == subDeclaration.symbol.name) {
-                    if (classIdPathIndex == classIdPathSegment.lastIndex) {
-                        if (expectedMemberDeclaration == null ||
-                            subDeclaration is FirRegularClass && expectedMemberDeclaration in subDeclaration.declarations
-                        ) {
-                            if (subDeclaration is FirRegularClass) {
-                                path += subDeclaration
-                            }
-
-                            result = subDeclaration
+                    subDeclaration is FirScript -> {
+                        val scriptDeclarations = subDeclaration.statements.asSequence().filterIsInstance<FirDeclaration>()
+                        if (find(scriptDeclarations.asIterable(), classIdPathIndex)) {
                             return true
                         }
 
                         continue
                     }
 
-                    if (subDeclaration is FirRegularClass) {
+                    subDeclaration is FirRegularClass && currentClassSegment == subDeclaration.symbol.name -> {
                         path += subDeclaration
                         if (find(subDeclaration.declarations, classIdPathIndex + 1)) {
                             return true
@@ -99,7 +123,10 @@ object FirElementFinder {
 
         find(firFile.declarations, classIdPathIndex = 0)
         return result?.let {
-            (path as List<FirRegularClass>) to it
+            FirDeclarationDesignation(
+                path = if (path.isEmpty()) emptyList() else path,
+                target = it,
+            )
         }
     }
 

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

@@ -5,7 +5,10 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.canBePartOfParentDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.containingDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
@@ -29,39 +32,59 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration(
     provider: FirProvider,
     containerFirFile: FirFile? = null,
 ): FirDeclaration {
-    //TODO test what way faster
-    findSourceNonLocalFirDeclarationByProvider(firFileBuilder, provider, containerFirFile)?.let { return it }
-    findSourceByTraversingWholeTree(firFileBuilder, containerFirFile)?.let { return it }
+    val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile)
+
+    // TODO test what way faster
+    if (isPhysical) {
+        // do not request providers with non-physical psi in order not to leak them there and
+        // to avoid inconsistency between physical psi and its copy during completion
+        findSourceNonLocalFirDeclarationByProvider(
+            firDeclarationProvider = { declaration ->
+                if (declaration is KtClassLikeDeclaration) {
+                    declaration.findFir(provider)
+                } else {
+                    val containingClassOrObject = declaration.containingClassOrObject
+                    val declarations = if (containingClassOrObject != null) {
+                        val containerClassFir = containingClassOrObject.findFir(provider) as? FirRegularClass
+                        containerClassFir?.declarations
+                    } else {
+                        if (declaration.containingKtFile.isScript()) {
+                            // .kts will have a single [FirScript] as a declaration. We need to unwrap statements in it.
+                            (firFile.declarations.singleOrNull() as? FirScript)?.statements?.filterIsInstance<FirDeclaration>()
+                        } else {
+                            firFile.declarations
+                        }
+                    }
+
+                    // It is possible that we will not be able to find the needed declaration here when the code is invalid
+                    // e.g., we have two conflicting declarations with the same name,
+                    // and we are searching for the wrong one
+                    declarations?.find { it.psi == declaration }
+                }
+            },
+        )?.let { return it }
+    }
+
+    findSourceNonLocalFirDeclarationByProvider(
+        firDeclarationProvider = { declaration ->
+            FirElementFinder.findDeclaration(firFile, declaration)
+        },
+    )?.let { return it }
+
     errorWithFirSpecificEntries(
         "No fir element was found for",
         psi = this,
-        fir = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile),
+        fir = firFile,
         additionalInfos = { withEntry("isPhysical", isPhysical.toString()) }
     )
 }
 
-internal fun KtDeclaration.findFirDeclarationForAnyFirSourceDeclaration(
-    firFileBuilder: LLFirFileBuilder,
-    provider: FirProvider,
-): FirDeclaration {
-    val nonLocalDeclaration = getNonLocalContainingOrThisDeclaration()
-        ?.findSourceNonLocalFirDeclaration(firFileBuilder, provider)
-        ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile)
-    val originalDeclaration = originalDeclaration
-    val fir = FirElementFinder.findElementIn<FirDeclaration>(nonLocalDeclaration) { firDeclaration ->
-        firDeclaration.psi == this || firDeclaration.psi == originalDeclaration
-    }
-    return fir
-        ?: errorWithFirSpecificEntries("FirDeclaration was not found", psi = this)
-}
-
-internal inline fun <reified F : FirDeclaration> KtDeclaration.findFirDeclarationForAnyFirSourceDeclarationOfType(
-    firFileBuilder: LLFirFileBuilder,
-    provider: FirProvider,
-): FirDeclaration {
-    val fir = findFirDeclarationForAnyFirSourceDeclaration(firFileBuilder, provider)
-    if (fir !is F) throwUnexpectedFirElementError(fir, this, F::class)
-    return fir
+@KtAnalysisApiInternals
+fun PsiElement.collectContainingDeclarationsIfNonLocal(session: LLFirResolveSession): List<FirDeclaration>? {
+    val ktFile = containingFile as? KtFile ?: return null
+    val ktDeclaration = getNonLocalContainingOrThisDeclaration { !it.canBePartOfParentDeclaration } ?: return null
+    val firFile = session.getOrBuildFirFile(ktFile)
+    return FirElementFinder.findPathToDeclarationWithTarget(firFile, ktDeclaration)
 }
 
 internal fun KtElement.findSourceByTraversingWholeTree(
@@ -81,56 +104,22 @@ internal fun KtElement.findSourceByTraversingWholeTree(
 }
 
 private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
-    firFileBuilder: LLFirFileBuilder,
-    provider: FirProvider,
-    containerFirFile: FirFile?,
+    firDeclarationProvider: (KtDeclaration) -> FirDeclaration?,
 ): FirDeclaration? {
-    if (!isPhysical) {
-        //do not request providers with non-physical psi in order not to leak them there and 
-        //to avoid inconsistency between physical psi and its copy during completion
-        return null
-    }
-    val candidate = when {
-        this is KtClassOrObject -> findFir(provider)
-        this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> {
-            val containerClass = containingClassOrObject
-            val declarations = if (containerClass != null) {
-                val containerClassFir = containerClass.findFir(provider) as? FirRegularClass
-                containerClassFir?.declarations
-            } else {
-                val ktFile = containingKtFile
-                val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(ktFile)
-                if (ktFile.isScript()) {
-                    // .kts will have a single [FirScript] as a declaration. We need to unwrap statements in it.
-                    (firFile.declarations.singleOrNull() as? FirScript)?.statements?.filterIsInstance<FirDeclaration>()
-                } else {
-                    firFile.declarations
-                }
-            }
+    val candidate = when (this) {
+        is KtClassOrObject,
+        is KtProperty,
+        is KtNamedFunction,
+        is KtConstructor<*>,
+        is KtClassInitializer,
+        is KtTypeAlias,
+        is KtDestructuringDeclaration,
+        is KtScript,
+        -> firDeclarationProvider(this)
 
-            /*
-            It is possible that we will not be able to find needed declaration here when the code is invalid,
-            e.g, we have two conflicting declarations with the same name and we are searching in the wrong one
-             */
-            declarations?.firstOrNull { it.psi == this }
-        }
-        this is KtConstructor<*> || this is KtClassInitializer -> {
-            val containingClass = containingClassOrObject
-                ?: errorWithFirSpecificEntries("Container class should be not null for KtConstructor", psi = this)
-            val containerClassFir = containingClass.findFir(provider) as? FirRegularClass ?: return null
-            containerClassFir.declarations.firstOrNull { it.psi == this }
-        }
-        this is KtTypeAlias -> findFir(provider)
-        this is KtDestructuringDeclaration -> {
-            val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile)
-            firFile.declarations.firstOrNull { it.psi == this }
-        }
-        this is KtScript -> containerFirFile?.declarations?.singleOrNull { it is FirScript }
-        this is KtPropertyAccessor -> {
+        is KtPropertyAccessor -> {
             val firPropertyDeclaration = property.findSourceNonLocalFirDeclarationByProvider(
-                firFileBuilder,
-                provider,
-                containerFirFile,
+                firDeclarationProvider,
             ) as? FirVariable ?: return null
 
             if (isGetter) {
@@ -139,32 +128,36 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
                 firPropertyDeclaration.setter
             }
         }
-        this is KtParameter -> {
-            val ownerFunction = ownerFunction
-                ?: errorWithFirSpecificEntries("Containing function should be not null for KtParameter", psi = this)
+
+        is KtParameter -> {
+            val ownerFunction = ownerFunction ?: errorWithFirSpecificEntries(
+                "Containing function should be not null for KtParameter",
+                psi = this,
+            )
 
             val firFunctionDeclaration = ownerFunction.findSourceNonLocalFirDeclarationByProvider(
-                firFileBuilder,
-                provider,
-                containerFirFile,
+                firDeclarationProvider,
             ) as? FirFunction ?: return null
 
             firFunctionDeclaration.valueParameters[parameterIndex()]
         }
-        this is KtTypeParameter -> {
-            val declaration = containingDeclaration
-                ?: errorWithFirSpecificEntries("Containing declaration should be not null for KtTypeParameter", psi = this)
+
+        is KtTypeParameter -> {
+            val declaration = containingDeclaration ?: errorWithFirSpecificEntries(
+                "Containing declaration should be not null for KtTypeParameter",
+                psi = this,
+            )
 
             val firTypeParameterOwner = declaration.findSourceNonLocalFirDeclarationByProvider(
-                firFileBuilder,
-                provider,
-                containerFirFile,
+                firDeclarationProvider,
             ) as? FirTypeParameterRefsOwner ?: return null
 
             firTypeParameterOwner.typeParameters.firstOrNull { it.psi == this } as FirDeclaration
         }
+
         else -> errorWithFirSpecificEntries("Invalid container", psi = this)
     }
+
     return candidate?.takeIf { it.realPsi == this }
 }