Support .kotlin_module files in kotlinx-metadata

#KT-23198

libraries/kotlinx-metadata/jvm/src/kotlinx/metadata/jvm/KotlinModuleMetadata.kt

@@ -0,0 +1,152 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package kotlinx.metadata.jvm
+
+import kotlinx.metadata.KmAnnotation
+import kotlinx.metadata.ClassName
+import kotlinx.metadata.InconsistentKotlinMetadataException
+import org.jetbrains.kotlin.metadata.jvm.JvmModuleProtoBuf
+import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping
+import org.jetbrains.kotlin.metadata.jvm.deserialization.PackageParts
+import org.jetbrains.kotlin.metadata.jvm.deserialization.serializeToByteArray
+
+/**
+ * Represents the parsed metadata of a Kotlin JVM module file.
+ *
+ * To create an instance of [KotlinModuleMetadata], load the contents of the `.kotlin_module` file into a byte array
+ * and call [KotlinModuleMetadata.read].
+ *
+ * @property bytes the byte array representing the contents of a `.kotlin_module` file
+ */
+class KotlinModuleMetadata(@Suppress("CanBeParameter", "MemberVisibilityCanBePrivate") val bytes: ByteArray) {
+    internal val data: ModuleMapping = ModuleMapping.loadModuleMapping(bytes, javaClass.name, isVersionCompatible = { version ->
+        // We only support metadata of version 1.1.* (this is Kotlin from 1.0 until today)
+        version.getOrNull(0) == 1 && version.getOrNull(1) == 1
+    }, skipMetadataVersionCheck = false, isJvmPackageNameSupported = true)
+
+    /**
+     * A [KmModuleVisitor] that generates the metadata of a Kotlin JVM module file.
+     */
+    class Writer : KmModuleVisitor() {
+        private val b = JvmModuleProtoBuf.Module.newBuilder()
+
+        override fun visitPackageParts(fqName: String, fileFacades: List<String>, multiFileClassParts: Map<String, String>) {
+            PackageParts(fqName).apply {
+                for (fileFacade in fileFacades) {
+                    addPart(fileFacade, null)
+                }
+                for ((multiFileClassPart, multiFileFacade) in multiFileClassParts) {
+                    addPart(multiFileClassPart, multiFileFacade)
+                }
+
+                addTo(b)
+            }
+        }
+
+        override fun visitAnnotation(annotation: KmAnnotation) {
+            /*
+            // TODO: move StringTableImpl to module 'metadata' and support module annotations here
+            b.addAnnotation(ProtoBuf.Annotation.newBuilder().apply {
+                id = annotation.className.name // <-- use StringTableImpl here
+            })
+            */
+        }
+
+        /**
+         * Returns the metadata of the module file that was written with this writer.
+         *
+         * @param metadataVersion metadata version to be written to the metadata (see [KotlinClassHeader.metadataVersion]),
+         *   [KotlinClassHeader.COMPATIBLE_METADATA_VERSION] by default
+         */
+        fun write(metadataVersion: IntArray = KotlinClassHeader.COMPATIBLE_METADATA_VERSION): KotlinModuleMetadata =
+            KotlinModuleMetadata(b.build().serializeToByteArray(metadataVersion))
+    }
+
+    /**
+     * Makes the given visitor visit metadata of this module file.
+     *
+     * @param v the visitor that must visit this module file
+     */
+    fun accept(v: KmModuleVisitor) {
+        for ((fqName, parts) in data.packageFqName2Parts) {
+            val (fileFacades, multiFileClassParts) = parts.parts.partition { parts.getMultifileFacadeName(it) == null }
+            v.visitPackageParts(fqName, fileFacades, multiFileClassParts.associate { it to parts.getMultifileFacadeName(it)!! })
+        }
+
+        for (annotation in data.moduleData.annotations) {
+            v.visitAnnotation(KmAnnotation(annotation, emptyMap()))
+        }
+
+        v.visitEnd()
+    }
+
+    companion object {
+        /**
+         * Parses the given byte array with the .kotlin_module file content and returns the [KotlinModuleMetadata] instance,
+         * or `null` if this byte array encodes a module with an unsupported metadata version.
+         *
+         * Throws [InconsistentKotlinMetadataException] if an error happened while parsing the given byte array,
+         * which means that it's either not the content of a .kotlin_module file, or it has been corrupted.
+         */
+        @JvmStatic
+        fun read(bytes: ByteArray): KotlinModuleMetadata? {
+            try {
+                val result = KotlinModuleMetadata(bytes)
+                if (result.data == ModuleMapping.EMPTY) return null
+
+                if (result.data == ModuleMapping.CORRUPTED) {
+                    throw InconsistentKotlinMetadataException("Data doesn't look like the content of a .kotlin_module file")
+                }
+
+                return result
+            } catch (e: InconsistentKotlinMetadataException) {
+                throw e
+            } catch (e: Throwable) {
+                throw InconsistentKotlinMetadataException("Exception occurred when reading Kotlin metadata", e)
+            }
+        }
+    }
+}
+
+/**
+ * A visitor to visit Kotlin JVM module files.
+ *
+ * When using this class, [visitEnd] must be called exactly once and after calls to all other visit* methods.
+ */
+abstract class KmModuleVisitor(private val delegate: KmModuleVisitor? = null) {
+    /**
+     * Visits the table of single- and multi-file facades declared in some package of this module. Note that the same package may be split
+     * in several "chunks", in which case this method is called several times for one package, once for each chunk.
+     *
+     * Packages are separated by '/' in the names of file facades.
+     *
+     * @param fqName the fully qualified name of the package
+     * @param fileFacades the list of single-file facades in this chunk of the package
+     * @param multiFileClassParts the map of multi-file classes where keys are names of multi-file class parts,
+     *   and values are names of the corresponding multi-file facades
+     */
+    open fun visitPackageParts(fqName: String, fileFacades: List<String>, multiFileClassParts: Map<String, String>) {
+        delegate?.visitPackageParts(fqName, fileFacades, multiFileClassParts)
+    }
+
+    /**
+     * Visits the annotation on the module.
+     *
+     * @param annotation annotation on the module
+     */
+    open fun visitAnnotation(annotation: KmAnnotation) {
+        delegate?.visitAnnotation(annotation)
+    }
+
+    /**
+     * Visits the end of the module.
+     */
+    open fun visitEnd() {
+        delegate?.visitEnd()
+    }
+
+    // TODO: JvmPackageName
+}

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