[K2, MPP] Remove redundant expect declarations from klib metadata

^KT-57250 Fixed

Introduce flat Fir2IrActualizedResult

It contains output from Fir2Ir and IrActualizer

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/pipeline/convertToIr.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.pipeline
 
+import org.jetbrains.kotlin.backend.common.actualizer.IrActualizationResult
 import org.jetbrains.kotlin.backend.common.actualizer.IrActualizer
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.jvm.serialization.JvmIdSignatureDescriptor
@@ -22,6 +23,7 @@ import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmDescriptorMangler
 import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmIrMangler
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImpl
 import org.jetbrains.kotlin.ir.util.IdSignatureComposer
 import org.jetbrains.kotlin.ir.util.KotlinMangler
@@ -34,13 +36,20 @@ data class ModuleCompilerAnalyzedOutput(
     val fir: List<FirFile>
 )
 
+data class Fir2IrActualizedResult(
+    val irModuleFragment: IrModuleFragment,
+    val components: Fir2IrComponents,
+    val pluginContext: Fir2IrPluginContext,
+    val irActualizationResult: IrActualizationResult?,
+)
+
 fun FirResult.convertToIrAndActualizeForJvm(
     fir2IrExtensions: Fir2IrExtensions,
     irGeneratorExtensions: Collection<IrGenerationExtension>,
     linkViaSignatures: Boolean,
     diagnosticReporter: DiagnosticReporter,
     languageVersionSettings: LanguageVersionSettings,
-): Fir2IrResult = this.convertToIrAndActualize(
+): Fir2IrActualizedResult = this.convertToIrAndActualize(
     fir2IrExtensions,
     irGeneratorExtensions,
     linkViaSignatures = linkViaSignatures,
@@ -63,8 +72,9 @@ fun FirResult.convertToIrAndActualize(
     diagnosticReporter: DiagnosticReporter,
     languageVersionSettings: LanguageVersionSettings,
     fir2IrResultPostCompute: Fir2IrResult.() -> Unit = {},
-): Fir2IrResult {
-    val result: Fir2IrResult
+): Fir2IrActualizedResult {
+    val fir2IrResult: Fir2IrResult
+    val actualizationResult: IrActualizationResult?
 
     val commonMemberStorage = Fir2IrCommonMemberStorage(
         generateSignatures = linkViaSignatures,
@@ -75,7 +85,7 @@ fun FirResult.convertToIrAndActualize(
     when (outputs.size) {
         0 -> error("No modules found")
         1 -> {
-            result = outputs.single().convertToIr(
+            fir2IrResult = outputs.single().convertToIr(
                 fir2IrExtensions,
                 irGeneratorExtensions,
                 linkViaSignatures = linkViaSignatures,
@@ -85,6 +95,7 @@ fun FirResult.convertToIrAndActualize(
                 visibilityConverter,
                 kotlinBuiltIns,
             )
+            actualizationResult = null
         }
         else -> {
             val platformOutput = outputs.last()
@@ -107,7 +118,7 @@ fun FirResult.convertToIrAndActualize(
                     }
                 }
             }
-            result = platformOutput.convertToIr(
+            fir2IrResult = platformOutput.convertToIr(
                 fir2IrExtensions,
                 irGeneratorExtensions,
                 linkViaSignatures = linkViaSignatures,
@@ -120,8 +131,8 @@ fun FirResult.convertToIrAndActualize(
                 fir2IrResultPostCompute(it)
             }
 
-            IrActualizer.actualize(
-                result.irModuleFragment,
+            actualizationResult = IrActualizer.actualize(
+                fir2IrResult.irModuleFragment,
                 commonIrOutputs.map { it.irModuleFragment },
                 diagnosticReporter,
                 languageVersionSettings
@@ -129,7 +140,8 @@ fun FirResult.convertToIrAndActualize(
         }
     }
 
-    return result
+    val (irModuleFragment, components, pluginContext) = fir2IrResult
+    return Fir2IrActualizedResult(irModuleFragment, components, pluginContext, actualizationResult)
 }
 
 private fun ModuleCompilerAnalyzedOutput.convertToIr(

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

@@ -69,15 +69,24 @@ class FirElementSerializer private constructor(
     private val contractSerializer = FirContractSerializer()
     private val extensionDeclarationProviders = session.extensionService.declarationForMetadataProviders
 
-    fun packagePartProto(packageFqName: FqName, files: List<FirFile>): ProtoBuf.Package.Builder {
+    fun packagePartProto(
+        packageFqName: FqName,
+        files: List<FirFile>,
+        actualizedExpectDeclarations: Set<FirDeclaration>?
+    ): ProtoBuf.Package.Builder {
         val builder = ProtoBuf.Package.newBuilder()
 
         fun addDeclaration(declaration: FirDeclaration, onUnsupportedDeclaration: (FirDeclaration) -> Unit) {
-            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 -> onUnsupportedDeclaration(declaration)
+            if (declaration is FirMemberDeclaration) {
+                if (!declaration.shouldBeSerialized(actualizedExpectDeclarations)) return
+                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 -> onUnsupportedDeclaration(declaration)
+                }
+            } else {
+                onUnsupportedDeclaration(declaration)
             }
         }
 

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.serialization.SerializableStringTable
 
 fun serializeSingleFirFile(
     file: FirFile, session: FirSession, scopeSession: ScopeSession,
+    actualizedExpectDeclarations: Set<FirDeclaration>?,
     serializerExtension: FirKLibSerializerExtension,
     languageVersionSettings: LanguageVersionSettings,
 ): ProtoBuf.PackageFragment {
@@ -31,21 +32,26 @@ fun serializeSingleFirFile(
     // TODO: split package fragment (see klib serializer)
     // TODO: handle incremental/monolothic (see klib serializer) - maybe externally
 
-    val packageProto = packageSerializer.packagePartProto(file.packageFqName, listOf(file)).build()
+    val packageProto = packageSerializer.packagePartProto(file.packageFqName, listOf(file), actualizedExpectDeclarations).build()
 
-    fun List<FirDeclaration>.makeClassesProtoWithNested(): List<Pair<ProtoBuf.Class, Int>> =
-        // TODO: filter out expects
-        filterIsInstance<FirClass>().sortedBy { it.classId.asFqNameString() }.flatMap {
-            val classSerializer = FirElementSerializer.create(session, scopeSession, it, serializerExtension, null,
-                approximator, languageVersionSettings)
-            val index = classSerializer.stringTable.getFqNameIndex(it)
-            listOf(classSerializer.classProto(it).build() to index) + it.declarations.makeClassesProtoWithNested()
-        }
+    fun List<FirDeclaration>.makeClassesProtoWithNested(): List<Pair<ProtoBuf.Class, Int>> {
+        return filterIsInstance<FirClass>().filter { it.shouldBeSerialized(actualizedExpectDeclarations) }
+            .sortedBy { it.classId.asFqNameString() }
+            .flatMap {
+                val classSerializer = FirElementSerializer.create(
+                    session, scopeSession, it, serializerExtension, null,
+                    approximator, languageVersionSettings
+                )
+                val index = classSerializer.stringTable.getFqNameIndex(it)
+                listOf(classSerializer.classProto(it).build() to index) + it.declarations.makeClassesProtoWithNested()
+            }
+    }
 
     val classesProto = file.declarations.makeClassesProtoWithNested()
 
     val hasTopLevelDeclarations = file.declarations.any {
-        it is FirProperty || it is FirSimpleFunction || it is FirTypeAlias
+        it is FirMemberDeclaration && it.shouldBeSerialized(actualizedExpectDeclarations) &&
+                (it is FirProperty || it is FirSimpleFunction || it is FirTypeAlias)
     }
 
     return buildKlibPackageFragment(

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

@@ -7,6 +7,9 @@ package org.jetbrains.kotlin.fir.serialization
 
 import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
+import org.jetbrains.kotlin.fir.declarations.utils.isExpect
 import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
 import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
@@ -43,3 +46,7 @@ fun ConeKotlinType.suspendFunctionTypeToFunctionTypeWithContinuation(session: Fi
         attributes = fullyExpandedType.attributes
     )
 }
+
+fun FirMemberDeclaration.shouldBeSerialized(actualizedExpectDeclaration: Set<FirDeclaration>?): Boolean {
+    return !isExpect || actualizedExpectDeclaration == null || this !in actualizedExpectDeclaration
+}

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.backend.jvm
 
+import org.jetbrains.kotlin.backend.common.actualizer.IrActualizationResult
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmBackendExtension
 import org.jetbrains.kotlin.backend.jvm.metadata.MetadataSerializer
@@ -15,7 +16,10 @@ import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.org.objectweb.asm.Type
 
-class FirJvmBackendExtension(private val components: Fir2IrComponents) : JvmBackendExtension {
+class FirJvmBackendExtension(
+    private val components: Fir2IrComponents,
+    private val irActualizationResult: IrActualizationResult?
+) : JvmBackendExtension {
     override fun createSerializer(
         context: JvmBackendContext,
         klass: IrClass,
@@ -23,7 +27,15 @@ class FirJvmBackendExtension(private val components: Fir2IrComponents) : JvmBack
         bindings: JvmSerializationBindings,
         parentSerializer: MetadataSerializer?
     ): MetadataSerializer {
-        return makeFirMetadataSerializerForIrClass(components.session, context, klass, bindings, components, parentSerializer)
+        return makeFirMetadataSerializerForIrClass(
+            components.session,
+            context,
+            klass,
+            bindings,
+            components,
+            parentSerializer,
+            irActualizationResult
+        )
     }
 
     override fun generateMetadataExtraFlags(abiStability: JvmAbiStability?): Int =

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.backend.jvm
 
+import org.jetbrains.kotlin.backend.common.actualizer.IrActualizationResult
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.metadata.MetadataSerializer
@@ -14,6 +15,7 @@ import org.jetbrains.kotlin.config.*
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.backend.Fir2IrComponents
 import org.jetbrains.kotlin.fir.backend.FirMetadataSource
+import org.jetbrains.kotlin.fir.backend.extractFirDeclarations
 import org.jetbrains.kotlin.fir.containingClassForLocal
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.*
@@ -45,7 +47,8 @@ fun makeFirMetadataSerializerForIrClass(
     irClass: IrClass,
     serializationBindings: JvmSerializationBindings,
     components: Fir2IrComponents,
-    parent: MetadataSerializer?
+    parent: MetadataSerializer?,
+    irActualizationResult: IrActualizationResult?
 ): FirMetadataSerializer {
     val approximator = TypeApproximatorForMetadataSerializer(session)
     val localDelegatedProperties = context.localDelegatedProperties[irClass.attributeOwnerId]?.map {
@@ -62,8 +65,9 @@ fun makeFirMetadataSerializerForIrClass(
         approximator,
         makeElementSerializer(
             irClass.metadata, components.session, components.scopeSession, firSerializerExtension, approximator, parent,
-            context.state.configuration.languageVersionSettings,
-        )
+            context.state.configuration.languageVersionSettings
+        ),
+        irActualizationResult
     )
 }
 
@@ -75,7 +79,8 @@ fun makeLocalFirMetadataSerializerForMetadataSource(
     globalSerializationBindings: JvmSerializationBindings,
     parent: MetadataSerializer?,
     targetId: TargetId,
-    configuration: CompilerConfiguration
+    configuration: CompilerConfiguration,
+    irActualizationResult: IrActualizationResult?
 ): FirMetadataSerializer {
     val serializationBindings = JvmSerializationBindings()
     val approximator = TypeApproximatorForMetadataSerializer(session)
@@ -106,8 +111,9 @@ fun makeLocalFirMetadataSerializerForMetadataSource(
         approximator,
         makeElementSerializer(
             metadata, session, scopeSession, firSerializerExtension, approximator, parent,
-            configuration.languageVersionSettings,
-        )
+            configuration.languageVersionSettings
+        ),
+        irActualizationResult
     )
 }
 
@@ -116,16 +122,17 @@ class FirMetadataSerializer(
     private val serializationBindings: JvmSerializationBindings,
     private val serializerExtension: FirJvmSerializerExtension,
     private val approximator: AbstractTypeApproximator,
-    internal val serializer: FirElementSerializer?
+    internal val serializer: FirElementSerializer?,
+    irActualizationResult: IrActualizationResult?
 ) : MetadataSerializer {
+    private val actualizedExpectDeclarations = irActualizationResult.extractFirDeclarations()
 
     override fun serialize(metadata: MetadataSource): Pair<MessageLite, JvmStringTable>? {
         val message = when (metadata) {
             is FirMetadataSource.Class -> serializer!!.classProto(metadata.fir).build()
             is FirMetadataSource.File ->
-                serializer!!.packagePartProto(metadata.files.first().packageFqName, metadata.files).apply {
-                    serializerExtension.serializeJvmPackage(this)
-                }.build()
+                serializer!!.packagePartProto(metadata.files.first().packageFqName, metadata.files, actualizedExpectDeclarations)
+                    .apply { serializerExtension.serializeJvmPackage(this) }.build()
             is FirMetadataSource.Function -> {
                 val withTypeParameters = metadata.fir.copyToFreeAnonymousFunction(approximator)
                 serializationBindings.get(FirJvmSerializerExtension.METHOD_FOR_FIR_FUNCTION, metadata.fir)?.let {

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.backend
 
 import com.intellij.psi.PsiCompiledElement
 import org.jetbrains.kotlin.*
+import org.jetbrains.kotlin.backend.common.actualizer.IrActualizationResult
 import org.jetbrains.kotlin.builtins.StandardNames.DATA_CLASS_COMPONENT_PREFIX
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.startOffsetSkippingComments
@@ -748,3 +749,7 @@ fun FirCallableDeclaration.contextReceiversForFunctionOrContainingProperty(): Li
         this.propertySymbol.fir.contextReceivers
     else
         this.contextReceivers
+
+fun IrActualizationResult?.extractFirDeclarations(): Set<FirDeclaration>? {
+    return this?.actualizedExpectDeclarations?.mapNotNullTo(mutableSetOf()) { ((it as IrMetadataSourceOwner).metadata as FirMetadataSource).fir }
+}

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLightTreeBlackBoxCodegenTestGenerated.java

@@ -33381,6 +33381,12 @@ public class FirLightTreeBlackBoxCodegenTestGenerated extends AbstractFirLightTr
                     runTest("compiler/testData/codegen/box/multiplatform/k2/basic/noArgActualConstructor.kt");
                 }
 
+                @Test
+                @TestMetadata("removeExpectDeclarationsFromMetadata.kt")
+                public void testRemoveExpectDeclarationsFromMetadata() throws Exception {
+                    runTest("compiler/testData/codegen/box/multiplatform/k2/basic/removeExpectDeclarationsFromMetadata.kt");
+                }
+
                 @Test
                 @TestMetadata("valueClasses.kt")
                 public void testValueClasses() throws Exception {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirPsiBlackBoxCodegenTestGenerated.java

@@ -33381,6 +33381,12 @@ public class FirPsiBlackBoxCodegenTestGenerated extends AbstractFirPsiBlackBoxCo
                     runTest("compiler/testData/codegen/box/multiplatform/k2/basic/noArgActualConstructor.kt");
                 }
 
+                @Test
+                @TestMetadata("removeExpectDeclarationsFromMetadata.kt")
+                public void testRemoveExpectDeclarationsFromMetadata() throws Exception {
+                    runTest("compiler/testData/codegen/box/multiplatform/k2/basic/removeExpectDeclarationsFromMetadata.kt");
+                }
+
                 @Test
                 @TestMetadata("valueClasses.kt")
                 public void testValueClasses() throws Exception {