[FIR] Fix container source for class members

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/deserialization/KotlinDeserializedJvmSymbolsProvider.kt

@@ -120,6 +120,7 @@ class KotlinDeserializedJvmSymbolsProvider(
                 FirDeserializationContext.createForPackage(
                     packageFqName, packageProto, nameResolver, session,
                     JvmBinaryAnnotationDeserializer(session),
+                    source
                 ),
                 source,
             )
@@ -332,7 +333,8 @@ class KotlinDeserializedJvmSymbolsProvider(
                 classId, classProto, symbol, nameResolver, session,
                 JvmBinaryAnnotationDeserializer(session),
                 kotlinScopeProvider,
-                parentContext, this::findAndDeserializeClass,
+                parentContext, KotlinJvmBinarySourceElement(kotlinJvmBinaryClass),
+                this::findAndDeserializeClass
             )
 
             classesCache[classId] = symbol
@@ -360,7 +362,7 @@ class KotlinDeserializedJvmSymbolsProvider(
         val functionIds = part.topLevelFunctionNameIndex[name] ?: return emptyList()
         return functionIds.map { part.proto.getFunction(it) }
             .map {
-                val firNamedFunction = part.context.memberDeserializer.loadFunction(it, containerSource = part.source) as FirSimpleFunctionImpl
+                val firNamedFunction = part.context.memberDeserializer.loadFunction(it) as FirSimpleFunctionImpl
                 firNamedFunction.symbol
             }
     }
@@ -369,7 +371,7 @@ class KotlinDeserializedJvmSymbolsProvider(
         val propertyIds = part.topLevelPropertyNameIndex[name] ?: return emptyList()
         return propertyIds.map { part.proto.getProperty(it) }
             .map {
-                val firProperty = part.context.memberDeserializer.loadProperty(it, containerSource = part.source)
+                val firProperty = part.context.memberDeserializer.loadProperty(it)
                 firProperty.symbol
             }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.metadata.deserialization.supertypes
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.serialization.deserialization.ProtoEnumFlags
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 import org.jetbrains.kotlin.serialization.deserialization.getName
 
 fun deserializeClassToSymbol(
@@ -45,6 +46,7 @@ fun deserializeClassToSymbol(
     defaultAnnotationDeserializer: AbstractAnnotationDeserializer?,
     scopeProvider: KotlinScopeProvider,
     parentContext: FirDeserializationContext? = null,
+    containerSource: DeserializedContainerSource? = null,
     deserializeNestedClass: (ClassId, FirDeserializationContext) -> FirRegularClassSymbol?
 ) {
     val flags = classProto.flags
@@ -71,7 +73,8 @@ fun deserializeClassToSymbol(
             classId.relativeClassName
         ) ?: FirDeserializationContext.createForClass(
             classId, classProto, nameResolver, session,
-            defaultAnnotationDeserializer ?: FirBuiltinAnnotationDeserializer(session)
+            defaultAnnotationDeserializer ?: FirBuiltinAnnotationDeserializer(session),
+            containerSource
         )
     classBuilder.apply {
         this.session = session

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt

@@ -37,6 +37,7 @@ class FirDeserializationContext(
     val relativeClassName: FqName?,
     val typeDeserializer: FirTypeDeserializer,
     val annotationDeserializer: AbstractAnnotationDeserializer,
+    val containerSource: DeserializedContainerSource?,
     val components: FirDeserializationComponents
 ) {
     fun childContext(
@@ -49,7 +50,7 @@ class FirDeserializationContext(
         FirTypeDeserializer(
             session, nameResolver, typeTable, typeParameterProtos, typeDeserializer
         ),
-        annotationDeserializer, components
+        annotationDeserializer, containerSource, components
     )
 
     val memberDeserializer: FirMemberDeserializer = FirMemberDeserializer(this)
@@ -60,7 +61,8 @@ class FirDeserializationContext(
             packageProto: ProtoBuf.Package,
             nameResolver: NameResolver,
             session: FirSession,
-            annotationDeserializer: AbstractAnnotationDeserializer
+            annotationDeserializer: AbstractAnnotationDeserializer,
+            containerSource: DeserializedContainerSource?
         ) = createRootContext(
             nameResolver,
             TypeTable(packageProto.typeTable),
@@ -68,7 +70,8 @@ class FirDeserializationContext(
             annotationDeserializer,
             fqName,
             relativeClassName = null,
-            typeParameterProtos = emptyList()
+            typeParameterProtos = emptyList(),
+            containerSource
         )
 
         fun createForClass(
@@ -76,7 +79,8 @@ class FirDeserializationContext(
             classProto: ProtoBuf.Class,
             nameResolver: NameResolver,
             session: FirSession,
-            annotationDeserializer: AbstractAnnotationDeserializer
+            annotationDeserializer: AbstractAnnotationDeserializer,
+            containerSource: DeserializedContainerSource?
         ) = createRootContext(
             nameResolver,
             TypeTable(classProto.typeTable),
@@ -84,7 +88,8 @@ class FirDeserializationContext(
             annotationDeserializer,
             classId.packageFqName,
             classId.relativeClassName,
-            classProto.typeParameterList
+            classProto.typeParameterList,
+            containerSource
         )
 
         private fun createRootContext(
@@ -94,7 +99,8 @@ class FirDeserializationContext(
             annotationDeserializer: AbstractAnnotationDeserializer,
             packageFqName: FqName,
             relativeClassName: FqName?,
-            typeParameterProtos: List<ProtoBuf.TypeParameter>
+            typeParameterProtos: List<ProtoBuf.TypeParameter>,
+            containerSource: DeserializedContainerSource?
         ): FirDeserializationContext {
             return FirDeserializationContext(
                 nameResolver, typeTable,
@@ -110,6 +116,7 @@ class FirDeserializationContext(
                     null
                 ),
                 annotationDeserializer,
+                containerSource,
                 FirDeserializationComponents()
             )
         }
@@ -150,7 +157,7 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
         }
     }
 
-    fun loadProperty(proto: ProtoBuf.Property, containerSource: DeserializedContainerSource? = null): FirProperty {
+    fun loadProperty(proto: ProtoBuf.Property): FirProperty {
         val flags = if (proto.hasFlags()) proto.flags else loadOldFlags(proto.oldFlags)
         val callableName = c.nameResolver.getName(proto.name)
         val symbol = FirPropertySymbol(CallableId(c.packageFqName, c.relativeClassName, callableName))
@@ -185,11 +192,11 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
             setter = if (isVar) {
                 FirDefaultPropertySetter(null, c.session, returnTypeRef, ProtoEnumFlags.visibility(Flags.VISIBILITY.get(setterFlags)))
             } else null
-            this.containerSource = containerSource
+            this.containerSource = c.containerSource
         }
     }
 
-    fun loadFunction(proto: ProtoBuf.Function, containerSource: DeserializedContainerSource? = null): FirSimpleFunction {
+    fun loadFunction(proto: ProtoBuf.Function): FirSimpleFunction {
         val flags = if (proto.hasFlags()) proto.flags else loadOldFlags(proto.oldFlags)
 
         val receiverAnnotations =
@@ -229,7 +236,7 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
             typeParameters += local.typeDeserializer.ownTypeParameters.map { it.fir }
             valueParameters += local.memberDeserializer.valueParameters(proto.valueParameterList)
             annotations += local.annotationDeserializer.loadFunctionAnnotations(proto, local.nameResolver)
-            this.containerSource = containerSource
+            this.containerSource = c.containerSource
         }
         if (proto.hasContract()) {
             val contractDescription = contractDeserializer.loadContract(proto.contract, simpleFunction)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/impl/FirBuiltinSymbolProvider.kt

@@ -78,6 +78,7 @@ class FirBuiltinSymbolProvider(val session: FirSession, val kotlinScopeProvider:
             FirDeserializationContext.createForPackage(
                 fqName, packageProto.`package`, nameResolver, session,
                 FirBuiltinAnnotationDeserializer(session),
+                containerSource = null
             ).memberDeserializer
         }
 
@@ -99,6 +100,7 @@ class FirBuiltinSymbolProvider(val session: FirSession, val kotlinScopeProvider:
                 deserializeClassToSymbol(
                     classId, classProto, symbol, nameResolver, session,
                     null, kotlinScopeProvider, parentContext,
+                    null,
                     this::findAndDeserializeClass,
                 )
             }