Support .kotlin_module files in kotlinx-metadata

#KT-23198

libraries/tools/kotlinp/src/org/jetbrains/kotlin/kotlinp/Kotlinp.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.kotlinp
 
 import kotlinx.metadata.InconsistentKotlinMetadataException
 import kotlinx.metadata.jvm.KotlinClassMetadata
+import kotlinx.metadata.jvm.KotlinModuleMetadata
 import java.io.File
 
 object Kotlinp {
@@ -32,4 +33,11 @@ object Kotlinp {
             throw KotlinpException("inconsistent Kotlin metadata: ${e.message}")
         }
     }
+
+    internal fun renderModuleFile(metadata: KotlinModuleMetadata?): String =
+        if (metadata != null) ModuleFilePrinter().print(metadata)
+        else buildString { appendln("unsupported file") }
+
+    internal fun readModuleFile(file: File): KotlinModuleMetadata? =
+        KotlinModuleMetadata.read(file.readBytes())
 }

libraries/tools/kotlinp/src/org/jetbrains/kotlin/kotlinp/Main.kt

@@ -28,16 +28,20 @@ object Main {
         }
 
         for (path in paths) {
-            if (!path.endsWith(".class")) throw KotlinpException("only .class files are supported")
-
             val file = File(path)
             if (!file.exists()) throw KotlinpException("file does not exist: $path")
 
-            try {
-                print(Kotlinp.renderClassFile(Kotlinp.readClassFile(file)))
+            val text = try {
+                when (file.extension) {
+                    "class" -> Kotlinp.renderClassFile(Kotlinp.readClassFile(file))
+                    "kotlin_module" -> Kotlinp.renderModuleFile(Kotlinp.readModuleFile(file))
+                    else -> throw KotlinpException("only .class and .kotlin_module files are supported")
+                }
             } catch (e: IOException) {
                 throw KotlinpException("I/O operation failed: ${e.message}")
             }
+
+            print(text)
         }
 
         if (paths.isEmpty()) {

libraries/tools/kotlinp/src/org/jetbrains/kotlin/kotlinp/printers.kt

@@ -647,6 +647,37 @@ class MultiFileClassFacadePrinter : AbstractPrinter<KotlinClassMetadata.MultiFil
         }
 }
 
+class ModuleFilePrinter : KmModuleVisitor() {
+    private val sb = StringBuilder().apply {
+        appendln("module {")
+    }
+
+    override fun visitPackageParts(fqName: String, fileFacades: List<String>, multiFileClassParts: Map<String, String>) {
+        val presentableFqName = if (fqName.isEmpty()) "<root>" else fqName
+        sb.appendln("  package $presentableFqName {")
+        for (fileFacade in fileFacades) {
+            sb.appendln("    $fileFacade")
+        }
+        for ((multiFileClassPart, facade) in multiFileClassParts) {
+            sb.appendln("    $multiFileClassPart ($facade)")
+        }
+        sb.appendln("  }")
+    }
+
+    override fun visitAnnotation(annotation: KmAnnotation) {
+        // TODO
+    }
+
+    override fun visitEnd() {
+        sb.appendln("}")
+    }
+
+    fun print(metadata: KotlinModuleMetadata): String {
+        metadata.accept(this)
+        return sb.toString()
+    }
+}
+
 private val VISIBILITY_FLAGS_MAP = mapOf(
     Flags.IS_INTERNAL to "internal",
     Flags.IS_PRIVATE to "private",

libraries/tools/kotlinp/test/org/jetbrains/kotlin/kotlinp/test/KotlinpTestUtils.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.kotlinp.test
 import com.intellij.openapi.Disposable
 import junit.framework.TestCase.assertEquals
 import kotlinx.metadata.jvm.KotlinClassMetadata
+import kotlinx.metadata.jvm.KotlinModuleMetadata
 import org.jetbrains.kotlin.checkers.setupLanguageVersionSettingsForCompilerTests
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
@@ -26,7 +27,13 @@ fun compileAndPrintAllFiles(file: File, disposable: Disposable, tmpdir: File, co
     compile(file, disposable, tmpdir) { outputFile ->
         when (outputFile.extension) {
             "kotlin_module" -> {
-                // TODO: support kotlin_module files
+                val moduleFile = Kotlinp.readModuleFile(outputFile)!!
+                val moduleFile2 = transformModuleFile(moduleFile)
+
+                for ((sb, moduleFileToRender) in listOf(read to moduleFile, readWriteRead to moduleFile2)) {
+                    sb.appendFileName(outputFile.relativeTo(tmpdir))
+                    sb.append(Kotlinp.renderModuleFile(moduleFileToRender))
+                }
             }
             "class" -> {
                 val classFile = Kotlinp.readClassFile(outputFile)!!
@@ -91,3 +98,6 @@ private fun transformClassFile(classFile: KotlinClassMetadata): KotlinClassMetad
             KotlinClassMetadata.MultiFileClassPart.Writer().apply(classFile::accept).write(classFile.facadeClassName)
         else -> classFile
     }
+
+private fun transformModuleFile(moduleFile: KotlinModuleMetadata): KotlinModuleMetadata =
+    KotlinModuleMetadata.Writer().apply(moduleFile::accept).write()

libraries/tools/kotlinp/testData/Lambda.txt

@@ -19,3 +19,10 @@ lambda {
   // signature: invoke(Lkotlin/Unit;[ILjava/util/Set;)Ljava/lang/String;
   local final fun kotlin/Unit.<no name provided>(s: kotlin/IntArray?, t: kotlin/collections/Set<kotlin/Double>): kotlin/String
 }
+// META-INF/test-module.kotlin_module
+// ------------------------------------------
+module {
+  package <root> {
+    LambdaKt
+  }
+}

libraries/tools/kotlinp/testData/MultiFileClass.txt

@@ -12,3 +12,10 @@ package {
   public final const val x: kotlin/Int /* = ... */
     public final get
 }
+// META-INF/test-module.kotlin_module
+// ------------------------------------------
+module {
+  package test {
+    test/File__MultiFileClassKt (test/File)
+  }
+}

libraries/tools/kotlinp/testData/SimplePackage.txt

@@ -16,3 +16,10 @@ package {
 
   public typealias G<T#0 /* S */> = F^<kotlin/collections/List<T#0>, kotlin/collections/Set<T#0>> /* = kotlin/collections/Map<kotlin/collections/List<T#0>, kotlin/Function1<java/lang/StringBuilder /* = kotlin/text/StringBuilder^ */, kotlin/collections/Set<T#0>?>> */
 }
+// META-INF/test-module.kotlin_module
+// ------------------------------------------
+module {
+  package <root> {
+    SimplePackageKt
+  }
+}

libraries/tools/kotlinp/testData/SyntheticClass.txt

@@ -8,3 +8,10 @@ package {
 // SyntheticClassKt$WhenMappings.class
 // ------------------------------------------
 synthetic class
+// META-INF/test-module.kotlin_module
+// ------------------------------------------
+module {
+  package <root> {
+    SyntheticClassKt
+  }
+}

libraries/tools/kotlinp/testData/VersionRequirement.txt

@@ -32,3 +32,10 @@ package {
   // requires API version 1.1.0 (level=HIDDEN, errorCode=314)
   public typealias Typealias = kotlin/String /* = kotlin/String */
 }
+// META-INF/test-module.kotlin_module
+// ------------------------------------------
+module {
+  package test {
+    test/VersionRequirementKt
+  }
+}