Implement writing visitors for kotlinx-metadata

#KT-23198

libraries/kotlinx-metadata/jvm/build.gradle.kts

@@ -24,6 +24,9 @@ dependencies {
     shadows(project(":core:metadata.jvm"))
     shadows(protobufLite())
     testCompile(commonDep("junit:junit"))
+    testCompile(intellijDep()) { includeJars("asm-all", rootProject = rootProject) }
+    testCompileOnly(project(":kotlin-reflect-api"))
+    testRuntime(projectDist(":kotlin-reflect"))
 }
 
 noDefaultJar()

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

@@ -9,7 +9,11 @@ import kotlinx.metadata.InconsistentKotlinMetadataException
 import kotlinx.metadata.KmClassVisitor
 import kotlinx.metadata.KmLambdaVisitor
 import kotlinx.metadata.KmPackageVisitor
+import kotlinx.metadata.impl.ClassWriter
+import kotlinx.metadata.impl.LambdaWriter
+import kotlinx.metadata.impl.PackageWriter
 import kotlinx.metadata.impl.accept
+import kotlinx.metadata.jvm.impl.writeProtoBufData
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
 import kotlin.LazyThreadSafetyMode.PUBLICATION
 
@@ -39,6 +43,34 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
             val (strings, proto) = classData
             proto.accept(v, strings)
         }
+
+        /**
+         * A [KmClassVisitor] that generates the metadata of a Kotlin class.
+         */
+        class Writer : ClassWriter() {
+            /**
+             * Returns the metadata of the class 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
+             * @param bytecodeVersion bytecode version to be written to the metadata (see [KotlinClassHeader.bytecodeVersion]),
+             *   [KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION] by default
+             * @param extraInt the value of the class-level flags to be written to the metadata (see [KotlinClassHeader.extraInt]),
+             *   0 by default
+             */
+            @JvmOverloads
+            fun write(
+                metadataVersion: IntArray = KotlinClassHeader.COMPATIBLE_METADATA_VERSION,
+                bytecodeVersion: IntArray = KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION,
+                extraInt: Int = 0
+            ): KotlinClassMetadata.Class {
+                val (d1, d2) = writeProtoBufData(t.build(), c)
+                val metadata = KotlinClassHeader(
+                    KotlinClassHeader.CLASS_KIND, metadataVersion, bytecodeVersion, d1, d2, null, null, extraInt
+                )
+                return KotlinClassMetadata.Class(metadata)
+            }
+        }
     }
 
     /**
@@ -60,6 +92,34 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
             val (strings, proto) = packageData
             proto.accept(v, strings)
         }
+
+        /**
+         * A [KmPackageVisitor] that generates the metadata of a Kotlin file facade.
+         */
+        class Writer : PackageWriter() {
+            /**
+             * Returns the metadata of the file facade 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
+             * @param bytecodeVersion bytecode version to be written to the metadata (see [KotlinClassHeader.bytecodeVersion]),
+             *   [KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION] by default
+             * @param extraInt the value of the class-level flags to be written to the metadata (see [KotlinClassHeader.extraInt]),
+             *   0 by default
+             */
+            @JvmOverloads
+            fun write(
+                metadataVersion: IntArray = KotlinClassHeader.COMPATIBLE_METADATA_VERSION,
+                bytecodeVersion: IntArray = KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION,
+                extraInt: Int = 0
+            ): KotlinClassMetadata.FileFacade {
+                val (d1, d2) = writeProtoBufData(t.build(), c)
+                val metadata = KotlinClassHeader(
+                    KotlinClassHeader.FILE_FACADE_KIND, metadataVersion, bytecodeVersion, d1, d2, null, null, extraInt
+                )
+                return KotlinClassMetadata.FileFacade(metadata)
+            }
+        }
     }
 
     /**
@@ -95,6 +155,39 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
             val (strings, proto) = functionData!!
             proto.accept(v, strings)
         }
+
+        /**
+         * A [KmLambdaVisitor] that generates the metadata of a synthetic class. To generate metadata of a Kotlin lambda,
+         * call [Writer.visitFunction] and [Writer.visitEnd] on a newly created instance of this writer. If these methods are not called,
+         * the resulting metadata will represent a _non-lambda_ synthetic class.
+         */
+        class Writer : LambdaWriter() {
+            /**
+             * Returns the metadata of the synthetic class 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
+             * @param bytecodeVersion bytecode version to be written to the metadata (see [KotlinClassHeader.bytecodeVersion]),
+             *   [KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION] by default
+             * @param extraInt the value of the class-level flags to be written to the metadata (see [KotlinClassHeader.extraInt]),
+             *   0 by default
+             */
+            @JvmOverloads
+            fun write(
+                metadataVersion: IntArray = KotlinClassHeader.COMPATIBLE_METADATA_VERSION,
+                bytecodeVersion: IntArray = KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION,
+                extraInt: Int = 0
+            ): KotlinClassMetadata.SyntheticClass {
+                val proto = t?.build()
+                val (d1, d2) =
+                        if (proto != null) writeProtoBufData(proto, c)
+                        else Pair(emptyArray(), emptyArray())
+                val metadata = KotlinClassHeader(
+                    KotlinClassHeader.SYNTHETIC_CLASS_KIND, metadataVersion, bytecodeVersion, d1, d2, null, null, extraInt
+                )
+                return KotlinClassMetadata.SyntheticClass(metadata)
+            }
+        }
     }
 
     /**
@@ -107,6 +200,36 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
          * JVM internal names of the part classes which this multi-file class combines.
          */
         val partClassNames: List<String> = header.data1.asList()
+
+        /**
+         * A writer that generates the metadata of a multi-file class facade.
+         */
+        class Writer {
+            /**
+             * Returns the metadata of the multi-file class facade that was written with this writer.
+             *
+             * @param partClassNames JVM internal names of the part classes which this multi-file class combines
+             * @param metadataVersion metadata version to be written to the metadata (see [KotlinClassHeader.metadataVersion]),
+             *   [KotlinClassHeader.COMPATIBLE_METADATA_VERSION] by default
+             * @param bytecodeVersion bytecode version to be written to the metadata (see [KotlinClassHeader.bytecodeVersion]),
+             *   [KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION] by default
+             * @param extraInt the value of the class-level flags to be written to the metadata (see [KotlinClassHeader.extraInt]),
+             *   0 by default
+             */
+            @JvmOverloads
+            fun write(
+                partClassNames: List<String>,
+                metadataVersion: IntArray = KotlinClassHeader.COMPATIBLE_METADATA_VERSION,
+                bytecodeVersion: IntArray = KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION,
+                extraInt: Int = 0
+            ): KotlinClassMetadata.MultiFileClassFacade {
+                val metadata = KotlinClassHeader(
+                    KotlinClassHeader.MULTI_FILE_CLASS_FACADE_KIND, metadataVersion, bytecodeVersion, partClassNames.toTypedArray(),
+                    null, null, null, extraInt
+                )
+                return KotlinClassMetadata.MultiFileClassFacade(metadata)
+            }
+        }
     }
 
     /**
@@ -137,6 +260,36 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
             val (strings, proto) = packageData
             proto.accept(v, strings)
         }
+
+        /**
+         * A [KmPackageVisitor] that generates the metadata of a multi-file class part.
+         */
+        class Writer : PackageWriter() {
+            /**
+             * Returns the metadata of the multi-file class part that was written with this writer.
+             *
+             * @param facadeClassName JVM internal name of the corresponding multi-file class facade
+             * @param metadataVersion metadata version to be written to the metadata (see [KotlinClassHeader.metadataVersion]),
+             *   [KotlinClassHeader.COMPATIBLE_METADATA_VERSION] by default
+             * @param bytecodeVersion bytecode version to be written to the metadata (see [KotlinClassHeader.bytecodeVersion]),
+             *   [KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION] by default
+             * @param extraInt the value of the class-level flags to be written to the metadata (see [KotlinClassHeader.extraInt]),
+             *   0 by default
+             */
+            @JvmOverloads
+            fun write(
+                facadeClassName: String,
+                metadataVersion: IntArray = KotlinClassHeader.COMPATIBLE_METADATA_VERSION,
+                bytecodeVersion: IntArray = KotlinClassHeader.COMPATIBLE_BYTECODE_VERSION,
+                extraInt: Int = 0
+            ): KotlinClassMetadata.MultiFileClassPart {
+                val (d1, d2) = writeProtoBufData(t.build(), c)
+                val metadata = KotlinClassHeader(
+                    KotlinClassHeader.MULTI_FILE_CLASS_PART_KIND, metadataVersion, bytecodeVersion, d1, d2, facadeClassName, null, extraInt
+                )
+                return KotlinClassMetadata.MultiFileClassPart(metadata)
+            }
+        }
     }
 
     /**

libraries/kotlinx-metadata/jvm/src/kotlinx/metadata/jvm/impl/JvmMetadataExtensions.kt

@@ -8,6 +8,7 @@ package kotlinx.metadata.jvm.impl
 import kotlinx.metadata.*
 import kotlinx.metadata.impl.extensions.MetadataExtensions
 import kotlinx.metadata.impl.readAnnotation
+import kotlinx.metadata.impl.writeAnnotation
 import kotlinx.metadata.jvm.*
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.NameResolver
@@ -15,6 +16,8 @@ import org.jetbrains.kotlin.metadata.deserialization.TypeTable
 import org.jetbrains.kotlin.metadata.deserialization.getExtensionOrNull
 import org.jetbrains.kotlin.metadata.jvm.JvmProtoBuf
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
+import org.jetbrains.kotlin.metadata.jvm.serialization.JvmStringTable
+import org.jetbrains.kotlin.metadata.serialization.StringTable
 
 internal class JvmMetadataExtensions : MetadataExtensions {
     override fun readFunctionExtensions(v: KmFunctionVisitor, proto: ProtoBuf.Function, strings: NameResolver, types: TypeTable) {
@@ -65,4 +68,115 @@ internal class JvmMetadataExtensions : MetadataExtensions {
         }
         ext.visitEnd()
     }
+
+    override fun createStringTable(): StringTable = JvmStringTable()
+
+    override fun writeFunctionExtensions(
+        type: KmExtensionType, proto: ProtoBuf.Function.Builder, strings: StringTable
+    ): KmFunctionExtensionVisitor? {
+        if (type != JvmFunctionExtensionVisitor.TYPE) return null
+        return object : JvmFunctionExtensionVisitor() {
+            override fun visit(desc: String?) {
+                if (desc != null) {
+                    proto.setExtension(JvmProtoBuf.methodSignature, desc.toJvmMethodSignature(strings))
+                }
+            }
+        }
+    }
+
+    override fun writePropertyExtensions(
+        type: KmExtensionType, proto: ProtoBuf.Property.Builder, strings: StringTable
+    ): KmPropertyExtensionVisitor? {
+        if (type != JvmPropertyExtensionVisitor.TYPE) return null
+        return object : JvmPropertyExtensionVisitor() {
+            var signature: JvmProtoBuf.JvmPropertySignature.Builder? = null
+
+            override fun visit(fieldName: String?, fieldTypeDesc: String?, getterDesc: String?, setterDesc: String?) {
+                if (fieldName == null && fieldTypeDesc == null && getterDesc == null && setterDesc == null) return
+
+                if (signature == null) {
+                    signature = JvmProtoBuf.JvmPropertySignature.newBuilder()
+                }
+                signature!!.apply {
+                    if (fieldName != null || fieldTypeDesc != null) {
+                        field = JvmProtoBuf.JvmFieldSignature.newBuilder().also { field ->
+                            if (fieldName != null) {
+                                field.name = strings.getStringIndex(fieldName)
+                            }
+                            if (fieldTypeDesc != null) {
+                                field.desc = strings.getStringIndex(fieldTypeDesc)
+                            }
+                        }.build()
+                    }
+                    if (getterDesc != null) {
+                        getter = getterDesc.toJvmMethodSignature(strings)
+                    }
+                    if (setterDesc != null) {
+                        setter = setterDesc.toJvmMethodSignature(strings)
+                    }
+                }
+            }
+
+            override fun visitSyntheticMethodForAnnotations(desc: String?) {
+                if (desc == null) return
+
+                if (signature == null) {
+                    signature = JvmProtoBuf.JvmPropertySignature.newBuilder()
+                }
+
+                signature!!.syntheticMethod = desc.toJvmMethodSignature(strings)
+            }
+
+            override fun visitEnd() {
+                if (signature != null) {
+                    proto.setExtension(JvmProtoBuf.propertySignature, signature!!.build())
+                }
+            }
+        }
+    }
+
+    override fun writeConstructorExtensions(
+        type: KmExtensionType, proto: ProtoBuf.Constructor.Builder, strings: StringTable
+    ): KmConstructorExtensionVisitor? {
+        if (type != JvmConstructorExtensionVisitor.TYPE) return null
+        return object : JvmConstructorExtensionVisitor() {
+            override fun visit(desc: String?) {
+                if (desc != null) {
+                    proto.setExtension(JvmProtoBuf.constructorSignature, desc.toJvmMethodSignature(strings))
+                }
+            }
+        }
+    }
+
+    override fun writeTypeParameterExtensions(
+        type: KmExtensionType, proto: ProtoBuf.TypeParameter.Builder, strings: StringTable
+    ): KmTypeParameterExtensionVisitor? {
+        if (type != JvmTypeParameterExtensionVisitor.TYPE) return null
+        return object : JvmTypeParameterExtensionVisitor() {
+            override fun visitAnnotation(annotation: KmAnnotation) {
+                proto.addExtension(JvmProtoBuf.typeParameterAnnotation, annotation.writeAnnotation(strings).build())
+            }
+        }
+    }
+
+    override fun writeTypeExtensions(type: KmExtensionType, proto: ProtoBuf.Type.Builder, strings: StringTable): KmTypeExtensionVisitor? {
+        if (type != JvmTypeExtensionVisitor.TYPE) return null
+        return object : JvmTypeExtensionVisitor() {
+            override fun visit(isRaw: Boolean) {
+                if (isRaw) {
+                    proto.setExtension(JvmProtoBuf.isRaw, true)
+                }
+            }
+
+            override fun visitAnnotation(annotation: KmAnnotation) {
+                proto.addExtension(JvmProtoBuf.typeAnnotation, annotation.writeAnnotation(strings).build())
+            }
+        }
+    }
+
+    private fun String.toJvmMethodSignature(strings: StringTable): JvmProtoBuf.JvmMethodSignature =
+        JvmProtoBuf.JvmMethodSignature.newBuilder().apply {
+            name = strings.getStringIndex(substringBefore('('))
+            desc = strings.getStringIndex("(" + substringAfter('('))
+        }.build()
 }

libraries/kotlinx-metadata/jvm/src/kotlinx/metadata/jvm/impl/jvmWriteUtils.kt

@@ -0,0 +1,19 @@
+/*
+ * 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.impl
+
+import kotlinx.metadata.impl.WriteContext
+import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
+import org.jetbrains.kotlin.metadata.jvm.serialization.JvmStringTable
+import org.jetbrains.kotlin.protobuf.MessageLite
+
+internal fun writeProtoBufData(message: MessageLite, c: WriteContext): Pair<Array<String>, Array<String>> {
+    val strings = c.strings as JvmStringTable
+    return Pair(
+        JvmProtoBufUtil.writeData(message, strings),
+        strings.strings.toTypedArray()
+    )
+}

libraries/kotlinx-metadata/jvm/test/kotlinx/metadata/test/MetadataSmokeTest.kt

@@ -6,11 +6,16 @@
 package kotlinx.metadata.test
 
 import kotlinx.metadata.*
+import kotlinx.metadata.jvm.JvmConstructorExtensionVisitor
 import kotlinx.metadata.jvm.JvmFunctionExtensionVisitor
 import kotlinx.metadata.jvm.KotlinClassHeader
 import kotlinx.metadata.jvm.KotlinClassMetadata
+import org.jetbrains.org.objectweb.asm.ClassWriter
+import org.jetbrains.org.objectweb.asm.Opcodes
 import org.junit.Assert.assertEquals
 import org.junit.Test
+import java.net.URLClassLoader
+import kotlin.reflect.full.primaryConstructor
 
 class MetadataSmokeTest {
     @Suppress("INVISIBLE_MEMBER", "INVISIBLE_REFERENCE")
@@ -55,4 +60,92 @@ class MetadataSmokeTest {
             inlineFunctions
         )
     }
+
+    @Test
+    fun produceKotlinClassFile() {
+        // First, produce a KotlinMetadata instance with the kotlinx-metadata API, for the following class:
+        //     class Hello {
+        //         fun hello(): String = "Hello, world!"
+        //     }
+
+        val header = KotlinClassMetadata.Class.Writer().run {
+            visit(Flags(Flags.IS_PUBLIC), "Hello")
+            visitConstructor(Flags(Flags.IS_PUBLIC, Flags.Constructor.IS_PRIMARY))!!.run {
+                (visitExtensions(JvmConstructorExtensionVisitor.TYPE) as JvmConstructorExtensionVisitor).run {
+                    visit("<init>()V")
+                }
+                visitEnd()
+            }
+            visitFunction(Flags(Flags.IS_PUBLIC, Flags.Function.IS_DECLARATION), "hello")!!.run {
+                visitReturnType(0)!!.run {
+                    visitClass("kotlin/String")
+                    visitEnd()
+                }
+                (visitExtensions(JvmFunctionExtensionVisitor.TYPE) as JvmFunctionExtensionVisitor).run {
+                    visit("hello()Ljava/lang/String;")
+                }
+                visitEnd()
+            }
+            visitEnd()
+
+            write().header
+        }
+
+        // Then, produce the bytecode of a .class file with ASM
+
+        val bytes = ClassWriter(0).run {
+            visit(Opcodes.V1_6, Opcodes.ACC_PUBLIC + Opcodes.ACC_FINAL + Opcodes.ACC_SUPER, "Hello", null, "java/lang/Object", null)
+
+            // Use the created KotlinMetadata instance to write @kotlin.Metadata annotation on the class file
+            visitAnnotation("Lkotlin/Metadata;", true).run {
+                visit("mv", header.metadataVersion)
+                visit("bv", header.bytecodeVersion)
+                visit("k", header.kind)
+                visitArray("d1").run {
+                    header.data1.forEach { visit(null, it) }
+                    visitEnd()
+                }
+                visitArray("d2").run {
+                    header.data2.forEach { visit(null, it) }
+                    visitEnd()
+                }
+                visitEnd()
+            }
+
+            visitMethod(Opcodes.ACC_PUBLIC + Opcodes.ACC_FINAL, "hello", "()Ljava/lang/String;", null, null).run {
+                visitAnnotation("Lorg/jetbrains/annotations/NotNull;", false).visitEnd()
+                visitCode()
+                visitLdcInsn("Hello, world!")
+                visitInsn(Opcodes.ARETURN)
+                visitMaxs(1, 1)
+                visitEnd()
+            }
+            visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null).run {
+                visitCode()
+                visitVarInsn(Opcodes.ALOAD, 0)
+                visitMethodInsn(Opcodes.INVOKESPECIAL, "java/lang/Object", "<init>", "()V", false)
+                visitInsn(Opcodes.RETURN)
+                visitMaxs(1, 1)
+                visitEnd()
+            }
+            visitEnd()
+
+            toByteArray()
+        }
+
+        // Finally, load the generated class, create its instance and invoke the method via Kotlin reflection.
+        // Kotlin reflection loads the metadata and builds a mapping from Kotlin symbols to JVM, so if the call succeeds,
+        // we can be sure that the metadata is consistent
+
+        val classLoader = object : URLClassLoader(emptyArray()) {
+            override fun findClass(name: String): Class<*> =
+                if (name == "Hello") defineClass(name, bytes, 0, bytes.size) else super.findClass(name)
+        }
+
+        val kClass = classLoader.loadClass("Hello").kotlin
+        val hello = kClass.primaryConstructor!!.call()
+        val result = kClass.members.single { it.name == "hello" }.call(hello) as String
+
+        assertEquals("Hello, world!", result)
+    }
 }

libraries/kotlinx-metadata/src/kotlinx/metadata/impl/extensions/MetadataExtensions.kt

@@ -9,6 +9,7 @@ import kotlinx.metadata.*
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.NameResolver
 import org.jetbrains.kotlin.metadata.deserialization.TypeTable
+import org.jetbrains.kotlin.metadata.serialization.StringTable
 import java.util.*
 
 interface MetadataExtensions {
@@ -22,6 +23,22 @@ interface MetadataExtensions {
 
     fun readTypeExtensions(v: KmTypeVisitor, proto: ProtoBuf.Type, strings: NameResolver)
 
+    fun createStringTable(): StringTable
+
+    fun writeFunctionExtensions(type: KmExtensionType, proto: ProtoBuf.Function.Builder, strings: StringTable): KmFunctionExtensionVisitor?
+
+    fun writePropertyExtensions(type: KmExtensionType, proto: ProtoBuf.Property.Builder, strings: StringTable): KmPropertyExtensionVisitor?
+
+    fun writeConstructorExtensions(
+        type: KmExtensionType, proto: ProtoBuf.Constructor.Builder, strings: StringTable
+    ): KmConstructorExtensionVisitor?
+
+    fun writeTypeParameterExtensions(
+        type: KmExtensionType, proto: ProtoBuf.TypeParameter.Builder, strings: StringTable
+    ): KmTypeParameterExtensionVisitor?
+
+    fun writeTypeExtensions(type: KmExtensionType, proto: ProtoBuf.Type.Builder, strings: StringTable): KmTypeExtensionVisitor?
+
     companion object {
         val INSTANCE: MetadataExtensions by lazy {
             ServiceLoader.load(MetadataExtensions::class.java).toList().firstOrNull()

libraries/kotlinx-metadata/src/kotlinx/metadata/impl/writeUtils.kt

@@ -0,0 +1,91 @@
+/*
+ * 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.impl
+
+import kotlinx.metadata.ClassName
+import kotlinx.metadata.KmAnnotation
+import kotlinx.metadata.KmAnnotationArgument
+import kotlinx.metadata.isLocal
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.serialization.StringTable
+
+fun KmAnnotation.writeAnnotation(strings: StringTable): ProtoBuf.Annotation.Builder =
+    ProtoBuf.Annotation.newBuilder().apply {
+        id = strings.getClassNameIndex(className)
+        for ((name, argument) in arguments) {
+            addArgument(ProtoBuf.Annotation.Argument.newBuilder().apply {
+                nameId = strings.getStringIndex(name)
+                value = argument.writeAnnotationArgument(strings).build()
+            })
+        }
+    }
+
+private fun KmAnnotationArgument<*>.writeAnnotationArgument(strings: StringTable): ProtoBuf.Annotation.Argument.Value.Builder =
+    ProtoBuf.Annotation.Argument.Value.newBuilder().apply {
+        when (this@writeAnnotationArgument) {
+            is KmAnnotationArgument.ByteValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.BYTE
+                this.intValue = value.toLong()
+            }
+            is KmAnnotationArgument.CharValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.CHAR
+                this.intValue = value.toLong()
+            }
+            is KmAnnotationArgument.ShortValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.SHORT
+                this.intValue = value.toLong()
+            }
+            is KmAnnotationArgument.IntValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.INT
+                this.intValue = value.toLong()
+            }
+            is KmAnnotationArgument.LongValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.LONG
+                this.intValue = value
+            }
+            is KmAnnotationArgument.FloatValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.FLOAT
+                this.floatValue = value
+            }
+            is KmAnnotationArgument.DoubleValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.DOUBLE
+                this.doubleValue = value
+            }
+            is KmAnnotationArgument.BooleanValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.BOOLEAN
+                this.intValue = if (value) 1 else 0
+            }
+            is KmAnnotationArgument.StringValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.STRING
+                this.stringValue = strings.getStringIndex(value)
+            }
+            is KmAnnotationArgument.KClassValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.CLASS
+                this.classId = strings.getClassNameIndex(value)
+            }
+            is KmAnnotationArgument.EnumValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.ENUM
+                this.classId = strings.getClassNameIndex(enumClassName)
+                this.enumValueId = strings.getStringIndex(enumEntryName)
+            }
+            is KmAnnotationArgument.AnnotationValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.ANNOTATION
+                this.annotation = value.writeAnnotation(strings).build()
+            }
+            is KmAnnotationArgument.ArrayValue -> {
+                this.type = ProtoBuf.Annotation.Argument.Value.Type.ARRAY
+                for (element in value) {
+                    this.addArrayElement(element.writeAnnotationArgument(strings))
+                }
+            }
+        }
+    }
+
+internal fun StringTable.getClassNameIndex(name: ClassName): Int =
+    if (name.isLocal)
+        getQualifiedClassNameIndex(name.substring(1), true)
+    else
+        getQualifiedClassNameIndex(name, false)

libraries/kotlinx-metadata/src/kotlinx/metadata/impl/writers.kt

@@ -0,0 +1,470 @@
+/*
+ * 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.impl
+
+import kotlinx.metadata.*
+import kotlinx.metadata.impl.extensions.MetadataExtensions
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.deserialization.VersionRequirement
+import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
+import org.jetbrains.kotlin.metadata.serialization.StringTable
+
+class WriteContext {
+    internal val extensions = MetadataExtensions.INSTANCE
+    val strings: StringTable = extensions.createStringTable()
+    val versionRequirements: MutableVersionRequirementTable = MutableVersionRequirementTable()
+
+    operator fun get(string: String): Int =
+        strings.getStringIndex(string)
+
+    fun getClassName(name: ClassName): Int =
+        strings.getClassNameIndex(name)
+}
+
+private fun writeTypeParameter(
+    c: WriteContext, flags: Int, name: String, id: Int, variance: KmVariance,
+    output: (ProtoBuf.TypeParameter.Builder) -> Unit
+): KmTypeParameterVisitor =
+    object : KmTypeParameterVisitor() {
+        private val t = ProtoBuf.TypeParameter.newBuilder()
+
+        override fun visitUpperBound(flags: Int): KmTypeVisitor? =
+            writeType(c, flags) { t.addUpperBound(it) }
+
+        override fun visitExtensions(type: KmExtensionType): KmTypeParameterExtensionVisitor? =
+            c.extensions.writeTypeParameterExtensions(type, t, c.strings)
+
+        override fun visitEnd() {
+            t.name = c[name]
+            t.id = id
+            val reified = Flags.TypeParameter.IS_REIFIED(flags)
+            if (reified != ProtoBuf.TypeParameter.getDefaultInstance().reified) {
+                t.reified = reified
+            }
+            if (variance == KmVariance.IN) {
+                t.variance = ProtoBuf.TypeParameter.Variance.IN
+            } else if (variance == KmVariance.OUT) {
+                t.variance = ProtoBuf.TypeParameter.Variance.OUT
+            }
+            output(t)
+        }
+    }
+
+private fun writeType(c: WriteContext, flags: Int, output: (ProtoBuf.Type.Builder) -> Unit): KmTypeVisitor =
+    object : KmTypeVisitor() {
+        private val t = ProtoBuf.Type.newBuilder()
+
+        override fun visitClass(name: ClassName) {
+            t.className = c.getClassName(name)
+        }
+
+        override fun visitTypeAlias(name: ClassName) {
+            t.typeAliasName = c.getClassName(name)
+        }
+
+        override fun visitStarProjection() {
+            t.addArgument(ProtoBuf.Type.Argument.newBuilder().apply {
+                projection = ProtoBuf.Type.Argument.Projection.STAR
+            })
+        }
+
+        override fun visitArgument(flags: Int, variance: KmVariance): KmTypeVisitor? =
+            writeType(c, flags) { argument ->
+                t.addArgument(ProtoBuf.Type.Argument.newBuilder().apply {
+                    if (variance == KmVariance.IN) {
+                        projection = ProtoBuf.Type.Argument.Projection.IN
+                    } else if (variance == KmVariance.OUT) {
+                        projection = ProtoBuf.Type.Argument.Projection.OUT
+                    }
+                    type = argument.build()
+                })
+            }
+
+        override fun visitTypeParameter(id: Int) {
+            t.typeParameter = id
+        }
+
+        override fun visitAbbreviatedType(flags: Int): KmTypeVisitor? =
+            writeType(c, flags) { t.abbreviatedType = it.build() }
+
+        override fun visitOuterType(flags: Int): KmTypeVisitor? =
+            writeType(c, flags) { t.outerType = it.build() }
+
+        override fun visitFlexibleTypeUpperBound(flags: Int, typeFlexibilityId: String?): KmTypeVisitor? =
+            writeType(c, flags) {
+                if (typeFlexibilityId != null) {
+                    t.flexibleTypeCapabilitiesId = c[typeFlexibilityId]
+                }
+                t.flexibleUpperBound = it.build()
+            }
+
+        override fun visitExtensions(type: KmExtensionType): KmTypeExtensionVisitor? =
+            c.extensions.writeTypeExtensions(type, t, c.strings)
+
+        override fun visitEnd() {
+            if (Flags.Type.IS_NULLABLE(flags)) {
+                t.nullable = true
+            }
+            val flagsToWrite = flags shr 1
+            if (flagsToWrite != ProtoBuf.Type.getDefaultInstance().flags) {
+                t.flags = flagsToWrite
+            }
+            output(t)
+        }
+    }
+
+private fun writeConstructor(c: WriteContext, flags: Int, output: (ProtoBuf.Constructor.Builder) -> Unit): KmConstructorVisitor =
+    object : KmConstructorVisitor() {
+        val t = ProtoBuf.Constructor.newBuilder()
+
+        override fun visitValueParameter(flags: Int, name: String): KmValueParameterVisitor? =
+            writeValueParameter(c, flags, name) { t.addValueParameter(it.build()) }
+
+        override fun visitVersionRequirement(): KmVersionRequirementVisitor? =
+            writeVersionRequirement(c) { t.versionRequirement = it }
+
+        override fun visitExtensions(type: KmExtensionType): KmConstructorExtensionVisitor? =
+            c.extensions.writeConstructorExtensions(type, t, c.strings)
+
+        override fun visitEnd() {
+            if (flags != ProtoBuf.Constructor.getDefaultInstance().flags) {
+                t.flags = flags
+            }
+            output(t)
+        }
+    }
+
+private fun writeFunction(c: WriteContext, flags: Int, name: String, output: (ProtoBuf.Function.Builder) -> Unit): KmFunctionVisitor =
+    object : KmFunctionVisitor() {
+        val t = ProtoBuf.Function.newBuilder()
+
+        override fun visitTypeParameter(flags: Int, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor? =
+            writeTypeParameter(c, flags, name, id, variance) { t.addTypeParameter(it) }
+
+        override fun visitReceiverParameterType(flags: Int): KmTypeVisitor? =
+            writeType(c, flags) { t.receiverType = it.build() }
+
+        override fun visitValueParameter(flags: Int, name: String): KmValueParameterVisitor? =
+            writeValueParameter(c, flags, name) { t.addValueParameter(it) }
+
+        override fun visitReturnType(flags: Int): KmTypeVisitor? =
+            writeType(c, flags) { t.returnType = it.build() }
+
+        override fun visitVersionRequirement(): KmVersionRequirementVisitor? =
+            writeVersionRequirement(c) { t.versionRequirement = it }
+
+        override fun visitContract(): KmContractVisitor? =
+            writeContract(c) { t.contract = it.build() }
+
+        override fun visitExtensions(type: KmExtensionType): KmFunctionExtensionVisitor? =
+            c.extensions.writeFunctionExtensions(type, t, c.strings)
+
+        override fun visitEnd() {
+            t.name = c[name]
+            if (flags != ProtoBuf.Function.getDefaultInstance().flags) {
+                t.flags = flags
+            }
+            output(t)
+        }
+    }
+
+private fun writeProperty(
+    c: WriteContext, flags: Int, name: String, getterFlags: Int, setterFlags: Int, output: (ProtoBuf.Property.Builder) -> Unit
+): KmPropertyVisitor = object : KmPropertyVisitor() {
+    val t = ProtoBuf.Property.newBuilder()
+
+    override fun visitTypeParameter(flags: Int, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor? =
+        writeTypeParameter(c, flags, name, id, variance) { t.addTypeParameter(it) }
+
+    override fun visitReceiverParameterType(flags: Int): KmTypeVisitor? =
+        writeType(c, flags) { t.receiverType = it.build() }
+
+    override fun visitSetterParameter(flags: Int, name: String): KmValueParameterVisitor? =
+        writeValueParameter(c, flags, name) { t.setterValueParameter = it.build() }
+
+    override fun visitReturnType(flags: Int): KmTypeVisitor? =
+        writeType(c, flags) { t.returnType = it.build() }
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor? =
+        writeVersionRequirement(c) { t.versionRequirement = it }
+
+    override fun visitExtensions(type: KmExtensionType): KmPropertyExtensionVisitor? =
+        c.extensions.writePropertyExtensions(type, t, c.strings)
+
+    override fun visitEnd() {
+        t.name = c[name]
+        if (flags != ProtoBuf.Property.getDefaultInstance().flags) {
+            t.flags = flags
+        }
+        // TODO: do not write getterFlags/setterFlags if not needed
+        t.getterFlags = getterFlags
+        t.setterFlags = setterFlags
+        output(t)
+    }
+}
+
+private fun writeValueParameter(
+    c: WriteContext, flags: Int, name: String,
+    output: (ProtoBuf.ValueParameter.Builder) -> Unit
+): KmValueParameterVisitor = object : KmValueParameterVisitor() {
+    val t = ProtoBuf.ValueParameter.newBuilder()
+
+    override fun visitType(flags: Int): KmTypeVisitor? =
+        writeType(c, flags) { t.type = it.build() }
+
+    override fun visitVarargElementType(flags: Int): KmTypeVisitor? =
+        writeType(c, flags) { t.varargElementType = it.build() }
+
+    override fun visitEnd() {
+        if (flags != ProtoBuf.ValueParameter.getDefaultInstance().flags) {
+            t.flags = flags
+        }
+        t.name = c[name]
+        output(t)
+    }
+}
+
+private fun writeTypeAlias(
+    c: WriteContext, flags: Int, name: String,
+    output: (ProtoBuf.TypeAlias.Builder) -> Unit
+): KmTypeAliasVisitor = object : KmTypeAliasVisitor() {
+    val t = ProtoBuf.TypeAlias.newBuilder()
+
+    override fun visitTypeParameter(flags: Int, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor? =
+        writeTypeParameter(c, flags, name, id, variance) { t.addTypeParameter(it) }
+
+    override fun visitUnderlyingType(flags: Int): KmTypeVisitor? =
+        writeType(c, flags) { t.underlyingType = it.build() }
+
+    override fun visitExpandedType(flags: Int): KmTypeVisitor? =
+        writeType(c, flags) { t.expandedType = it.build() }
+
+    override fun visitAnnotation(annotation: KmAnnotation) {
+        t.addAnnotation(annotation.writeAnnotation(c.strings))
+    }
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor? =
+        writeVersionRequirement(c) { t.versionRequirement = it }
+
+    override fun visitEnd() {
+        if (flags != ProtoBuf.TypeAlias.getDefaultInstance().flags) {
+            t.flags = flags
+        }
+        t.name = c[name]
+        output(t)
+    }
+}
+
+private fun writeVersionRequirement(
+    c: WriteContext, output: (Int) -> Unit
+): KmVersionRequirementVisitor = object : KmVersionRequirementVisitor() {
+    var t: ProtoBuf.VersionRequirement.Builder? = null
+
+    override fun visit(kind: KmVersionRequirementVersionKind, level: KmVersionRequirementLevel, errorCode: Int?, message: String?) {
+        t = ProtoBuf.VersionRequirement.newBuilder().apply {
+            if (kind != defaultInstanceForType.versionKind) {
+                this.versionKind = when (kind) {
+                    KmVersionRequirementVersionKind.LANGUAGE_VERSION -> ProtoBuf.VersionRequirement.VersionKind.LANGUAGE_VERSION
+                    KmVersionRequirementVersionKind.COMPILER_VERSION -> ProtoBuf.VersionRequirement.VersionKind.COMPILER_VERSION
+                    KmVersionRequirementVersionKind.API_VERSION -> ProtoBuf.VersionRequirement.VersionKind.API_VERSION
+                }
+            }
+            if (level != defaultInstanceForType.level) {
+                this.level = when (level) {
+                    KmVersionRequirementLevel.WARNING -> ProtoBuf.VersionRequirement.Level.WARNING
+                    KmVersionRequirementLevel.ERROR -> ProtoBuf.VersionRequirement.Level.ERROR
+                    KmVersionRequirementLevel.HIDDEN -> ProtoBuf.VersionRequirement.Level.HIDDEN
+                }
+            }
+            if (errorCode != null) {
+                this.errorCode = errorCode
+            }
+            if (message != null) {
+                this.message = c[message]
+            }
+        }
+    }
+
+    override fun visitVersion(major: Int, minor: Int, patch: Int) {
+        if (t == null) {
+            throw IllegalStateException("KmVersionRequirementVisitor.visit has not been called")
+        }
+        VersionRequirement.Version(major, minor, patch).encode(
+            writeVersion = { t!!.version = it },
+            writeVersionFull = { t!!.versionFull = it }
+        )
+    }
+
+    override fun visitEnd() {
+        if (t == null) {
+            throw IllegalStateException("KmVersionRequirementVisitor.visit has not been called")
+        }
+        output(c.versionRequirements[t!!])
+    }
+}
+
+private fun writeContract(c: WriteContext, output: (ProtoBuf.Contract.Builder) -> Unit): KmContractVisitor =
+    object : KmContractVisitor() {
+        val t = ProtoBuf.Contract.newBuilder()
+
+        override fun visitEffect(type: KmEffectType, invocationKind: KmEffectInvocationKind?): KmEffectVisitor? =
+            writeEffect(c, type, invocationKind) { t.addEffect(it) }
+
+        override fun visitEnd() {
+            output(t)
+        }
+    }
+
+private fun writeEffect(
+    c: WriteContext, type: KmEffectType, invocationKind: KmEffectInvocationKind?,
+    output: (ProtoBuf.Effect.Builder) -> Unit
+): KmEffectVisitor = object : KmEffectVisitor() {
+    val t = ProtoBuf.Effect.newBuilder()
+
+    override fun visitConstructorArgument(): KmEffectExpressionVisitor? =
+        writeEffectExpression(c) { t.addEffectConstructorArgument(it) }
+
+    override fun visitConclusionOfConditionalEffect(): KmEffectExpressionVisitor? =
+        writeEffectExpression(c) { t.conclusionOfConditionalEffect = it.build() }
+
+    @Suppress("UNUSED_VARIABLE") // force exhaustive whens
+    override fun visitEnd() {
+        val unused = when (type) {
+            KmEffectType.RETURNS_CONSTANT -> t.effectType = ProtoBuf.Effect.EffectType.RETURNS_CONSTANT
+            KmEffectType.CALLS -> t.effectType = ProtoBuf.Effect.EffectType.CALLS
+            KmEffectType.RETURNS_NOT_NULL -> t.effectType = ProtoBuf.Effect.EffectType.RETURNS_NOT_NULL
+        }
+        val unused2 = when (invocationKind) {
+            KmEffectInvocationKind.AT_MOST_ONCE -> t.kind = ProtoBuf.Effect.InvocationKind.AT_MOST_ONCE
+            KmEffectInvocationKind.EXACTLY_ONCE -> t.kind = ProtoBuf.Effect.InvocationKind.EXACTLY_ONCE
+            KmEffectInvocationKind.AT_LEAST_ONCE -> t.kind = ProtoBuf.Effect.InvocationKind.AT_LEAST_ONCE
+            null -> null
+        }
+        output(t)
+    }
+}
+
+private fun writeEffectExpression(c: WriteContext, output: (ProtoBuf.Expression.Builder) -> Unit): KmEffectExpressionVisitor =
+    object : KmEffectExpressionVisitor() {
+        val t = ProtoBuf.Expression.newBuilder()
+
+        override fun visit(flags: Int, parameterIndex: Int?) {
+            if (flags != ProtoBuf.Expression.getDefaultInstance().flags) {
+                t.flags = flags
+            }
+            if (parameterIndex != null) {
+                t.valueParameterReference = parameterIndex
+            }
+        }
+
+        override fun visitConstantValue(value: Any?) {
+            @Suppress("UNUSED_VARIABLE") // force exhaustive when
+            val unused = when (value) {
+                true -> t.constantValue = ProtoBuf.Expression.ConstantValue.TRUE
+                false -> t.constantValue = ProtoBuf.Expression.ConstantValue.FALSE
+                null -> null
+                else -> throw IllegalArgumentException("Only true, false or null constant values are allowed for effects (was=$value)")
+            }
+        }
+
+        override fun visitIsInstanceType(flags: Int): KmTypeVisitor? =
+            writeType(c, flags) { t.isInstanceType = it.build() }
+
+        override fun visitAndArgument(): KmEffectExpressionVisitor? =
+            writeEffectExpression(c) { t.addAndArgument(it) }
+
+        override fun visitOrArgument(): KmEffectExpressionVisitor? =
+            writeEffectExpression(c) { t.addOrArgument(it) }
+
+        override fun visitEnd() {
+            output(t)
+        }
+    }
+
+open class ClassWriter : KmClassVisitor() {
+    val t = ProtoBuf.Class.newBuilder()!!
+    val c = WriteContext()
+
+    override fun visit(flags: Int, name: ClassName) {
+        if (flags != ProtoBuf.Class.getDefaultInstance().flags) {
+            t.flags = flags
+        }
+        t.fqName = c.getClassName(name)
+    }
+
+    override fun visitTypeParameter(flags: Int, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor? =
+        writeTypeParameter(c, flags, name, id, variance) { t.addTypeParameter(it) }
+
+    override fun visitSupertype(flags: Int): KmTypeVisitor? =
+        writeType(c, flags) { t.addSupertype(it) }
+
+    override fun visitConstructor(flags: Int): KmConstructorVisitor? =
+        writeConstructor(c, flags) { t.addConstructor(it) }
+
+    override fun visitFunction(flags: Int, name: String): KmFunctionVisitor? =
+        writeFunction(c, flags, name) { t.addFunction(it) }
+
+    override fun visitProperty(flags: Int, name: String, getterFlags: Int, setterFlags: Int): KmPropertyVisitor? =
+        writeProperty(c, flags, name, getterFlags, setterFlags) { t.addProperty(it) }
+
+    override fun visitTypeAlias(flags: Int, name: String): KmTypeAliasVisitor? =
+        writeTypeAlias(c, flags, name) { t.addTypeAlias(it) }
+
+    override fun visitCompanionObject(name: String) {
+        t.companionObjectName = c[name]
+    }
+
+    override fun visitNestedClass(name: String) {
+        t.addNestedClassName(c[name])
+    }
+
+    override fun visitEnumEntry(name: String) {
+        t.addEnumEntry(ProtoBuf.EnumEntry.newBuilder().also { enumEntry ->
+            enumEntry.name = c[name]
+        })
+    }
+
+    override fun visitSealedSubclass(name: ClassName) {
+        t.addSealedSubclassFqName(c.getClassName(name))
+    }
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor? =
+        writeVersionRequirement(c) { t.versionRequirement = it }
+
+    override fun visitEnd() {
+        c.versionRequirements.serialize()?.let {
+            t.versionRequirementTable = it
+        }
+    }
+}
+
+open class PackageWriter : KmPackageVisitor() {
+    val t = ProtoBuf.Package.newBuilder()!!
+    val c = WriteContext()
+
+    override fun visitFunction(flags: Int, name: String): KmFunctionVisitor? =
+        writeFunction(c, flags, name) { t.addFunction(it) }
+
+    override fun visitProperty(flags: Int, name: String, getterFlags: Int, setterFlags: Int): KmPropertyVisitor? =
+        writeProperty(c, flags, name, getterFlags, setterFlags) { t.addProperty(it) }
+
+    override fun visitTypeAlias(flags: Int, name: String): KmTypeAliasVisitor? =
+        writeTypeAlias(c, flags, name) { t.addTypeAlias(it) }
+
+    override fun visitEnd() {
+        c.versionRequirements.serialize()?.let {
+            t.versionRequirementTable = it
+        }
+    }
+}
+
+open class LambdaWriter : KmLambdaVisitor() {
+    var t: ProtoBuf.Function.Builder? = null
+    val c = WriteContext()
+
+    override fun visitFunction(flags: Int, name: String): KmFunctionVisitor? =
+        writeFunction(c, flags, name) { t = it }
+}

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

@@ -10,6 +10,6 @@ import java.io.File
 
 abstract class AbstractKotlinpTest : TestCaseWithTmpdir() {
     protected fun doTest(fileName: String) {
-        compileAndPrintAllFiles(File(fileName), testRootDisposable, tmpdir, compareWithTxt = true)
+        compileAndPrintAllFiles(File(fileName), testRootDisposable, tmpdir, compareWithTxt = true, readWriteAndCompare = true)
     }
 }

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

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.kotlinp.test
 
 import com.intellij.openapi.Disposable
+import junit.framework.TestCase.assertEquals
+import kotlinx.metadata.jvm.KotlinClassMetadata
 import org.jetbrains.kotlin.checkers.setupLanguageVersionSettingsForCompilerTests
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
@@ -17,8 +19,9 @@ import org.jetbrains.kotlin.test.TestJdkKind
 import java.io.File
 import kotlin.test.fail
 
-fun compileAndPrintAllFiles(file: File, disposable: Disposable, tmpdir: File, compareWithTxt: Boolean) {
+fun compileAndPrintAllFiles(file: File, disposable: Disposable, tmpdir: File, compareWithTxt: Boolean, readWriteAndCompare: Boolean) {
-    val sb = StringBuilder()
+    val read = StringBuilder()
+    val readWriteRead = StringBuilder()
 
     compile(file, disposable, tmpdir) { outputFile ->
         when (outputFile.extension) {
@@ -26,15 +29,24 @@ fun compileAndPrintAllFiles(file: File, disposable: Disposable, tmpdir: File, co
                 // TODO: support kotlin_module files
             }
             "class" -> {
-                sb.appendFileName(outputFile.relativeTo(tmpdir))
+                val classFile = Kotlinp.readClassFile(outputFile)!!
-                sb.append(Kotlinp.renderClassFile(Kotlinp.readClassFile(outputFile)!!))
+                val classFile2 = transformClassFile(classFile)
+
+                for ((sb, classFileToRender) in listOf(read to classFile, readWriteRead to classFile2)) {
+                    sb.appendFileName(outputFile.relativeTo(tmpdir))
+                    sb.append(Kotlinp.renderClassFile(classFileToRender))
+                }
             }
             else -> fail("Unknown file: $outputFile")
         }
     }
 
     if (compareWithTxt) {
-        KotlinTestUtils.assertEqualsToFile(File(file.path.replace(".kt", ".txt")), sb.toString())
+        KotlinTestUtils.assertEqualsToFile(File(file.path.replace(".kt", ".txt")), read.toString())
+    }
+
+    if (readWriteAndCompare) {
+        assertEquals(read.toString(), readWriteRead.toString())
     }
 }
 
@@ -59,3 +71,23 @@ private fun StringBuilder.appendFileName(file: File) {
     appendln("// $file")
     appendln("// ------------------------------------------")
 }
+
+// Reads the class file and writes it back with *Writer visitors.
+// The resulting class file should be the same from the point of view of any metadata reader, including kotlinp
+// (the exact bytes may differ though, because there are multiple ways to encode the same metadata)
+private fun transformClassFile(classFile: KotlinClassMetadata): KotlinClassMetadata =
+    when (classFile) {
+        is KotlinClassMetadata.Class -> KotlinClassMetadata.Class.Writer().apply(classFile::accept).write()
+        is KotlinClassMetadata.FileFacade -> KotlinClassMetadata.FileFacade.Writer().apply(classFile::accept).write()
+        is KotlinClassMetadata.SyntheticClass -> {
+            val writer = KotlinClassMetadata.SyntheticClass.Writer()
+            if (classFile.isLambda) {
+                classFile.accept(writer)
+            }
+            writer.write()
+        }
+        is KotlinClassMetadata.MultiFileClassFacade -> KotlinClassMetadata.MultiFileClassFacade.Writer().write(classFile.partClassNames)
+        is KotlinClassMetadata.MultiFileClassPart ->
+            KotlinClassMetadata.MultiFileClassPart.Writer().apply(classFile::accept).write(classFile.facadeClassName)
+        else -> classFile
+    }