[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/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
+}