[AA LL] LLFirProvider: avoid index search if class declaration presents

^KT-55669 Fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProvider.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtClassLikeDeclaration
 
 @ThreadSafeMutableState
 internal class LLFirProvider(
@@ -45,11 +46,20 @@ internal class LLFirProvider(
 
     override val isPhasedFirAllowed: Boolean get() = true
 
-    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? {
+    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? =
-        return SyntheticFirClassProvider.getInstance(session).getFirClassifierByFqName(classId)
+        getFirClassifierByFqNameAndDeclaration(classId, classLikeDeclaration = null)
-            ?: providerHelper.getFirClassifierByFqName(classId)
+
+    fun getFirClassifierByDeclaration(classLikeDeclaration: KtClassLikeDeclaration): FirClassLikeDeclaration? {
+        val classId = classLikeDeclaration.getClassId() ?: return null
+        return getFirClassifierByFqNameAndDeclaration(classId, classLikeDeclaration)
     }
 
+    private fun getFirClassifierByFqNameAndDeclaration(
+        classId: ClassId,
+        classLikeDeclaration: KtClassLikeDeclaration?,
+    ): FirClassLikeDeclaration? = SyntheticFirClassProvider.getInstance(session).getFirClassifierByFqName(classId)
+        ?: providerHelper.getFirClassifierByFqNameAndDeclaration(classId, classLikeDeclaration)
+
     override fun getFirClassifierContainerFile(fqName: ClassId): FirFile {
         return getFirClassifierContainerFileIfAny(fqName)
             ?: error("Couldn't find container for $fqName")

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtClassLikeDeclaration
 
 internal class LLFirProviderHelper(
     firSession: FirSession,
@@ -37,14 +38,15 @@ internal class LLFirProviderHelper(
     private val allowKotlinPackage = canContainKotlinPackage ||
             firSession.languageVersionSettings.getFlag(AnalysisFlags.allowKotlinPackage)
 
-    private val classifierByClassId = firSession.firCachesFactory.createCache<ClassId, FirClassLikeDeclaration?> { classId ->
+    private val classifierByClassId =
-        val ktClass = declarationProvider.getClassLikeDeclarationByClassId(classId)
+        firSession.firCachesFactory.createCache<ClassId, FirClassLikeDeclaration?, KtClassLikeDeclaration?> { classId, context ->
-            ?: return@createCache null
+            val ktClass = context ?: declarationProvider.getClassLikeDeclarationByClassId(classId) ?: return@createCache null
-        if (ktClass.getClassId() == null) return@createCache null
+
-        val firFile = firFileBuilder.buildRawFirFileWithCaching(ktClass.containingKtFile)
+            if (ktClass.getClassId() == null) return@createCache null
-        FirElementFinder.findClassifierWithClassId(firFile, classId)
+            val firFile = firFileBuilder.buildRawFirFileWithCaching(ktClass.containingKtFile)
-            ?: error("Classifier $classId was found in file ${ktClass.containingKtFile.virtualFilePath} but was not found in FirFile")
+            FirElementFinder.findClassifierWithClassId(firFile, classId)
-    }
+                ?: error("Classifier $classId was found in file ${ktClass.containingKtFile.virtualFilePath} but was not found in FirFile")
+        }
 
 
     private val callablesByCallableId = firSession.firCachesFactory.createCache<CallableId, List<FirCallableSymbol<*>>> { callableId ->
@@ -57,10 +59,13 @@ internal class LLFirProviderHelper(
         }
     }
 
-    fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? {
+    fun getFirClassifierByFqNameAndDeclaration(
+        classId: ClassId,
+        classLikeDeclaration: KtClassLikeDeclaration?,
+    ): FirClassLikeDeclaration? {
         if (classId.isLocal) return null
         if (!allowKotlinPackage && classId.isKotlinPackage()) return null
-        return classifierByClassId.getValue(classId)
+        return classifierByClassId.getValue(classId, classLikeDeclaration)
     }
 
     fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
+import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirProvider
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
@@ -117,8 +118,14 @@ var KtFile.originalKtFile by UserDataProperty(ORIGINAL_KT_FILE_KEY)
 
 
 private fun KtClassLikeDeclaration.findFir(provider: FirProvider): FirClassLikeDeclaration? {
-    val classId = getClassId() ?: return null
+    val declaration = if (provider is LLFirProvider) {
-    return provider.getFirClassifierByFqName(classId) as? FirRegularClass
+        provider.getFirClassifierByDeclaration(this)
+    } else {
+        val classId = getClassId() ?: return null
+        provider.getFirClassifierByFqName(classId)
+    }
+
+    return declaration as? FirRegularClass
 }