[FIR] Handle expect declarations in JVM backend

1. Do not generate bytecode for expect declarations

2. Serialize @OptionalExpectation annotations into .kotlin_module file

^KT-62931: Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/ClassFileFactory.java

@@ -116,7 +116,6 @@ public class ClassFileFactory implements OutputFileCollection {
     public void done() {
         if (!isDone) {
             isDone = true;
-            writeModuleMappings();
             for (ClassFileFactoryFinalizerExtension extension : finalizers) {
                 extension.finalizeClassFactory(this);
             }
@@ -127,20 +126,8 @@ public class ClassFileFactory implements OutputFileCollection {
         generators.clear();
     }
 
-    private void writeModuleMappings() {
-        JvmModuleProtoBuf.Module.Builder builder = JvmModuleProtoBuf.Module.newBuilder();
-        String outputFilePath = getMappingFileName(state.getModuleName());
-
-        StringTableImpl stringTable = new StringTableImpl();
-        ClassFileUtilsKt.addDataFromCompiledModule(builder, packagePartRegistry, stringTable, state);
-
-        Pair<ProtoBuf.StringTable, ProtoBuf.QualifiedNameTable> tables = stringTable.buildProto();
-        builder.setStringTable(tables.getFirst());
-        builder.setQualifiedNameTable(tables.getSecond());
-
-        JvmModuleProtoBuf.Module moduleProto = builder.build();
-
-        generators.put(outputFilePath, new OutAndSourceFileList(CollectionsKt.toList(sourceFiles)) {
+    public void setModuleMapping(JvmModuleProtoBuf.Module moduleProto) {
+        generators.put(getMappingFileName(state.getModuleName()), new OutAndSourceFileList(CollectionsKt.toList(sourceFiles)) {
             @Override
             public byte[] asBytes(ClassBuilderFactory factory) {
                 int flags = 0;

compiler/backend/src/org/jetbrains/kotlin/codegen/CodegenFactory.kt

@@ -23,10 +23,12 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.fileClasses.JvmFileClassUtil
+import org.jetbrains.kotlin.metadata.jvm.JvmModuleProtoBuf
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.serialization.StringTableImpl
 
 interface CodegenFactory {
     fun convertToIr(input: IrConversionInput): BackendInput
@@ -126,7 +128,20 @@ object DefaultCodegenFactory : CodegenFactory {
     }
 
     override fun invokeCodegen(input: CodegenFactory.CodegenInput) {
-        // Do nothing
+        generateModuleMetadata(input)
+    }
+
+    private fun generateModuleMetadata(result: CodegenFactory.CodegenInput) {
+        val builder = JvmModuleProtoBuf.Module.newBuilder()
+
+        val stringTable = StringTableImpl()
+        builder.addDataFromCompiledModule(stringTable, result.state)
+
+        val (stringTableProto, qualifiedNameTableProto) = stringTable.buildProto()
+        builder.setStringTable(stringTableProto)
+        builder.setQualifiedNameTable(qualifiedNameTableProto)
+
+        result.state.factory.setModuleMapping(builder.build())
     }
 
     private fun generateMultifileClass(state: GenerationState, multifileClassFqName: FqName, files: Collection<KtFile>) {

compiler/backend/src/org/jetbrains/kotlin/codegen/PackagePartRegistry.kt

@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.name.FqName
 
 class PackagePartRegistry {
     val parts = mutableMapOf<FqName, PackageParts>()
+
+    // Drop after old BE removal
     val optionalAnnotations = mutableListOf<ClassDescriptor>()
 
     fun addPart(packageFqName: FqName, partInternalName: String, facadeInternalName: String?) {

compiler/backend/src/org/jetbrains/kotlin/codegen/classFileUtils.kt

@@ -37,9 +37,8 @@ fun List<OutputFile>.filterClassFiles(): List<OutputFile> {
     return filter { it.relativePath.endsWith(".class") }
 }
 
-fun JvmModuleProtoBuf.Module.Builder.addDataFromCompiledModule(
-    registry: PackagePartRegistry, stringTable: StringTableImpl, state: GenerationState
-) {
+fun JvmModuleProtoBuf.Module.Builder.addDataFromCompiledModule(stringTable: StringTableImpl, state: GenerationState) {
+    val registry = state.factory.packagePartRegistry
     for (part in registry.parts.values.addCompiledPartsAndSort(state)) {
         part.addTo(this)
     }
@@ -74,7 +73,7 @@ class JvmOptionalAnnotationSerializerExtension(
     override fun shouldUseTypeTable(): Boolean = true
 }
 
-private fun Iterable<PackageParts>.addCompiledPartsAndSort(state: GenerationState): List<PackageParts> =
+fun Iterable<PackageParts>.addCompiledPartsAndSort(state: GenerationState): List<PackageParts> =
     addCompiledParts(state).sortedBy { it.packageFqName }
 
 private fun Iterable<PackageParts>.addCompiledParts(state: GenerationState): List<PackageParts> {
@@ -94,7 +93,7 @@ private fun Iterable<PackageParts>.addCompiledParts(state: GenerationState): Lis
         }
 }
 
-private fun GenerationState.loadCompiledModule(): ModuleMapping? {
+fun GenerationState.loadCompiledModule(): ModuleMapping? {
     val moduleMappingData = incrementalCacheForThisTarget?.getModuleMappingData() ?: return null
     return ModuleMapping.loadModuleMapping(moduleMappingData, "<incremental>", deserializationConfiguration) { version ->
         throw IllegalStateException("Version of the generated module cannot be incompatible: $version")

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/KlibBasedAnnotationDeserializer.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.session
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.deserialization.AbstractAnnotationDeserializerWithTypeAnnotations
 import org.jetbrains.kotlin.library.metadata.KlibMetadataSerializerProtocol
 
 class KlibBasedAnnotationDeserializer(session: FirSession) :

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/AbstractAnnotationDeserializerWithTypeAnnotations.kt

@@ -1,12 +1,11 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 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.
  */
 
-package org.jetbrains.kotlin.fir.session
+package org.jetbrains.kotlin.fir.deserialization
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.deserialization.AbstractAnnotationDeserializer
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.NameResolver

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/MetadataBasedAnnotationDeserializer.kt

@@ -1,9 +1,9 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 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.
  */
 
-package org.jetbrains.kotlin.fir.session
+package org.jetbrains.kotlin.fir.deserialization
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerializerProtocol

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

@@ -7,8 +7,6 @@ package org.jetbrains.kotlin.fir.serialization
 
 import org.jetbrains.kotlin.constant.ConstantValue
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
-import org.jetbrains.kotlin.fir.FirAnnotationContainer
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.serialization.constant.toConstantValue
@@ -35,7 +33,7 @@ abstract class FirSerializerExtensionBase(
         versionRequirementTable: MutableVersionRequirementTable,
         childSerializer: FirElementSerializer
     ) {
-        klass.serializeAnnotations(proto, protocol.classAnnotation)
+        klass.serializeAnnotations(session, additionalMetadataProvider, annotationSerializer, proto, protocol.classAnnotation)
     }
 
     override fun serializeScript(
@@ -50,7 +48,7 @@ abstract class FirSerializerExtensionBase(
         proto: ProtoBuf.Constructor.Builder,
         childSerializer: FirElementSerializer
     ) {
-        constructor.serializeAnnotations(proto, protocol.constructorAnnotation)
+        constructor.serializeAnnotations(session, additionalMetadataProvider, annotationSerializer, proto, protocol.constructorAnnotation)
     }
 
     override fun serializeFunction(
@@ -59,8 +57,14 @@ abstract class FirSerializerExtensionBase(
         versionRequirementTable: MutableVersionRequirementTable?,
         childSerializer: FirElementSerializer
     ) {
-        function.serializeAnnotations(proto, protocol.functionAnnotation)
-        function.receiverParameter?.serializeAnnotations(proto, protocol.functionExtensionReceiverAnnotation)
+        function.serializeAnnotations(session, additionalMetadataProvider, annotationSerializer, proto, protocol.functionAnnotation)
+        function.receiverParameter?.serializeAnnotations(
+            session,
+            additionalMetadataProvider,
+            annotationSerializer,
+            proto,
+            protocol.functionExtensionReceiverAnnotation
+        )
     }
 
     override fun serializeProperty(
@@ -72,7 +76,7 @@ abstract class FirSerializerExtensionBase(
         val fieldPropertyAnnotations = mutableListOf<FirAnnotation>()
         val delegatePropertyAnnotations = mutableListOf<FirAnnotation>()
 
-        for (annotation in property.backingField?.allRequiredAnnotations(session).orEmpty()) {
+        for (annotation in property.backingField?.allRequiredAnnotations(session, additionalMetadataProvider).orEmpty()) {
             val destination = when (annotation.useSiteTarget) {
                 AnnotationUseSiteTarget.PROPERTY_DELEGATE_FIELD -> delegatePropertyAnnotations
                 else -> fieldPropertyAnnotations
@@ -80,13 +84,31 @@ abstract class FirSerializerExtensionBase(
             destination += annotation
         }
 
-        property.allRequiredAnnotations(session).serializeAnnotations(proto, protocol.propertyAnnotation)
+        property.allRequiredAnnotations(session, additionalMetadataProvider).serializeAnnotations(proto, protocol.propertyAnnotation)
         fieldPropertyAnnotations.serializeAnnotations(proto, protocol.propertyBackingFieldAnnotation)
         delegatePropertyAnnotations.serializeAnnotations(proto, protocol.propertyDelegatedFieldAnnotation)
 
-        property.getter?.serializeAnnotations(proto, protocol.propertyGetterAnnotation)
-        property.setter?.serializeAnnotations(proto, protocol.propertySetterAnnotation)
-        property.receiverParameter?.serializeAnnotations(proto, protocol.propertyExtensionReceiverAnnotation)
+        property.getter?.serializeAnnotations(
+            session,
+            additionalMetadataProvider,
+            annotationSerializer,
+            proto,
+            protocol.propertyGetterAnnotation
+        )
+        property.setter?.serializeAnnotations(
+            session,
+            additionalMetadataProvider,
+            annotationSerializer,
+            proto,
+            protocol.propertySetterAnnotation
+        )
+        property.receiverParameter?.serializeAnnotations(
+            session,
+            additionalMetadataProvider,
+            annotationSerializer,
+            proto,
+            protocol.propertyExtensionReceiverAnnotation
+        )
 
         if (!Flags.HAS_CONSTANT.get(proto.flags)) return
         property.initializer?.toConstantValue<ConstantValue<*>>(session, constValueProvider)?.let {
@@ -95,11 +117,11 @@ abstract class FirSerializerExtensionBase(
     }
 
     override fun serializeEnumEntry(enumEntry: FirEnumEntry, proto: ProtoBuf.EnumEntry.Builder) {
-        enumEntry.serializeAnnotations(proto, protocol.enumEntryAnnotation)
+        enumEntry.serializeAnnotations(session, additionalMetadataProvider, annotationSerializer, proto, protocol.enumEntryAnnotation)
     }
 
     override fun serializeValueParameter(parameter: FirValueParameter, proto: ProtoBuf.ValueParameter.Builder) {
-        parameter.serializeAnnotations(proto, protocol.parameterAnnotation)
+        parameter.serializeAnnotations(session, additionalMetadataProvider, annotationSerializer, proto, protocol.parameterAnnotation)
     }
 
     override fun serializeTypeAnnotations(annotations: List<FirAnnotation>, proto: ProtoBuf.Type.Builder) {
@@ -107,19 +129,13 @@ abstract class FirSerializerExtensionBase(
     }
 
     override fun serializeTypeParameter(typeParameter: FirTypeParameter, proto: ProtoBuf.TypeParameter.Builder) {
-        typeParameter.serializeAnnotations(proto, protocol.typeParameterAnnotation)
-    }
-
-    @Suppress("Reformat")
-    private fun <
-        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
-        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
-    > FirAnnotationContainer.serializeAnnotations(
-        proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
-        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?
-    ) {
-        if (extension == null) return
-        this.allRequiredAnnotations(session).serializeAnnotations(proto, extension)
+        typeParameter.serializeAnnotations(
+            session,
+            additionalMetadataProvider,
+            annotationSerializer,
+            proto,
+            protocol.typeParameterAnnotation
+        )
     }
 
     @Suppress("Reformat")
@@ -135,28 +151,4 @@ abstract class FirSerializerExtensionBase(
             proto.addExtensionOrNull(extension, annotationSerializer.serializeAnnotation(annotation))
         }
     }
-
-    @Suppress("Reformat")
-    private fun <
-        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
-        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
-        Type
-    > GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>.addExtensionOrNull(
-        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<Type>>,
-        value: Type?
-    ) {
-        if (value != null) {
-            addExtension(extension, value)
-        }
-    }
-
-    private fun FirAnnotationContainer.allRequiredAnnotations(session: FirSession): List<FirAnnotation> {
-        val nonSourceAnnotations = nonSourceAnnotations(session)
-        val additionalMetadataAnnotationsProvider = additionalMetadataProvider
-        return if (this is FirDeclaration && additionalMetadataAnnotationsProvider != null) {
-            nonSourceAnnotations + additionalMetadataAnnotationsProvider.findGeneratedAnnotationsFor(this)
-        } else {
-            nonSourceAnnotations
-        }
-    }
 }

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

@@ -6,19 +6,24 @@
 package org.jetbrains.kotlin.fir.serialization
 
 import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
+import org.jetbrains.kotlin.fir.declarations.nonSourceAnnotations
 import org.jetbrains.kotlin.fir.declarations.utils.isExpect
 import org.jetbrains.kotlin.fir.declarations.utils.isInterface
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.protobuf.GeneratedMessageLite
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
 import org.jetbrains.kotlin.types.model.SimpleTypeMarker
 
@@ -59,3 +64,56 @@ fun FirMemberDeclaration.isNotPrivateOrShouldBeSerialized(produceHeaderKlib: Boo
             // Always keep private interfaces as they can be part of public type hierarchies.
             || (this as? FirClass)?.isInterface == true
 }
+
+fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        > FirAnnotationContainer.serializeAnnotations(
+    session: FirSession,
+    additionalMetadataProvider: FirAdditionalMetadataProvider?,
+    annotationSerializer: FirAnnotationSerializer,
+    proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+    extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?,
+) {
+    if (extension == null) return
+    allRequiredAnnotations(session, additionalMetadataProvider).serializeAnnotations(annotationSerializer, proto, extension)
+}
+
+fun FirAnnotationContainer.allRequiredAnnotations(
+    session: FirSession,
+    additionalMetadataProvider: FirAdditionalMetadataProvider?,
+): List<FirAnnotation> {
+    val nonSourceAnnotations = nonSourceAnnotations(session)
+    return if (this is FirDeclaration && additionalMetadataProvider != null) {
+        nonSourceAnnotations + additionalMetadataProvider.findGeneratedAnnotationsFor(this)
+    } else {
+        nonSourceAnnotations
+    }
+}
+
+fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        > List<FirAnnotation>.serializeAnnotations(
+    annotationSerializer: FirAnnotationSerializer,
+    proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+    extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?,
+) {
+    if (extension == null) return
+    for (annotation in this) {
+        proto.addExtensionOrNull(extension, annotationSerializer.serializeAnnotation(annotation))
+    }
+}
+
+fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        Type,
+        > GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>.addExtensionOrNull(
+    extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<Type>>,
+    value: Type?,
+) {
+    if (value != null) {
+        addExtension(extension, value)
+    }
+}

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

@@ -7,11 +7,23 @@ package org.jetbrains.kotlin.fir.backend.jvm
 
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmBackendExtension
+import org.jetbrains.kotlin.backend.jvm.ModuleMetadataSerializer
 import org.jetbrains.kotlin.backend.jvm.metadata.MetadataSerializer
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings
 import org.jetbrains.kotlin.fir.backend.Fir2IrComponents
+import org.jetbrains.kotlin.fir.backend.FirMetadataSource
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.serialization.FirElementAwareStringTable
+import org.jetbrains.kotlin.fir.serialization.FirElementSerializer
+import org.jetbrains.kotlin.fir.serialization.TypeApproximatorForMetadataSerializer
+import org.jetbrains.kotlin.fir.serialization.serializeAnnotations
 import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.MetadataSource
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
+import org.jetbrains.kotlin.serialization.StringTableImpl
+import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerializerProtocol
 import org.jetbrains.org.objectweb.asm.Type
 
 class FirJvmBackendExtension(
@@ -35,4 +47,59 @@ class FirJvmBackendExtension(
             actualizedExpectDeclarations
         )
     }
+
+    override fun createModuleMetadataSerializer(context: JvmBackendContext) = object : ModuleMetadataSerializer {
+        override fun serializeOptionalAnnotationClass(metadata: MetadataSource.Class, stringTable: StringTableImpl): ProtoBuf.Class {
+
+            require(metadata is FirMetadataSource.Class) { "Metadata is expected to be ${FirMetadataSource.Class::class.simpleName}" }
+
+            val session = components.session
+            val fir = metadata.fir
+
+            val typeApproximator = TypeApproximatorForMetadataSerializer(session)
+            // Get rid of special serializer extension after KT-57919, i.e. when we serialize all the annotations by default
+            val firSerializerExtension = object : FirJvmSerializerExtension(
+                session,
+                JvmSerializationBindings(),
+                context.state,
+                metadata,
+                // annotation can't have local delegated properties, it is safe to pass empty list
+                localDelegatedProperties = emptyList(),
+                typeApproximator,
+                components,
+                object : FirElementAwareStringTable {
+                    override fun getQualifiedClassNameIndex(className: String, isLocal: Boolean): Int =
+                        stringTable.getQualifiedClassNameIndex(className, isLocal)
+
+                    override fun getStringIndex(string: String): Int = stringTable.getStringIndex(string)
+                }
+            ) {
+                override fun serializeClass(
+                    klass: FirClass,
+                    proto: ProtoBuf.Class.Builder,
+                    versionRequirementTable: MutableVersionRequirementTable,
+                    childSerializer: FirElementSerializer,
+                ) {
+                    klass.serializeAnnotations(
+                        session,
+                        additionalMetadataProvider,
+                        annotationSerializer,
+                        proto,
+                        BuiltInSerializerProtocol.classAnnotation
+                    )
+                    super.serializeClass(klass, proto, versionRequirementTable, childSerializer)
+                }
+            }
+            val serializer = FirElementSerializer.create(
+                session,
+                components.scopeSession,
+                fir,
+                firSerializerExtension,
+                parentSerializer = null,
+                typeApproximator,
+                context.config.languageVersionSettings
+            )
+            return serializer.classProto(fir).build()
+        }
+    }
 }

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.backend.jvm
 
-import org.jetbrains.kotlin.backend.jvm.mapping.IrTypeMapper
 import org.jetbrains.kotlin.codegen.ClassBuilderMode
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings
 import org.jetbrains.kotlin.codegen.serialization.JvmSignatureSerializer
@@ -33,7 +32,6 @@ import org.jetbrains.kotlin.fir.serialization.FirElementSerializer
 import org.jetbrains.kotlin.fir.serialization.FirSerializerExtension
 import org.jetbrains.kotlin.fir.serialization.constant.ConstValueProvider
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.ir.declarations.IrAttributeContainer
 import org.jetbrains.kotlin.ir.declarations.MetadataSource
 import org.jetbrains.kotlin.load.kotlin.NON_EXISTENT_CLASS_NAME
 import org.jetbrains.kotlin.metadata.ProtoBuf
@@ -51,7 +49,7 @@ import org.jetbrains.kotlin.types.AbstractTypeApproximator
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.Method
 
-class FirJvmSerializerExtension(
+open class FirJvmSerializerExtension(
     override val session: FirSession,
     private val bindings: JvmSerializationBindings,
     private val metadata: MetadataSource?,
@@ -66,7 +64,7 @@ class FirJvmSerializerExtension(
     private val unifiedNullChecks: Boolean,
     override val metadataVersion: BinaryVersion,
     private val jvmDefaultMode: JvmDefaultMode,
-    override val stringTable: FirElementAwareStringTable,
+    final override val stringTable: FirElementAwareStringTable,
     override val constValueProvider: ConstValueProvider?,
     override val additionalMetadataProvider: FirAdditionalMetadataProvider?,
 ) : FirSerializerExtension() {
@@ -78,10 +76,9 @@ class FirJvmSerializerExtension(
         state: GenerationState,
         metadata: MetadataSource?,
         localDelegatedProperties: List<FirProperty>,
-        localPoppedUpClasses: List<IrAttributeContainer>,
         approximator: AbstractTypeApproximator,
-        typeMapper: IrTypeMapper,
-        components: Fir2IrComponents
+        components: Fir2IrComponents,
+        stringTable: FirElementAwareStringTable
     ) : this(
         session,
         bindings,
@@ -97,7 +94,7 @@ class FirJvmSerializerExtension(
         state.config.unifiedNullChecks,
         state.config.metadataVersion,
         state.jvmDefaultMode,
-        FirJvmElementAwareStringTable(typeMapper, components, localPoppedUpClasses),
+        stringTable,
         ConstValueProviderImpl(components),
         components.annotationsFromPluginRegistrar.createAdditionalMetadataProvider()
     )

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

@@ -51,8 +51,14 @@ fun makeFirMetadataSerializerForIrClass(
         (it.owner.metadata as? FirMetadataSource.Property)?.fir?.copyToFreeProperty(approximator)
     } ?: emptyList()
     val firSerializerExtension = FirJvmSerializerExtension(
-        session, serializationBindings, context.state, irClass.metadata, localDelegatedProperties, context.isEnclosedInConstructor.toList(),
-        approximator, context.defaultTypeMapper, components
+        session,
+        serializationBindings,
+        context.state,
+        irClass.metadata,
+        localDelegatedProperties,
+        approximator,
+        components,
+        FirJvmElementAwareStringTable(context.defaultTypeMapper, components, context.isEnclosedInConstructor.toList())
     )
     return FirMetadataSerializer(
         context.state.globalSerializationBindings,

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

@@ -7,10 +7,7 @@ package org.jetbrains.kotlin.fir.java.deserialization
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
-import org.jetbrains.kotlin.fir.deserialization.AbstractFirDeserializedSymbolProvider
-import org.jetbrains.kotlin.fir.deserialization.FirDeserializationContext
-import org.jetbrains.kotlin.fir.deserialization.ModuleDataProvider
-import org.jetbrains.kotlin.fir.deserialization.PackagePartsCacheData
+import org.jetbrains.kotlin.fir.deserialization.*
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.metadata.ProtoBuf
@@ -32,6 +29,8 @@ class OptionalAnnotationClassesProvider(
     session, moduleDataProvider, kotlinScopeProvider, defaultDeserializationOrigin, BuiltInSerializerProtocol
 ) {
 
+    private val annotationDeserializer = MetadataBasedAnnotationDeserializer(session)
+
     private val optionalAnnotationClassesAndPackages by lazy(LazyThreadSafetyMode.PUBLICATION) {
         val optionalAnnotationClasses = mutableMapOf<ClassId, ClassData>()
         val optionalAnnotationPackages = mutableSetOf<String>()
@@ -71,7 +70,7 @@ class OptionalAnnotationClassesProvider(
         return ClassMetadataFindResult.Metadata(
             optionalAnnotationClass.nameResolver,
             optionalAnnotationClass.classProto,
-            null,
+            annotationDeserializer,
             moduleDataProvider.allModuleData.last(),
             null,
             classPostProcessor = null

compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt

@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.backend.jvm.ir.getKtFile
 import org.jetbrains.kotlin.backend.jvm.serialization.DisabledIdSignatureDescriptor
 import org.jetbrains.kotlin.backend.jvm.serialization.JvmIdSignatureDescriptor
 import org.jetbrains.kotlin.codegen.CodegenFactory
+import org.jetbrains.kotlin.codegen.addCompiledPartsAndSort
+import org.jetbrains.kotlin.codegen.loadCompiledModule
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
@@ -43,6 +45,7 @@ import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.library.metadata.DeserializedKlibModuleOrigin
 import org.jetbrains.kotlin.library.metadata.KlibModuleOrigin
+import org.jetbrains.kotlin.metadata.jvm.JvmModuleProtoBuf
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi2ir.Psi2IrConfiguration
 import org.jetbrains.kotlin.psi2ir.Psi2IrTranslator
@@ -53,6 +56,7 @@ import org.jetbrains.kotlin.psi2ir.generators.fragments.EvaluatorFragmentInfo
 import org.jetbrains.kotlin.psi2ir.generators.fragments.FragmentContext
 import org.jetbrains.kotlin.psi2ir.preprocessing.SourceDeclarationsPreprocessor
 import org.jetbrains.kotlin.resolve.CleanableBindingContext
+import org.jetbrains.kotlin.serialization.StringTableImpl
 import org.jetbrains.kotlin.utils.IDEAPlatforms
 import org.jetbrains.kotlin.utils.IDEAPluginsCompatibilityAPI
 
@@ -370,6 +374,50 @@ open class JvmIrCodegenFactory(
         // TODO: split classes into groups connected by inline calls; call this after every group
         //       and clear `JvmBackendContext.classCodegens`
         state.afterIndependentPart()
+
+        generateModuleMetadata(input)
+    }
+
+    private fun generateModuleMetadata(result: CodegenFactory.CodegenInput) {
+        val backendContext = (result as JvmIrCodegenInput).context
+        val builder = JvmModuleProtoBuf.Module.newBuilder()
+        val stringTable = StringTableImpl()
+
+        backendContext.state.loadCompiledModule()?.moduleData?.run {
+            // In incremental compilation scenario, we might already have some serialized optionalAnnotations from the previous run
+            // In this case, we first initialize string table with the serialized one
+            // See jps/jps-plugin/testData/incremental/multiModule/multiplatform/custom/modifyOptionalAnnotationUsage for example
+            val nameResolver = nameResolver
+            repeat(nameResolver.strings.stringCount) { stringIndex ->
+                stringTable.addString(nameResolver.strings.getString(stringIndex))
+            }
+            repeat(nameResolver.qualifiedNames.qualifiedNameCount) { nameIndex ->
+                val qualifiedName = nameResolver.qualifiedNames.getQualifiedName(nameIndex)
+                stringTable.addQualifiedName(qualifiedName)
+            }
+            // Then add the annotations themselves, unless they are in dirty sources, i.e. contained in backendContext.optionalAnnotations
+            for (proto in optionalAnnotations) {
+                val name = nameResolver.getQualifiedClassName(proto.fqName)
+                if (backendContext.optionalAnnotations.none { metadata -> metadata.name?.asString() == name }) {
+                    builder.addOptionalAnnotationClass(proto)
+                }
+            }
+        }
+
+        for (part in backendContext.state.factory.packagePartRegistry.parts.values.addCompiledPartsAndSort(backendContext.state)) {
+            part.addTo(builder)
+        }
+
+        for (metadata in backendContext.optionalAnnotations) {
+            val serializer = backendContext.backendExtension.createModuleMetadataSerializer(backendContext)
+            builder.addOptionalAnnotationClass(serializer.serializeOptionalAnnotationClass(metadata, stringTable))
+        }
+
+        val (stringTableProto, qualifiedNameTableProto) = stringTable.buildProto()
+        builder.setStringTable(stringTableProto)
+        builder.setQualifiedNameTable(qualifiedNameTableProto)
+
+        backendContext.state.factory.setModuleMapping(builder.build())
     }
 
     fun generateModuleInFrontendIRMode(

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/JvmLoweringPhases.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.util.PatchDeclarationParentsVisitor
 import org.jetbrains.kotlin.ir.util.isAnonymousObject
+import org.jetbrains.kotlin.ir.util.isExpect
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.util.resolveFakeOverride
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
@@ -102,7 +103,11 @@ private val arrayConstructorPhase = makeIrFilePhase(
 
 internal val expectDeclarationsRemovingPhase = makeIrModulePhase(
     { context: JvmBackendContext ->
-        if (context.config.useFir) FileLoweringPass.Empty
+        if (context.config.useFir) object : FileLoweringPass {
+            override fun lower(irFile: IrFile) {
+                irFile.declarations.removeIf { it.isExpect }
+            }
+        }
         else ExpectDeclarationRemover(context)
     },
     name = "ExpectDeclarationsRemoving",

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/ProcessOptionalAnnotations.kt

@@ -12,15 +12,47 @@ import org.jetbrains.kotlin.backend.jvm.ir.isOptionalAnnotationClass
 import org.jetbrains.kotlin.ir.declarations.DescriptorMetadataSource
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.MetadataSource
+import org.jetbrains.kotlin.ir.util.hasAnnotation
+import org.jetbrains.kotlin.ir.util.isAnnotationClass
 import org.jetbrains.kotlin.resolve.multiplatform.OptionalAnnotationUtil
 
 internal val processOptionalAnnotationsPhase = makeIrModulePhase(
-    ::ProcessOptionalAnnotations,
+    { context: JvmBackendContext ->
+        if (context.config.useFir) ProcessOptionalAnnotationsFir(context) else ProcessOptionalAnnotationsDescriptors(context)
+    },
     name = "ProcessOptionalAnnotations",
     description = "Record metadata of @OptionalExpectation-annotated classes to backend-specific storage, later written to .kotlin_module"
 )
 
-class ProcessOptionalAnnotations(private val context: JvmBackendContext) : FileLoweringPass {
+class ProcessOptionalAnnotationsDescriptors(private val context: JvmBackendContext) : ProcessOptionalAnnotations() {
+    override fun IrClass.processClassFrontendSpecific() {
+        val classMetadata = metadata
+
+        require(classMetadata is DescriptorMetadataSource.Class?) { "IrClass has unexpected metadata: ${classMetadata!!::class.simpleName}" }
+
+        if (classMetadata != null) {
+            val descriptor = classMetadata.descriptor
+            if (OptionalAnnotationUtil.shouldGenerateExpectClass(descriptor)) {
+                context.state.factory.packagePartRegistry.optionalAnnotations += descriptor
+                context.optionalAnnotations += classMetadata
+            }
+        }
+    }
+}
+
+class ProcessOptionalAnnotationsFir(private val context: JvmBackendContext) : ProcessOptionalAnnotations() {
+    override fun IrClass.processClassFrontendSpecific() {
+        val classMetadata = metadata
+        require(classMetadata is MetadataSource.Class?) { "IrClass has unexpected metadata: ${classMetadata!!::class.simpleName}" }
+
+        if (classMetadata != null && isAnnotationClass && isExpect && hasAnnotation(OptionalAnnotationUtil.OPTIONAL_EXPECTATION_FQ_NAME)) {
+            context.optionalAnnotations += classMetadata
+        }
+    }
+}
+
+abstract class ProcessOptionalAnnotations : FileLoweringPass {
     override fun lower(irFile: IrFile) {
         for (declaration in irFile.declarations) {
             if (declaration !is IrClass || !declaration.isOptionalAnnotationClass) continue
@@ -28,14 +60,14 @@ class ProcessOptionalAnnotations(private val context: JvmBackendContext) : FileL
         }
     }
 
+    abstract fun IrClass.processClassFrontendSpecific()
+
     private fun IrClass.registerOptionalAnnotations() {
-        // TODO FirMetadataSource.Class
-        val metadataSource = (metadata as? DescriptorMetadataSource.Class)?.descriptor ?: return
-        if (!OptionalAnnotationUtil.shouldGenerateExpectClass(metadataSource)) return
-        context.state.factory.packagePartRegistry.optionalAnnotations += metadataSource
+        processClassFrontendSpecific()
 
         declarations.forEach {
             if (it is IrClass && it.isOptionalAnnotationClass) it.registerOptionalAnnotations()
         }
     }
+
 }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -193,6 +193,8 @@ class JvmBackendContext(
 
     val visitedDeclarationsForRegenerationLowering: MutableSet<IrDeclaration> = ConcurrentHashMap.newKeySet()
 
+    val optionalAnnotations = mutableListOf<MetadataSource.Class>()
+
     init {
         state.mapInlineClass = { descriptor ->
             defaultTypeMapper.mapType(referenceClass(descriptor).defaultType)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendExtension.kt

@@ -7,10 +7,16 @@ package org.jetbrains.kotlin.backend.jvm
 
 import org.jetbrains.kotlin.backend.jvm.metadata.DescriptorMetadataSerializer
 import org.jetbrains.kotlin.backend.jvm.metadata.MetadataSerializer
+import org.jetbrains.kotlin.codegen.JvmOptionalAnnotationSerializerExtension
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings
 import org.jetbrains.kotlin.config.JvmAbiStability
+import org.jetbrains.kotlin.ir.declarations.DescriptorMetadataSource
 import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.MetadataSource
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.serialization.DescriptorSerializer
+import org.jetbrains.kotlin.serialization.StringTableImpl
 import org.jetbrains.org.objectweb.asm.Type
 
 interface JvmBackendExtension {
@@ -29,8 +35,24 @@ interface JvmBackendExtension {
             return DescriptorMetadataSerializer(context, klass, type, bindings, parentSerializer)
         }
 
+        override fun createModuleMetadataSerializer(context: JvmBackendContext) = object : ModuleMetadataSerializer {
+            override fun serializeOptionalAnnotationClass(metadata: MetadataSource.Class, stringTable: StringTableImpl): ProtoBuf.Class {
+                require(metadata is DescriptorMetadataSource.Class)
+                return DescriptorSerializer.createTopLevel(
+                    JvmOptionalAnnotationSerializerExtension(stringTable), context.state.config.languageVersionSettings,
+                ).classProto(metadata.descriptor).build()
+            }
+        }
+
+
         fun generateMetadataExtraFlags(abiStability: JvmAbiStability?): Int =
             JvmAnnotationNames.METADATA_JVM_IR_FLAG or
                     (if (abiStability != JvmAbiStability.UNSTABLE) JvmAnnotationNames.METADATA_JVM_IR_STABLE_ABI_FLAG else 0)
     }
+
+    fun createModuleMetadataSerializer(context: JvmBackendContext): ModuleMetadataSerializer
+}
+
+interface ModuleMetadataSerializer {
+    fun serializeOptionalAnnotationClass(metadata: MetadataSource.Class, stringTable: StringTableImpl): ProtoBuf.Class
 }

compiler/serialization/src/org/jetbrains/kotlin/serialization/StringTableImpl.kt

@@ -46,6 +46,14 @@ open class SerializableStringTable : StringTable {
     private val strings = Interner<String>()
     private val qualifiedNames = Interner<FqNameProto>()
 
+    fun addString(string: String) {
+        strings.intern(string)
+    }
+
+    fun addQualifiedName(qualifiedName: ProtoBuf.QualifiedNameTable.QualifiedName) {
+        qualifiedNames.intern(FqNameProto(qualifiedName.toBuilder()))
+    }
+
     override fun getStringIndex(string: String): Int = strings.intern(string)
 
     override fun getQualifiedClassNameIndex(className: String, isLocal: Boolean): Int =

compiler/testData/codegen/box/multiplatform/optionalExpectationJvm.jvm_abi.txt

@@ -1,6 +1,4 @@
 MODULE main
-	Missing in K1
-		Anno.class
 	CLASS Foo.class
 		CLASS METADATA
 			PROPERTY getX()I
@@ -9,9 +7,3 @@ MODULE main
 						<set-?>: kotlin/Int
 					K2
 						value: kotlin/Int
-	MODULE METADATA
-		Property: module.metadata.optionalAnnotations
-			K1
-				[Anno]
-			K2
-				[]

compiler/testData/codegen/box/multiplatform/optionalExpectationJvm.kt

@@ -2,7 +2,7 @@
 // !OPT_IN: kotlin.ExperimentalMultiplatform
 // TARGET_BACKEND: JVM
 // WITH_STDLIB
-// JVM_ABI_K1_K2_DIFF: KT-62931, KT-63984
+// JVM_ABI_K1_K2_DIFF: KT-63984
 
 // FILE: common.kt
 

compiler/testData/codegen/bytecodeListing/multiplatform/optionalExpectation.jvm_abi.txt

@@ -0,0 +1,9 @@
+MODULE main
+	CLASS Foo.class
+		CLASS METADATA
+			PROPERTY getX()I
+				Property: class.metadata.property.setterValueParameter
+					K1
+						<set-?>: kotlin/Int
+					K2
+						value: kotlin/Int

compiler/testData/codegen/bytecodeListing/multiplatform/optionalExpectation.kt

@@ -2,9 +2,8 @@
 // !OPT_IN: kotlin.ExperimentalMultiplatform
 // TARGET_BACKEND: JVM
 // WITH_STDLIB
+// JVM_ABI_K1_K2_DIFF: KT-63984
 
-// IGNORE_BACKEND_K2: JVM_IR
-// FIR status: K2 incorrectly generates Anno.class, need to investigate after KT-57243 is fixed.
 
 @file:Suppress("OPTIONAL_DECLARATION_USAGE_IN_NON_COMMON_SOURCE") // TODO: support common sources in the test infrastructure
 

core/metadata.jvm/src/org/jetbrains/kotlin/metadata/jvm/deserialization/BinaryModuleData.kt

@@ -6,7 +6,7 @@
 package org.jetbrains.kotlin.metadata.jvm.deserialization
 
 import org.jetbrains.kotlin.metadata.ProtoBuf
-import org.jetbrains.kotlin.metadata.deserialization.NameResolver
+import org.jetbrains.kotlin.metadata.deserialization.NameResolverImpl
 
 /**
  * @param annotations list of module annotations, in the format: "org/foo/bar/Baz.Inner" (see [ClassId.fromString])
@@ -16,5 +16,5 @@ import org.jetbrains.kotlin.metadata.deserialization.NameResolver
 class BinaryModuleData(
     val annotations: List<String>,
     val optionalAnnotations: List<ProtoBuf.Class>,
-    val nameResolver: NameResolver,
+    val nameResolver: NameResolverImpl,
 )