kotlin-reflect: split K2/K1 branches in KotlinModuleShadowTransformer

libraries/reflect/build.gradle.kts

@@ -60,7 +60,7 @@ dependencies {
 }
 
 @CacheableTransformer
-class KotlinModuleShadowTransformer(private val logger: Logger) : Transformer {
+class KotlinModuleShadowTransformer(private val logger: Logger, private val useK2: Boolean) : Transformer {
     @Suppress("ArrayInDataClass")
     private data class Entry(val path: String, val bytes: ByteArray)
 
@@ -75,23 +75,36 @@ class KotlinModuleShadowTransformer(private val logger: Logger) : Transformer {
         fun relocate(content: String): String =
             context.relocators.fold(content) { acc, relocator -> relocator.applyToSourceContent(acc) }
 
-        val internalData = org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping.loadModuleMapping(
-            context.`is`.readBytes(), javaClass.name, skipMetadataVersionCheck = true, isJvmPackageNameSupported = true
-        ) {
-        }
         val writer = KotlinModuleMetadata.Writer()
         logger.info("Transforming ${context.path}")
-        val visitor = object : KmModuleVisitor(writer) {
+        if (useK2) {
-            override fun visitPackageParts(fqName: String, fileFacades: List<String>, multiFileClassParts: Map<String, String>) {
+            // TODO: remove this branch after migration to version 1.9
-                assert(multiFileClassParts.isEmpty()) { multiFileClassParts } // There are no multi-file class parts in core
+            val internalData = org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping.loadModuleMapping(
-                super.visitPackageParts(relocate(fqName), fileFacades.map(::relocate), multiFileClassParts)
+                context.`is`.readBytes(), javaClass.name, skipMetadataVersionCheck = true, isJvmPackageNameSupported = true
+            ) {
             }
+            val visitor = object : KmModuleVisitor(writer) {
+                override fun visitPackageParts(fqName: String, fileFacades: List<String>, multiFileClassParts: Map<String, String>) {
+                    assert(multiFileClassParts.isEmpty()) { multiFileClassParts } // There are no multi-file class parts in core
+                    super.visitPackageParts(relocate(fqName), fileFacades.map(::relocate), multiFileClassParts)
+                }
+            }
+            for ((fqName, parts) in internalData.packageFqName2Parts) {
+                val (fileFacades, multiFileClassParts) = parts.parts.partition { parts.getMultifileFacadeName(it) == null }
+                visitor.visitPackageParts(fqName, fileFacades, multiFileClassParts.associateWith { parts.getMultifileFacadeName(it)!! })
+            }
+            visitor.visitEnd()
+        } else {
+            val metadata = KotlinModuleMetadata.read(context.`is`.readBytes())
+                ?: error("Not a .kotlin_module file: ${context.path}")
+            // TODO: writer declaration and logger.info call from above should be move here after migration to version 1.9
+            metadata.accept(object : KmModuleVisitor(writer) {
+                override fun visitPackageParts(fqName: String, fileFacades: List<String>, multiFileClassParts: Map<String, String>) {
+                    assert(multiFileClassParts.isEmpty()) { multiFileClassParts } // There are no multi-file class parts in core
+                    super.visitPackageParts(relocate(fqName), fileFacades.map(::relocate), multiFileClassParts)
+                }
+            })
         }
-        for ((fqName, parts) in internalData.packageFqName2Parts) {
-            val (fileFacades, multiFileClassParts) = parts.parts.partition { parts.getMultifileFacadeName(it) == null }
-            visitor.visitPackageParts(fqName, fileFacades, multiFileClassParts.associateWith { parts.getMultifileFacadeName(it)!! })
-        }
-        visitor.visitEnd()
 
         data += Entry(context.path, writer.write().bytes)
     }
@@ -120,7 +133,7 @@ val reflectShadowJar by task<ShadowJar> {
 
     if (kotlinBuildProperties.relocation) {
         mergeServiceFiles()
-        transform(KotlinModuleShadowTransformer(logger))
+        transform(KotlinModuleShadowTransformer(logger, project.kotlinBuildProperties.useFir))
         relocate("org.jetbrains.kotlin", "kotlin.reflect.jvm.internal.impl")
         relocate("javax.inject", "kotlin.reflect.jvm.internal.impl.javax.inject")
     }