[FIR] Allow FirElementSerializer.packagePartProto accept multiple files

This is needed for cases when there are multiple files with same name
 and `@JvmMultifileClass` annotation

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirElementSerializer.kt

@@ -67,15 +67,17 @@ class FirElementSerializer private constructor(
 ) {
     private val contractSerializer = FirContractSerializer()
 
-    fun packagePartProto(packageFqName: FqName, file: FirFile): ProtoBuf.Package.Builder {
+    fun packagePartProto(packageFqName: FqName, files: List<FirFile>): ProtoBuf.Package.Builder {
         val builder = ProtoBuf.Package.newBuilder()
 
-        for (declaration in file.declarations) {
-            when (declaration) {
-                is FirProperty -> propertyProto(declaration)?.let { builder.addProperty(it) }
-                is FirSimpleFunction -> functionProto(declaration)?.let { builder.addFunction(it) }
-                is FirTypeAlias -> typeAliasProto(declaration)?.let { builder.addTypeAlias(it) }
-                else -> {}
+        for (file in files) {
+            for (declaration in file.declarations) {
+                when (declaration) {
+                    is FirProperty -> propertyProto(declaration)?.let { builder.addProperty(it) }
+                    is FirSimpleFunction -> functionProto(declaration)?.let { builder.addFunction(it) }
+                    is FirTypeAlias -> typeAliasProto(declaration)?.let { builder.addTypeAlias(it) }
+                    else -> {}
+                }
             }
         }
 

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/firKlibSerialization.kt

@@ -26,7 +26,7 @@ fun serializeSingleFirFile(file: FirFile, session: FirSession, scopeSession: Sco
     // TODO: split package fragment (see klib serializer)
     // TODO: handle incremental/monolothic (see klib serializer) - maybe externally
 
-    val packageProto = packageSerializer.packagePartProto(file.packageFqName, file).build()
+    val packageProto = packageSerializer.packagePartProto(file.packageFqName, listOf(file)).build()
 
     fun List<FirDeclaration>.makeClassesProtoWithNested(): List<Pair<ProtoBuf.Class, Int>> =
         // TODO: filter out expects

compiler/fir/fir2ir/jvm-backend/src/org/jetbrains/kotlin/fir/backend/jvm/FirMetadataSerializer.kt

@@ -118,7 +118,7 @@ class FirMetadataSerializer(
         val message = when (metadata) {
             is FirMetadataSource.Class -> serializer!!.classProto(metadata.fir).build()
             is FirMetadataSource.File ->
-                serializer!!.packagePartProto(metadata.fir.packageFqName, metadata.fir).apply {
+                serializer!!.packagePartProto(metadata.files.first().packageFqName, metadata.files).apply {
                     serializerExtension.serializeJvmPackage(this)
                 }.build()
             is FirMetadataSource.Function -> {

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -96,7 +96,7 @@ class Fir2IrVisitor(
                 it.toIrDeclaration()
             }
             annotationGenerator.generate(this, file)
-            metadata = FirMetadataSource.File(file)
+            metadata = FirMetadataSource.File(listOf(file))
         }
     }
 

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/FirMetadataSource.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.backend
 
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isConst
 import org.jetbrains.kotlin.ir.declarations.MetadataSource
@@ -13,10 +12,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
 
 sealed class FirMetadataSource : MetadataSource {
-    abstract val fir: FirDeclaration
-
-    val declarationSiteSession: FirSession
-        get() = fir.moduleData.session
+    abstract val fir: FirDeclaration?
 
     override val name: Name?
         get() = when (val fir = fir) {
@@ -27,8 +23,11 @@ sealed class FirMetadataSource : MetadataSource {
             else -> null
         }
 
-    class File(override val fir: FirFile) : FirMetadataSource(), MetadataSource.File {
+    class File(val files: List<FirFile>) : FirMetadataSource(), MetadataSource.File {
         override var serializedIr: ByteArray? = null
+
+        override val fir: FirDeclaration?
+            get() = null
     }
 
     class Class(override val fir: FirClass) : FirMetadataSource(), MetadataSource.Class {

compiler/incremental-compilation-impl/src/org/jetbrains/kotlin/incremental/incrementalFirCacheUtils.kt

@@ -70,7 +70,7 @@ internal fun collectNewDirtySources(
             }
 
             override fun visitFile(file: FirFile, data: MutableList<MetadataSerializer>) {
-                val metadata = FirMetadataSource.File(file)
+                val metadata = FirMetadataSource.File(listOf(file))
                 withMetadataSerializer(metadata, data) {
                     file.acceptChildren(this, data)
                     // TODO: compare package fragments?