Rework KotlinModuleMetadata to align it with KotlinClassMetadata:

- Public constructor with KmModule and JvmMetadataVersion
- make write() member function

libraries/kotlinx-metadata/jvm/ReadMe.md

@@ -202,11 +202,11 @@ The only difference is that the source for the reader (and the result of the wri
 ```kotlin
 // Read the module metadata
 val bytes = File("META-INF/main.kotlin_module").readBytes()
-val metadata = KotlinModuleMetadata.read(bytes)
+val moduleMetadata = KotlinModuleMetadata.read(bytes)
-val module = metadata.kmModule
+val module = moduleMetadata.kmModule
 ...
 
 // Write the module metadata
-val bytes = KotlinModuleMetadata.write(module)
+val bytes = moduleMetadata.write()
 File("META-INF/main.kotlin_module").writeBytes(bytes)
 ```

libraries/kotlinx-metadata/jvm/api/kotlinx-metadata-jvm.api

@@ -1638,11 +1638,15 @@ public final class kotlinx/metadata/jvm/KotlinClassMetadata$Unknown$Companion {
 
 public final class kotlinx/metadata/jvm/KotlinModuleMetadata {
 	public static final field Companion Lkotlinx/metadata/jvm/KotlinModuleMetadata$Companion;
-	public synthetic fun <init> ([BLkotlinx/metadata/internal/metadata/jvm/deserialization/ModuleMapping;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Lkotlinx/metadata/jvm/KmModule;Lkotlinx/metadata/jvm/JvmMetadataVersion;)V
 	public final fun accept (Lkotlinx/metadata/jvm/KmModuleVisitor;)V
 	public final fun getKmModule ()Lkotlinx/metadata/jvm/KmModule;
+	public final fun getVersion ()Lkotlinx/metadata/jvm/JvmMetadataVersion;
 	public static final fun read ([B)Lkotlinx/metadata/jvm/KotlinModuleMetadata;
+	public final fun setKmModule (Lkotlinx/metadata/jvm/KmModule;)V
+	public final fun setVersion (Lkotlinx/metadata/jvm/JvmMetadataVersion;)V
 	public final fun toKmModule ()Lkotlinx/metadata/jvm/KmModule;
+	public final fun write ()[B
 	public static final fun write (Lkotlinx/metadata/jvm/KmModule;)[B
 	public static final fun write (Lkotlinx/metadata/jvm/KmModule;Lkotlinx/metadata/jvm/JvmMetadataVersion;)[B
 }

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

@@ -4,27 +4,24 @@
  */
 
 @file:Suppress(
-    "MemberVisibilityCanBePrivate",
+    "DEPRECATION", // COMPATIBLE_METADATA_VERSION
-    "DEPRECATION", // For KmModule.annotations — remove reading when deprecation is raised to error
+    "DEPRECATION_ERROR", // KmModule.annotations
-    "DEPRECATION_ERROR", // deprecated .accept implementation
-    "INVISIBLE_MEMBER", // InconsistentKotlinMetadataException
-    "INVISIBLE_REFERENCE",
     "UNUSED_PARAMETER" // For deprecated Writer.write
 )
 
 package kotlinx.metadata.jvm
 
 import kotlinx.metadata.*
-import kotlinx.metadata.internal.toKmClass
 import kotlinx.metadata.jvm.KotlinClassMetadata.Companion.COMPATIBLE_METADATA_VERSION
+import kotlinx.metadata.jvm.internal.JvmReadUtils.readModuleMetadataImpl
 import kotlinx.metadata.jvm.internal.JvmReadUtils.throwIfNotCompatible
 import kotlinx.metadata.jvm.internal.wrapIntoMetadataExceptionWhenNeeded
 import kotlinx.metadata.jvm.internal.wrapWriteIntoIAE
 import org.jetbrains.kotlin.metadata.jvm.JvmModuleProtoBuf
-import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMetadataVersion as CompilerMetadataVersion
 import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping
 import org.jetbrains.kotlin.metadata.jvm.deserialization.PackageParts
 import org.jetbrains.kotlin.metadata.jvm.deserialization.serializeToByteArray
+import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMetadataVersion as CompilerMetadataVersion
 
 /**
  * Represents the parsed metadata of a Kotlin JVM module file.
@@ -34,22 +31,21 @@ import org.jetbrains.kotlin.metadata.jvm.deserialization.serializeToByteArray
  *
  * `.kotlin_module` file is produced per Kotlin compilation, and contains auxiliary information, such as a map of all single- and multi-file facades ([KmModule.packageParts]),
  *  `@OptionalExpectation` declarations ([KmModule.optionalAnnotationClasses]), and module annotations ([KmModule.annotations).
- *
- * @property bytes the byte array representing the contents of a `.kotlin_module` file
  */
 @UnstableMetadataApi
-public class KotlinModuleMetadata private constructor(
+public class KotlinModuleMetadata public constructor(
-    private val bytes: ByteArray,
-    private val data: ModuleMapping,
-) {
-
     /**
-     * Returns the [KmModule] representation of this metadata.
+     * [KmModule] representation of this metadata.
      *
      * Returns the same (mutable) [KmModule] instance every time.
      */
-    public val kmModule: KmModule = readImpl()
+    public var kmModule: KmModule,
 
+    /**
+     * Version of this metadata.
+     */
+    public var version: JvmMetadataVersion,
+) {
     /**
      * Visits metadata of this module with a new [KmModule] instance and returns that instance.
      */
@@ -60,6 +56,49 @@ public class KotlinModuleMetadata private constructor(
     )
     public fun toKmModule(): KmModule = KmModule().apply { kmModule.accept(this) }
 
+    /**
+     * Encodes and writes this metadata of the Kotlin module file.
+     *
+     * This method encodes all available data, including [version].
+     *
+     * @throws IllegalArgumentException if [kmModule] is not correct and cannot be written or if [version] is not supported for writing.
+     */
+    public fun write(): ByteArray {
+        val b = JvmModuleProtoBuf.Module.newBuilder()
+        kmModule.packageParts.forEach { (fqName, packageParts) ->
+            PackageParts(fqName).apply {
+                for (fileFacade in packageParts.fileFacades) {
+                    addPart(fileFacade, null)
+                }
+                for ((multiFileClassPart, multiFileFacade) in packageParts.multiFileClassParts) {
+                    addPart(multiFileClassPart, multiFileFacade)
+                }
+
+                addTo(b)
+            }
+        }
+
+        // visitAnnotation
+        /*
+        // TODO: move StringTableImpl to module 'metadata' and support module annotations here
+        b.addAnnotation(ProtoBuf.Annotation.newBuilder().apply {
+            id = annotation.className.name // <-- use StringTableImpl here
+        })
+        */
+
+        // visitOptionalAnnotationClass
+        /*
+        return object : ClassWriter(TODO() /* use StringTableImpl here */) {
+            override fun visitEnd() {
+                b.addOptionalAnnotationClass(t)
+            }
+        }
+        */
+
+        // The only flag present in module metadata is STRICT_SEMANTICS, and it seems fine to ignore it and set 0 here.
+        return b.build().serializeToByteArray(CompilerMetadataVersion(version.toIntArray(), false), 0)
+    }
+
 
     /**
      * A [KmModuleVisitor] that generates the metadata of a Kotlin JVM module file.
@@ -69,40 +108,6 @@ public class KotlinModuleMetadata private constructor(
         level = DeprecationLevel.ERROR
     )
     public class Writer {
-        private val b = JvmModuleProtoBuf.Module.newBuilder()
-
-        private fun writeModule(kmModule: KmModule) {
-            kmModule.packageParts.forEach { (fqName, packageParts) ->
-                PackageParts(fqName).apply {
-                    for (fileFacade in packageParts.fileFacades) {
-                        addPart(fileFacade, null)
-                    }
-                    for ((multiFileClassPart, multiFileFacade) in packageParts.multiFileClassParts) {
-                        addPart(multiFileClassPart, multiFileFacade)
-                    }
-
-                    addTo(b)
-                }
-            }
-
-            // visitAnnotation
-            /*
-            // TODO: move StringTableImpl to module 'metadata' and support module annotations here
-            b.addAnnotation(ProtoBuf.Annotation.newBuilder().apply {
-                id = annotation.className.name // <-- use StringTableImpl here
-            })
-            */
-
-            // visitOptionalAnnotationClass
-            /*
-            return object : ClassWriter(TODO() /* use StringTableImpl here */) {
-                override fun visitEnd() {
-                    b.addOptionalAnnotationClass(t)
-                }
-            }
-            */
-        }
-
         /**
          * Returns the metadata of the module file that was written with this writer.
          *
@@ -126,26 +131,6 @@ public class KotlinModuleMetadata private constructor(
     @Deprecated(VISITOR_API_MESSAGE, level = DeprecationLevel.ERROR)
     public fun accept(v: KmModuleVisitor): Unit = kmModule.accept(v)
 
-    private fun readImpl(): KmModule {
-        val v = KmModule()
-        for ((fqName, parts) in data.packageFqName2Parts) {
-            val (fileFacades, multiFileClassParts) = parts.parts.partition { parts.getMultifileFacadeName(it) == null }
-            v.packageParts[fqName] = KmPackageParts(
-                fileFacades.toMutableList(),
-                multiFileClassParts.associateWith { parts.getMultifileFacadeName(it)!! }.toMutableMap()
-            )
-        }
-
-        for (annotation in data.moduleData.annotations) {
-            v.annotations.add(KmAnnotation(annotation, emptyMap()))
-        }
-
-        for (classProto in data.moduleData.optionalAnnotations) {
-            v.optionalAnnotationClasses.add(classProto.toKmClass(data.moduleData.nameResolver))
-        }
-        return v
-    }
-
     /**
      * Collection of methods for reading and writing [KotlinModuleMetadata].
      */
@@ -161,12 +146,16 @@ public class KotlinModuleMetadata private constructor(
         @UnstableMetadataApi
         public fun read(bytes: ByteArray): KotlinModuleMetadata {
             return wrapIntoMetadataExceptionWhenNeeded {
-                val result = dataFromBytes(bytes)
+                val result = ModuleMapping.loadModuleMapping(
+                    bytes, "KotlinModuleMetadata", skipMetadataVersionCheck = false,
+                    isJvmPackageNameSupported = true
+                ) { throwIfNotCompatible(it, lenient = false /* TODO */) }
                 when (result) {
                     ModuleMapping.EMPTY, ModuleMapping.CORRUPTED ->
-                        throw InconsistentKotlinMetadataException("Data is not the content of a .kotlin_module file, or it has been corrupted.")
+                        throw IllegalArgumentException("Data is not the content of a .kotlin_module file, or it has been corrupted.")
                 }
-                KotlinModuleMetadata(bytes, result)
+                val module = readModuleMetadataImpl(result)
+                KotlinModuleMetadata(module, JvmMetadataVersion(result.version.toArray()))
             }
         }
 
@@ -179,18 +168,11 @@ public class KotlinModuleMetadata private constructor(
          * @throws IllegalArgumentException if [kmModule] is not correct and cannot be written or if [metadataVersion] is not supported for writing.
          */
         @UnstableMetadataApi
+        @Deprecated("Use a KotlinModuleMetadata instance and its write() member function", level = DeprecationLevel.WARNING)
         @JvmStatic
         @JvmOverloads
         public fun write(kmModule: KmModule, metadataVersion: JvmMetadataVersion = JvmMetadataVersion.CURRENT): ByteArray = wrapWriteIntoIAE {
-            val w = Writer().also { it.writeModule(kmModule) }
+            return KotlinModuleMetadata(kmModule, metadataVersion).write()
-            return w.b.build().serializeToByteArray(CompilerMetadataVersion(metadataVersion.toIntArray(), false), 0)
-        }
-
-        private fun dataFromBytes(bytes: ByteArray): ModuleMapping {
-            return ModuleMapping.loadModuleMapping(
-                bytes, "KotlinModuleMetadata", skipMetadataVersionCheck = false,
-                isJvmPackageNameSupported = true
-            ) { throwIfNotCompatible(it, lenient = false /* TODO */) }
         }
     }
 }

libraries/kotlinx-metadata/jvm/src/kotlinx/metadata/jvm/internal/JvmReadUtils.kt

@@ -5,15 +5,20 @@
 
 package kotlinx.metadata.jvm.internal
 
+import kotlinx.metadata.KmAnnotation
 import kotlinx.metadata.KmClass
 import kotlinx.metadata.KmLambda
 import kotlinx.metadata.KmPackage
 import kotlinx.metadata.internal.toKmClass
 import kotlinx.metadata.internal.toKmLambda
 import kotlinx.metadata.internal.toKmPackage
+import kotlinx.metadata.jvm.KmModule
+import kotlinx.metadata.jvm.KmPackageParts
 import kotlinx.metadata.jvm.KotlinClassMetadata
+import kotlinx.metadata.jvm.UnstableMetadataApi
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMetadataVersion
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
+import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping
 
 internal object JvmReadUtils {
     internal fun readKmClass(annotationData: Metadata): KmClass {
@@ -50,7 +55,29 @@ internal object JvmReadUtils {
         }
     }
 
-    internal fun checkMetadataVersionForRead(annotationData: Metadata, lenient: Boolean) {
+    @UnstableMetadataApi
+    internal fun readModuleMetadataImpl(data: ModuleMapping): KmModule {
+        val v = KmModule()
+        for ((fqName, parts) in data.packageFqName2Parts) {
+            val (fileFacades, multiFileClassParts) = parts.parts.partition { parts.getMultifileFacadeName(it) == null }
+            v.packageParts[fqName] = KmPackageParts(
+                fileFacades.toMutableList(),
+                multiFileClassParts.associateWith { parts.getMultifileFacadeName(it)!! }.toMutableMap()
+            )
+        }
+
+        for (annotation in data.moduleData.annotations) {
+            @Suppress("DEPRECATION_ERROR")
+            v.annotations.add(KmAnnotation(annotation, emptyMap()))
+        }
+
+        for (classProto in data.moduleData.optionalAnnotations) {
+            v.optionalAnnotationClasses.add(classProto.toKmClass(data.moduleData.nameResolver))
+        }
+        return v
+    }
+
+    private fun checkMetadataVersionForRead(annotationData: Metadata, lenient: Boolean) {
         if (annotationData.metadataVersion.isEmpty())
             throw IllegalArgumentException("Provided Metadata instance does not have metadataVersion in it and therefore is malformed and cannot be read.")
         val jvmMetadataVersion = JvmMetadataVersion(

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

@@ -44,7 +44,7 @@ class Kotlinp(private val settings: KotlinpSettings) {
 
     @OptIn(UnstableMetadataApi::class)
     internal fun readModuleFile(file: File): KotlinModuleMetadata? =
-        KotlinModuleMetadata.read(file.readBytes())
+        runCatching { KotlinModuleMetadata.read(file.readBytes()) }.getOrNull()
 }
 
 data class KotlinpSettings(

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

@@ -69,7 +69,7 @@ private fun compileAndPrintAllFiles(
         when (outputFile.extension) {
             "kotlin_module" -> {
                 val moduleFile = kotlinp.readModuleFile(outputFile)!!
-                val transformedWithNodes = KotlinModuleMetadata.read(transformModuleFileWithNodes(moduleFile))
+                val transformedWithNodes = KotlinModuleMetadata.read(moduleFile.write())
 
                 for ((sb, moduleFileToRender) in listOf(
                     main to moduleFile, afterNodes to transformedWithNodes
@@ -128,7 +128,3 @@ private fun StringBuilder.appendFileName(file: File) {
     appendLine("// ${file.invariantSeparatorsPath}")
     appendLine("// ------------------------------------------")
 }
-
-@OptIn(UnstableMetadataApi::class)
-private fun transformModuleFileWithNodes(moduleFile: KotlinModuleMetadata): ByteArray =
-    KotlinModuleMetadata.write(moduleFile.kmModule)