[LL FIR] StubBasedClassDeserialization: provide psi class there is possible

We can provide psi declaration for nested classes to avoid search
through providers and improve stability

^KTIJ-26670 Fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedClassDeserialization.kt

@@ -185,6 +185,7 @@ internal fun deserializeClassToSymbol(
             addDeclaration(constructor)
             constructor
         }
+
         classOrObject.body?.declarations?.forEach { declaration ->
             when (declaration) {
                 is KtConstructor<*> -> addDeclaration(memberDeserializer.loadConstructor(declaration, classOrObject, this))
@@ -192,17 +193,16 @@ internal fun deserializeClassToSymbol(
                 is KtProperty -> addDeclaration(memberDeserializer.loadProperty(declaration, symbol))
                 is KtEnumEntry -> addDeclaration(memberDeserializer.loadEnumEntry(declaration, symbol, classId))
                 is KtClassOrObject -> {
-                    val name = declaration.name
-                        ?: errorWithAttachment("Class doesn't have name $declaration") {
-                            withPsiEntry("class", declaration)
-                        }
-                    val nestedClassId =
-                        classId.createNestedClassId(Name.identifier(name))
-                    deserializeNestedClass(nestedClassId, context)?.fir?.let { addDeclaration(it) }
+                    val name = declaration.name ?: errorWithAttachment("Class doesn't have name $declaration") {
+                        withPsiEntry("class", declaration)
+                    }
+
+                    val nestedClassId = classId.createNestedClassId(Name.identifier(name))
+                    // Add declaration to the context to avoid redundant provider access to the class map
+                    deserializeNestedClass(nestedClassId, context.withClassLikeDeclaration(declaration))?.fir?.let(this::addDeclaration)
                 }
-                is KtTypeAlias -> addDeclaration(
-                    memberDeserializer.loadTypeAlias(declaration, FirTypeAliasSymbol(classId))
-                )
+
+                is KtTypeAlias -> addDeclaration(memberDeserializer.loadTypeAlias(declaration, FirTypeAliasSymbol(classId)))
             }
         }
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirMemberDeserializer.kt

@@ -53,7 +53,7 @@ internal class StubBasedFirDeserializationContext(
     private val initialOrigin: FirDeclarationOrigin,
     val classLikeDeclaration: KtClassLikeDeclaration? = null
 ) {
-    val session: FirSession = moduleData.session
+    val session: FirSession get() = moduleData.session
 
     val allTypeParameters: List<FirTypeParameterSymbol> =
         typeDeserializer.ownTypeParameters + outerTypeParameters
@@ -67,10 +67,10 @@ internal class StubBasedFirDeserializationContext(
         capturesTypeParameters: Boolean = true,
         containingDeclarationSymbol: FirBasedSymbol<*>? = this.outerClassSymbol
     ): StubBasedFirDeserializationContext = StubBasedFirDeserializationContext(
-        moduleData,
-        packageFqName,
-        relativeClassName,
-        StubBasedFirTypeDeserializer(
+        moduleData = moduleData,
+        packageFqName = packageFqName,
+        relativeClassName = relativeClassName,
+        typeDeserializer = StubBasedFirTypeDeserializer(
             moduleData,
             annotationDeserializer,
             typeDeserializer,
@@ -78,11 +78,26 @@ internal class StubBasedFirDeserializationContext(
             owner,
             initialOrigin
         ),
-        annotationDeserializer,
-        containerSource,
-        outerClassSymbol,
-        if (capturesTypeParameters) allTypeParameters else emptyList(),
-        initialOrigin
+        annotationDeserializer = annotationDeserializer,
+        containerSource = containerSource,
+        outerClassSymbol = outerClassSymbol,
+        outerTypeParameters = if (capturesTypeParameters) allTypeParameters else emptyList(),
+        initialOrigin = initialOrigin
+    )
+
+    fun withClassLikeDeclaration(
+        classLikeDeclaration: KtClassLikeDeclaration,
+    ): StubBasedFirDeserializationContext = StubBasedFirDeserializationContext(
+        moduleData = moduleData,
+        packageFqName = packageFqName,
+        relativeClassName = relativeClassName,
+        typeDeserializer = typeDeserializer,
+        annotationDeserializer = annotationDeserializer,
+        containerSource = containerSource,
+        outerClassSymbol = outerClassSymbol,
+        outerTypeParameters = outerTypeParameters,
+        initialOrigin = initialOrigin,
+        classLikeDeclaration = classLikeDeclaration,
     )
 
     val memberDeserializer: StubBasedFirMemberDeserializer = StubBasedFirMemberDeserializer(this, initialOrigin)