[LL FIR] StubBasedFirDeserializedSymbolProvider: fix loading of nested classes

We cannot load nested class without an outer as in this case we will
miss a symbol of the outer class, which is crucial in some cases
(e.g., during status calculation).
So we have to load the topmost class as it will load all nested classes
with the correct context as well

^KT-62891 Fixed

analysis/analysis-api/testData/components/diagnosticsProvider/diagnostics/overrideProtectedClassReturnFromLibrary.standalone.fir.txt

@@ -1 +0,0 @@
-Diagnostics from elements:

analysis/analysis-api/testData/components/diagnosticsProvider/diagnostics/overrideProtectedClassReturnFromLibrary.txt

@@ -1,4 +1 @@
 Diagnostics from elements:
-  for PSI element of type KtNamedFunction at (15,5-64)
-    EXPOSED_FUNCTION_RETURN_TYPE      text ranges: [(103,123)]
-      PSI: KtNamedFunction at (15,5-64)

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.load.kotlin.FacadeClassSource
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getTopmostParentOfType
 import org.jetbrains.kotlin.psi.stubs.impl.*
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerializerProtocol
@@ -136,12 +137,9 @@ internal open class StubBasedFirDeserializedSymbolProvider(
         classId: ClassId,
         parentContext: StubBasedFirDeserializationContext?,
     ): FirRegularClassSymbol? {
-        val (classLikeDeclaration, context) =
-            if (parentContext?.classLikeDeclaration != null) {
-                parentContext.classLikeDeclaration to null
-            } else {
-                (declarationProvider.getClassLikeDeclarationByClassId(classId) ?: return null) to parentContext
-            }
+        val classLikeDeclaration = parentContext?.classLikeDeclaration
+            ?: declarationProvider.getClassLikeDeclarationByClassId(classId)
+            ?: return null
 
         val symbol = FirRegularClassSymbol(classId)
         if (classLikeDeclaration is KtClassOrObject) {
@@ -153,11 +151,12 @@ internal open class StubBasedFirDeserializedSymbolProvider(
                 moduleData,
                 StubBasedAnnotationDeserializer(session),
                 kotlinScopeProvider,
-                parentContext = context,
+                parentContext = parentContext,
                 containerSource = JvmStubDeserializedContainerSource(classId),
                 deserializeNestedClass = this::getClass,
-                initialOrigin = getDeclarationOriginFor(classLikeDeclaration.containingKtFile)
+                initialOrigin = parentContext?.initialOrigin ?: getDeclarationOriginFor(classLikeDeclaration.containingKtFile)
             )
+
             return symbol
         }
         return null
@@ -315,11 +314,35 @@ internal open class StubBasedFirDeserializedSymbolProvider(
 
     override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
         if (!symbolNamesProvider.mayHaveTopLevelClassifier(classId)) return null
+
+        classId.takeIf(ClassId::isNestedClass)?.outermostClassId?.let { outermostClassId ->
+            // We have to load root declaration to initialize nested classes correctly
+            getClassLikeSymbolByClassId(outermostClassId)
+
+            // Nested declarations already loaded
+            val computedValue = classCache.getValueIfComputed(classId) ?: typeAliasCache.getValueIfComputed(classId)
+            computedValue?.let { return it }
+        }
+
         return getClass(classId) ?: getTypeAlias(classId)
     }
 
     @FirSymbolProviderInternals
     override fun getClassLikeSymbolByClassId(classId: ClassId, classLikeDeclaration: KtClassLikeDeclaration): FirClassLikeSymbol<*>? {
+        val topmostClassLikeDeclaration = classLikeDeclaration.takeIf {
+            classId.isNestedClass
+        }?.getTopmostParentOfType<KtClassLikeDeclaration>()
+
+        val outermostClassId = topmostClassLikeDeclaration?.getClassId()
+        val cache = if (classLikeDeclaration is KtClassOrObject) classCache else typeAliasCache
+        if (outermostClassId != null) {
+            // We have to load root declaration to initialize nested classes correctly
+            getClassLikeSymbolByClassId(outermostClassId, topmostClassLikeDeclaration)
+
+            // Nested declarations already loaded
+            cache.getValueIfComputed(classId)?.let { return it }
+        }
+
         val annotationDeserializer = StubBasedAnnotationDeserializer(session)
         val classOrigin = getDeclarationOriginFor(classLikeDeclaration.containingKtFile)
         val deserializationContext = StubBasedFirDeserializationContext(
@@ -342,7 +365,6 @@ internal open class StubBasedFirDeserializedSymbolProvider(
             classLikeDeclaration,
         )
 
-        val cache = if (classLikeDeclaration is KtClassOrObject) classCache else typeAliasCache
         return cache.getNotNullValueForNotNullContext(classId, deserializationContext)
     }
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2022 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.
  */
 
@@ -50,7 +50,7 @@ internal class StubBasedFirDeserializationContext(
     val containerSource: DeserializedContainerSource?,
     val outerClassSymbol: FirRegularClassSymbol?,
     val outerTypeParameters: List<FirTypeParameterSymbol>,
-    private val initialOrigin: FirDeclarationOrigin,
+    val initialOrigin: FirDeclarationOrigin,
     val classLikeDeclaration: KtClassLikeDeclaration? = null
 ) {
     val session: FirSession get() = moduleData.session
@@ -65,7 +65,7 @@ internal class StubBasedFirDeserializationContext(
         outerClassSymbol: FirRegularClassSymbol? = this.outerClassSymbol,
         annotationDeserializer: StubBasedAnnotationDeserializer = this.annotationDeserializer,
         capturesTypeParameters: Boolean = true,
-        containingDeclarationSymbol: FirBasedSymbol<*>? = this.outerClassSymbol
+        containingDeclarationSymbol: FirBasedSymbol<*>? = outerClassSymbol
     ): StubBasedFirDeserializationContext = StubBasedFirDeserializationContext(
         moduleData = moduleData,
         packageFqName = packageFqName,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirTypeDeserializer.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.
  */
 
@@ -63,6 +63,7 @@ internal class StubBasedFirTypeDeserializer(
                 val symbol = FirTypeParameterSymbol().also {
                     typeParametersByName[name.asString()] = it
                 }
+
                 builders += FirTypeParameterBuilder().apply {
                     source = KtRealPsiSourceElement(typeParameter)
                     moduleData = this@StubBasedFirTypeDeserializer.moduleData
@@ -70,7 +71,11 @@ internal class StubBasedFirTypeDeserializer(
                     origin = initialOrigin
                     this.name = name
                     this.symbol = symbol
-                    this.containingDeclarationSymbol = containingSymbol ?: error("Top-level type parameter ???")
+                    this.containingDeclarationSymbol = containingSymbol ?: errorWithAttachment("Top-level type parameter ???") {
+                        withPsiEntry("owner", owner)
+                        withPsiEntry("parameter", typeParameter)
+                    }
+
                     variance = typeParameter.variance
                     isReified = typeParameter.hasModifier(KtTokens.REIFIED_KEYWORD)
                     annotations += annotationDeserializer.loadAnnotations(typeParameter)