JVM IR: generate correct metadata for multi-file classes

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

@@ -307,27 +307,17 @@ class MultifileClassCodegenImpl(
 
         val extraFlags = if (shouldGeneratePartHierarchy) JvmAnnotationNames.METADATA_MULTIFILE_PARTS_INHERIT_FLAG else 0
 
-        writeKotlinMetadata(classBuilder, state, KotlinClassHeader.Kind.MULTIFILE_CLASS, extraFlags) { av ->
+        val kotlinPackageFqName =
-            val arv = av.visitArray(JvmAnnotationNames.METADATA_DATA_FIELD_NAME)
+            packageFragment?.fqName ?: compiledPackageFragment?.fqName
-            for (internalName in partInternalNamesSorted) {
+            ?: error("Either source package or compiled package should not be null: $facadeClassType ($files)")
-                arv.visit(null, internalName)
-            }
-            arv.visitEnd()
 
-            val kotlinPackageFqName =
+        if (files.any { it.packageFqName != kotlinPackageFqName })
-                packageFragment?.fqName ?: compiledPackageFragment?.fqName
+            throw UnsupportedOperationException(
-                ?: error("Either source package or compiled package should not be null: $facadeClassType ($files)")
+                "Multi-file parts of a facade with JvmPackageName should all lie in the same Kotlin package:\n  " +
+                        files.joinToString("\n  ") { file -> "$file: package ${file.packageFqName}" }
+            )
 
-            if (files.any { it.packageFqName != kotlinPackageFqName })
+        writeMetadata(classBuilder, state, extraFlags, partInternalNamesSorted, facadeClassType, kotlinPackageFqName)
-                throw UnsupportedOperationException(
-                    "Multi-file parts of a facade with JvmPackageName should all lie in the same Kotlin package:\n  " +
-                            files.joinToString("\n  ") { file -> "$file: package ${file.packageFqName}" }
-                )
-
-            if (kotlinPackageFqName != JvmClassName.byInternalName(facadeClassType.internalName).packageFqName) {
-                av.visit(JvmAnnotationNames.METADATA_PACKAGE_NAME_FIELD_NAME, kotlinPackageFqName.asString())
-            }
-        }
     }
 
     private fun createCodegenForDelegatesInMultifileFacade(facadeContext: FieldOwnerContext<*>): MemberCodegen<KtFile> =
@@ -375,5 +365,26 @@ class MultifileClassCodegenImpl(
 
             return incrementalPackageFragment?.getPackageFragmentForMultifileClass(facadeFqName)
         }
+
+        fun writeMetadata(
+            classBuilder: ClassBuilder,
+            state: GenerationState,
+            flags: Int,
+            partInternalNames: List<String>,
+            facadeClassType: Type,
+            kotlinPackageFqName: FqName
+        ) {
+            writeKotlinMetadata(classBuilder, state, KotlinClassHeader.Kind.MULTIFILE_CLASS, flags) { av ->
+                val arv = av.visitArray(JvmAnnotationNames.METADATA_DATA_FIELD_NAME)
+                for (internalName in partInternalNames) {
+                    arv.visit(null, internalName)
+                }
+                arv.visitEnd()
+
+                if (kotlinPackageFqName != JvmClassName.byInternalName(facadeClassType.internalName).packageFqName) {
+                    av.visit(JvmAnnotationNames.METADATA_PACKAGE_NAME_FIELD_NAME, kotlinPackageFqName.asString())
+                }
+            }
+        }
     }
 }

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

@@ -65,6 +65,7 @@ class JvmBackendContext(
     internal val localDelegatedProperties = mutableMapOf<IrClass, List<IrLocalDelegatedPropertySymbol>>()
 
     internal val multifileFacadesToAdd = mutableMapOf<JvmClassName, MutableList<IrClass>>()
+    internal val multifileFacadeForPart = mutableMapOf<IrClass, JvmClassName>()
 
     override var inVerbosePhase: Boolean = false
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.load.java.JvmAbi
+import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
@@ -144,13 +145,31 @@ open class ClassCodegen protected constructor(
 
                 serializerExtension.serializeJvmPackage(packageProto, type)
 
-                writeKotlinMetadata(visitor, state, KotlinClassHeader.Kind.FILE_FACADE, 0) {
+                val facadeClassName = context.multifileFacadeForPart[irClass.attributeOwnerId]
-                    AsmUtil.writeAnnotationData(it, serializer, packageProto.build())
+                val kind = if (facadeClassName != null) KotlinClassHeader.Kind.MULTIFILE_CLASS_PART else KotlinClassHeader.Kind.FILE_FACADE
+                writeKotlinMetadata(visitor, state, kind, 0) { av ->
+                    AsmUtil.writeAnnotationData(av, serializer, packageProto.build())
+
+                    if (facadeClassName != null) {
+                        av.visit(JvmAnnotationNames.METADATA_MULTIFILE_CLASS_NAME_FIELD_NAME, facadeClassName.internalName)
+                    }
+
                     // TODO: JvmPackageName
                 }
             }
             else -> {
-                writeSyntheticClassMetadata(visitor, state)
+                val entry = irClass.fileParent.fileEntry
+                if (entry is MultifileFacadeFileEntry) {
+                    val partInternalNames = entry.partFiles.mapNotNull { partFile ->
+                        val fileClass = partFile.declarations.singleOrNull { it.origin == IrDeclarationOrigin.FILE_CLASS } as IrClass?
+                        if (fileClass != null) typeMapper.mapClass(fileClass).internalName else null
+                    }
+                    MultifileClassCodegenImpl.writeMetadata(
+                        visitor, state, 0 /* TODO */, partInternalNames, type, irClass.fqNameWhenAvailable!!.parent()
+                    )
+                } else {
+                    writeSyntheticClassMetadata(visitor, state)
+                }
             }
         }
     }

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

@@ -85,6 +85,8 @@ private fun generateMultifileFacades(module: ModuleDescriptor, context: JvmBacke
         file.declarations.add(facadeClass)
 
         for (partClass in partClasses) {
+            context.multifileFacadeForPart[partClass.attributeOwnerId as IrClass] = jvmClassName
+
             for (member in partClass.declarations) {
                 member.createMultifileDelegateIfNeeded(context, facadeClass)
             }

jps-plugin/testData/incremental/withJava/other/multifileClassAddTopLevelFunWithDefault/directives.txt

@@ -1 +0,0 @@
-// IGNORE_BACKEND: JVM_IR

jps-plugin/testData/incremental/withJava/other/multifileClassFileAdded/directives.txt

@@ -1 +0,0 @@
-// IGNORE_BACKEND: JVM_IR

jps-plugin/testData/incremental/withJava/other/multifileClassFileChanged/directives.txt

@@ -1 +0,0 @@
-// IGNORE_BACKEND: JVM_IR

jps-plugin/testData/incremental/withJava/other/multifilePackagePartMethodAdded/directives.txt

@@ -1 +0,0 @@
-// IGNORE_BACKEND: JVM_IR

jps-plugin/testData/incremental/withJava/other/multifilePartsWithProperties/directives.txt

@@ -1 +0,0 @@
-// IGNORE_BACKEND: JVM_IR