Introduce new API for kotlinx-metadata-jvm

Kdocs are provided in a subsequent commit

 #KT-26602 Fixed

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

@@ -5,10 +5,7 @@
 
 package kotlinx.metadata.jvm
 
-import kotlinx.metadata.InconsistentKotlinMetadataException
+import kotlinx.metadata.*
-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
@@ -36,6 +33,12 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
             JvmProtoBufUtil.readClassDataFrom(data1, header.data2)
         }
 
+        /**
+         * Visits metadata of this class with a new [KmClass] instance and returns that instance.
+         */
+        fun toKmClass(): KmClass =
+            KmClass().apply(this::accept)
+
         /**
          * Makes the given visitor visit metadata of this class.
          *
@@ -85,6 +88,12 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
             JvmProtoBufUtil.readPackageDataFrom(data1, header.data2)
         }
 
+        /**
+         * Visits metadata of this file facade with a new [KmPackage] instance and returns that instance.
+         */
+        fun toKmPackage(): KmPackage =
+            KmPackage().apply(this::accept)
+
         /**
          * Makes the given visitor visit metadata of this file facade.
          *
@@ -135,6 +144,14 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
             }
         }
 
+        /**
+         * Visits metadata of this synthetic class with a new [KmLambda] instance and returns that instance.
+         *
+         * Returns `null` if this synthetic class does not represent a lambda.
+         */
+        fun toKmLambda(): KmLambda? =
+            if (isLambda) KmLambda().apply(this::accept) else null
+
         /**
          * Returns `true` if this synthetic class is a class file compiled for a Kotlin lambda.
          */
@@ -253,6 +270,12 @@ sealed class KotlinClassMetadata(val header: KotlinClassHeader) {
         val facadeClassName: String
             get() = header.extraString
 
+        /**
+         * Visits metadata of this multi-file class part with a new [KmPackage] instance and returns that instance.
+         */
+        fun toKmPackage(): KmPackage =
+            KmPackage().apply(this::accept)
+
         /**
          * Makes the given visitor visit metadata of this multi-file class part.
          *

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

@@ -28,6 +28,12 @@ class KotlinModuleMetadata(@Suppress("CanBeParameter", "MemberVisibilityCanBePri
         // TODO: report incorrect versions of modules
     }
 
+    /**
+     * Visits metadata of this module with a new [KmModule] instance and returns that instance.
+     */
+    fun toKmModule(): KmModule =
+        KmModule().apply(this::accept)
+
     /**
      * A [KmModuleVisitor] that generates the metadata of a Kotlin JVM module file.
      */
@@ -151,3 +157,28 @@ abstract class KmModuleVisitor(private val delegate: KmModuleVisitor? = null) {
 
     // TODO: JvmPackageName
 }
+
+class KmModule : KmModuleVisitor() {
+    val packageParts: MutableMap<String, KmPackageParts> = LinkedHashMap()
+    val annotations: MutableList<KmAnnotation> = ArrayList(0)
+
+    override fun visitPackageParts(fqName: String, fileFacades: List<String>, multiFileClassParts: Map<String, String>) {
+        packageParts[fqName] = KmPackageParts(fileFacades.toMutableList(), multiFileClassParts.toMutableMap())
+    }
+
+    override fun visitAnnotation(annotation: KmAnnotation) {
+        annotations.add(annotation)
+    }
+
+    fun accept(visitor: KmModuleVisitor) {
+        for ((fqName, parts) in packageParts) {
+            visitor.visitPackageParts(fqName, parts.fileFacades, parts.multiFileClassParts)
+        }
+        annotations.forEach(visitor::visitAnnotation)
+    }
+}
+
+class KmPackageParts(
+    val fileFacades: MutableList<String>,
+    val multiFileClassParts: MutableMap<String, String>
+)

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

@@ -7,7 +7,7 @@ package kotlinx.metadata.jvm.impl
 
 import kotlinx.metadata.*
 import kotlinx.metadata.impl.*
-import kotlinx.metadata.impl.extensions.MetadataExtensions
+import kotlinx.metadata.impl.extensions.*
 import kotlinx.metadata.jvm.*
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.getExtensionOrNull
@@ -258,6 +258,20 @@ internal class JvmMetadataExtensions : MetadataExtensions {
         }
     }
 
+    override fun createClassExtension(): KmClassExtension = JvmClassExtension()
+
+    override fun createPackageExtension(): KmPackageExtension = JvmPackageExtension()
+
+    override fun createFunctionExtension(): KmFunctionExtension = JvmFunctionExtension()
+
+    override fun createPropertyExtension(): KmPropertyExtension = JvmPropertyExtension()
+
+    override fun createConstructorExtension(): KmConstructorExtension = JvmConstructorExtension()
+
+    override fun createTypeParameterExtension(): KmTypeParameterExtension = JvmTypeParameterExtension()
+
+    override fun createTypeExtension(): KmTypeExtension = JvmTypeExtension()
+
     private fun JvmMemberSignature.toJvmMethodSignature(c: WriteContext): JvmProtoBuf.JvmMethodSignature =
         JvmProtoBuf.JvmMethodSignature.newBuilder().apply {
             name = c[this@toJvmMethodSignature.name]

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

@@ -0,0 +1,178 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.*
+import kotlinx.metadata.impl.extensions.*
+import kotlinx.metadata.jvm.*
+
+internal val KmClass.jvm: JvmClassExtension
+    get() = visitExtensions(JvmClassExtensionVisitor.TYPE) as JvmClassExtension
+
+internal val KmPackage.jvm: JvmPackageExtension
+    get() = visitExtensions(JvmPackageExtensionVisitor.TYPE) as JvmPackageExtension
+
+internal val KmFunction.jvm: JvmFunctionExtension
+    get() = visitExtensions(JvmFunctionExtensionVisitor.TYPE) as JvmFunctionExtension
+
+internal val KmProperty.jvm: JvmPropertyExtension
+    get() = visitExtensions(JvmPropertyExtensionVisitor.TYPE) as JvmPropertyExtension
+
+internal val KmConstructor.jvm: JvmConstructorExtension
+    get() = visitExtensions(JvmConstructorExtensionVisitor.TYPE) as JvmConstructorExtension
+
+internal val KmTypeParameter.jvm: JvmTypeParameterExtension
+    get() = visitExtensions(JvmTypeParameterExtensionVisitor.TYPE) as JvmTypeParameterExtension
+
+internal val KmType.jvm: JvmTypeExtension
+    get() = visitExtensions(JvmTypeExtensionVisitor.TYPE) as JvmTypeExtension
+
+
+internal class JvmClassExtension : JvmClassExtensionVisitor(), KmClassExtension {
+    val localDelegatedProperties: MutableList<KmProperty> = ArrayList(0)
+    var moduleName: String? = null
+    var anonymousObjectOriginName: String? = null
+
+    override fun visitLocalDelegatedProperty(flags: Flags, name: String, getterFlags: Flags, setterFlags: Flags): KmPropertyVisitor =
+        KmProperty(flags, name, getterFlags, setterFlags).also { localDelegatedProperties.add(it) }
+
+    override fun visitModuleName(name: String) {
+        this.moduleName = name
+    }
+
+    override fun visitAnonymousObjectOriginName(internalName: String) {
+        this.anonymousObjectOriginName = internalName
+    }
+
+    override fun accept(visitor: KmClassExtensionVisitor) {
+        require(visitor is JvmClassExtensionVisitor)
+        localDelegatedProperties.forEach {
+            visitor.visitLocalDelegatedProperty(it.flags, it.name, it.getterFlags, it.setterFlags)?.let(it::accept)
+        }
+        moduleName?.let(visitor::visitModuleName)
+        anonymousObjectOriginName?.let(visitor::visitAnonymousObjectOriginName)
+        visitor.visitEnd()
+    }
+}
+
+internal class JvmPackageExtension : JvmPackageExtensionVisitor(), KmPackageExtension {
+    val localDelegatedProperties: MutableList<KmProperty> = ArrayList(0)
+    var moduleName: String? = null
+
+    override fun visitLocalDelegatedProperty(flags: Flags, name: String, getterFlags: Flags, setterFlags: Flags): KmPropertyVisitor =
+        KmProperty(flags, name, getterFlags, setterFlags).also { localDelegatedProperties.add(it) }
+
+    override fun visitModuleName(name: String) {
+        this.moduleName = name
+    }
+
+    override fun accept(visitor: KmPackageExtensionVisitor) {
+        require(visitor is JvmPackageExtensionVisitor)
+        localDelegatedProperties.forEach {
+            visitor.visitLocalDelegatedProperty(it.flags, it.name, it.getterFlags, it.setterFlags)?.let(it::accept)
+        }
+        moduleName?.let(visitor::visitModuleName)
+        visitor.visitEnd()
+    }
+}
+
+internal class JvmFunctionExtension : JvmFunctionExtensionVisitor(), KmFunctionExtension {
+    var signature: JvmMethodSignature? = null
+    var lambdaClassOriginName: String? = null
+
+    override fun visit(signature: JvmMethodSignature?) {
+        this.signature = signature
+    }
+
+    override fun visitLambdaClassOriginName(internalName: String) {
+        this.lambdaClassOriginName = internalName
+    }
+
+    override fun accept(visitor: KmFunctionExtensionVisitor) {
+        require(visitor is JvmFunctionExtensionVisitor)
+        visitor.visit(signature)
+        lambdaClassOriginName?.let(visitor::visitLambdaClassOriginName)
+        visitor.visitEnd()
+    }
+}
+
+internal class JvmPropertyExtension : JvmPropertyExtensionVisitor(), KmPropertyExtension {
+    var jvmFlags: Flags = 0
+    var fieldSignature: JvmFieldSignature? = null
+    var getterSignature: JvmMethodSignature? = null
+    var setterSignature: JvmMethodSignature? = null
+    var syntheticMethodForAnnotations: JvmMethodSignature? = null
+
+    override fun visit(
+        jvmFlags: Flags,
+        fieldSignature: JvmFieldSignature?,
+        getterSignature: JvmMethodSignature?,
+        setterSignature: JvmMethodSignature?
+    ) {
+        this.jvmFlags = jvmFlags
+        this.fieldSignature = fieldSignature
+        this.getterSignature = getterSignature
+        this.setterSignature = setterSignature
+    }
+
+    override fun visitSyntheticMethodForAnnotations(signature: JvmMethodSignature?) {
+        this.syntheticMethodForAnnotations = signature
+    }
+
+    override fun accept(visitor: KmPropertyExtensionVisitor) {
+        require(visitor is JvmPropertyExtensionVisitor)
+        visitor.visit(jvmFlags, fieldSignature, getterSignature, setterSignature)
+        visitor.visitSyntheticMethodForAnnotations(syntheticMethodForAnnotations)
+        visitor.visitEnd()
+    }
+}
+
+internal class JvmConstructorExtension : JvmConstructorExtensionVisitor(), KmConstructorExtension {
+    var signature: JvmMethodSignature? = null
+
+    override fun visit(signature: JvmMethodSignature?) {
+        this.signature = signature
+    }
+
+    override fun accept(visitor: KmConstructorExtensionVisitor) {
+        require(visitor is JvmConstructorExtensionVisitor)
+        visitor.visit(signature)
+    }
+}
+
+internal class JvmTypeParameterExtension : JvmTypeParameterExtensionVisitor(), KmTypeParameterExtension {
+    val annotations: MutableList<KmAnnotation> = mutableListOf()
+
+    override fun visitAnnotation(annotation: KmAnnotation) {
+        annotations.add(annotation)
+    }
+
+    override fun accept(visitor: KmTypeParameterExtensionVisitor) {
+        require(visitor is JvmTypeParameterExtensionVisitor)
+        annotations.forEach(visitor::visitAnnotation)
+        visitor.visitEnd()
+    }
+}
+
+internal class JvmTypeExtension : JvmTypeExtensionVisitor(), KmTypeExtension {
+    var isRaw: Boolean = false
+    val annotations: MutableList<KmAnnotation> = mutableListOf()
+
+    override fun visit(isRaw: Boolean) {
+        this.isRaw = isRaw
+    }
+
+    override fun visitAnnotation(annotation: KmAnnotation) {
+        annotations.add(annotation)
+    }
+
+    override fun accept(visitor: KmTypeExtensionVisitor) {
+        require(visitor is JvmTypeExtensionVisitor)
+        visitor.visit(isRaw)
+        annotations.forEach(visitor::visitAnnotation)
+        visitor.visitEnd()
+    }
+}

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

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.*
+import kotlinx.metadata.jvm.impl.jvm
+
+val KmClass.localDelegatedProperties: MutableList<KmProperty>
+    get() = jvm.localDelegatedProperties
+
+var KmClass.moduleName: String?
+    get() = jvm.moduleName
+    set(value) {
+        jvm.moduleName = value
+    }
+
+var KmClass.anonymousObjectOriginName: String?
+    get() = jvm.anonymousObjectOriginName
+    set(value) {
+        jvm.anonymousObjectOriginName = value
+    }
+
+val KmPackage.localDelegatedProperties: MutableList<KmProperty>
+    get() = jvm.localDelegatedProperties
+
+var KmPackage.moduleName: String?
+    get() = jvm.moduleName
+    set(value) {
+        jvm.moduleName = value
+    }
+
+var KmFunction.signature: JvmMethodSignature?
+    get() = jvm.signature
+    set(value) {
+        jvm.signature = value
+    }
+
+var KmFunction.lambdaClassOriginName: String?
+    get() = jvm.lambdaClassOriginName
+    set(value) {
+        jvm.lambdaClassOriginName = value
+    }
+
+var KmProperty.jvmFlags: Flags
+    get() = jvm.jvmFlags
+    set(value) {
+        jvm.jvmFlags = value
+    }
+
+var KmProperty.fieldSignature: JvmFieldSignature?
+    get() = jvm.fieldSignature
+    set(value) {
+        jvm.fieldSignature = value
+    }
+
+var KmProperty.getterSignature: JvmMethodSignature?
+    get() = jvm.getterSignature
+    set(value) {
+        jvm.getterSignature = value
+    }
+
+var KmProperty.setterSignature: JvmMethodSignature?
+    get() = jvm.setterSignature
+    set(value) {
+        jvm.setterSignature = value
+    }
+
+var KmProperty.syntheticMethodForAnnotations: JvmMethodSignature?
+    get() = jvm.syntheticMethodForAnnotations
+    set(value) {
+        jvm.syntheticMethodForAnnotations = value
+    }
+
+var KmConstructor.signature: JvmMethodSignature?
+    get() = jvm.signature
+    set(value) {
+        jvm.signature = value
+    }
+
+val KmTypeParameter.annotations: MutableList<KmAnnotation>
+    get() = jvm.annotations
+
+var KmType.isRaw: Boolean
+    get() = jvm.isRaw
+    set(value) {
+        jvm.isRaw = value
+    }
+
+val KmType.annotations: MutableList<KmAnnotation>
+    get() = jvm.annotations

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

@@ -44,6 +44,20 @@ interface MetadataExtensions {
 
     fun writeTypeExtensions(type: KmExtensionType, proto: ProtoBuf.Type.Builder, c: WriteContext): KmTypeExtensionVisitor?
 
+    fun createClassExtension(): KmClassExtension
+
+    fun createPackageExtension(): KmPackageExtension
+
+    fun createFunctionExtension(): KmFunctionExtension
+
+    fun createPropertyExtension(): KmPropertyExtension
+
+    fun createConstructorExtension(): KmConstructorExtension
+
+    fun createTypeParameterExtension(): KmTypeParameterExtension
+
+    fun createTypeExtension(): KmTypeExtension
+
     companion object {
         val INSTANCES: List<MetadataExtensions> by lazy {
             ServiceLoader.load(MetadataExtensions::class.java, MetadataExtensions::class.java.classLoader).toList().also {

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.extensions
+
+import kotlinx.metadata.*
+
+interface KmExtension<V : KmExtensionVisitor> : KmExtensionVisitor {
+    fun accept(visitor: V)
+}
+
+interface KmClassExtension : KmClassExtensionVisitor, KmExtension<KmClassExtensionVisitor>
+
+interface KmPackageExtension : KmPackageExtensionVisitor, KmExtension<KmPackageExtensionVisitor>
+
+interface KmFunctionExtension : KmFunctionExtensionVisitor, KmExtension<KmFunctionExtensionVisitor>
+
+interface KmPropertyExtension : KmPropertyExtensionVisitor, KmExtension<KmPropertyExtensionVisitor>
+
+interface KmConstructorExtension : KmConstructorExtensionVisitor, KmExtension<KmConstructorExtensionVisitor>
+
+interface KmTypeParameterExtension : KmTypeParameterExtensionVisitor, KmExtension<KmTypeParameterExtensionVisitor>
+
+interface KmTypeExtension : KmTypeExtensionVisitor, KmExtension<KmTypeExtensionVisitor>

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.extensions
+
+import kotlinx.metadata.KmExtensionType
+import kotlinx.metadata.KmExtensionVisitor
+
+internal fun <T : KmExtensionVisitor> applySingleExtension(type: KmExtensionType, block: MetadataExtensions.() -> T?): T? {
+    var result: T? = null
+    for (extension in MetadataExtensions.INSTANCES) {
+        val current = block(extension) ?: continue
+        if (result != null) {
+            throw IllegalStateException("Multiple extensions handle the same extension type: $type")
+        }
+        result = current
+    }
+    return result
+}
+
+internal fun <N : KmExtension<*>> Collection<N>.singleOfType(type: KmExtensionType): N {
+    var result: N? = null
+    for (node in this) {
+        if (node.type != type) continue
+        if (result != null) {
+            throw IllegalStateException("Multiple extensions handle the same extension type: $type")
+        }
+        result = node
+    }
+    if (result == null) {
+        throw IllegalStateException("No extensions handle the extension type: $type")
+    }
+    return result
+}

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

@@ -6,14 +6,13 @@
 package kotlinx.metadata.impl
 
 import kotlinx.metadata.*
-import kotlinx.metadata.impl.extensions.MetadataExtensions
+import kotlinx.metadata.impl.extensions.applySingleExtension
 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(val strings: StringTable) {
-    private val extensions = MetadataExtensions.INSTANCES
     val versionRequirements: MutableVersionRequirementTable = MutableVersionRequirementTable()
 
     operator fun get(string: String): Int =
@@ -21,18 +20,6 @@ class WriteContext(val strings: StringTable) {
 
     fun getClassName(name: ClassName): Int =
         strings.getClassNameIndex(name)
-
-    internal fun <T : KmExtensionVisitor> applySingleExtension(type: KmExtensionType, block: MetadataExtensions.() -> T?): T? {
-        var result: T? = null
-        for (extension in extensions) {
-            val current = block(extension) ?: continue
-            if (result != null) {
-                throw IllegalStateException("Multiple extensions handle the same extension type: $type")
-            }
-            result = current
-        }
-        return result
-    }
 }
 
 private fun writeTypeParameter(
@@ -46,7 +33,7 @@ private fun writeTypeParameter(
             writeType(c, flags) { t.addUpperBound(it) }
 
         override fun visitExtensions(type: KmExtensionType): KmTypeParameterExtensionVisitor? =
-            c.applySingleExtension(type) {
+            applySingleExtension(type) {
                 writeTypeParameterExtensions(type, t, c)
             }
 
@@ -115,7 +102,7 @@ private fun writeType(c: WriteContext, flags: Flags, output: (ProtoBuf.Type.Buil
             }
 
         override fun visitExtensions(type: KmExtensionType): KmTypeExtensionVisitor? =
-            c.applySingleExtension(type) {
+            applySingleExtension(type) {
                 writeTypeExtensions(type, t, c)
             }
 
@@ -142,7 +129,7 @@ private fun writeConstructor(c: WriteContext, flags: Flags, output: (ProtoBuf.Co
             writeVersionRequirement(c) { t.addVersionRequirement(it) }
 
         override fun visitExtensions(type: KmExtensionType): KmConstructorExtensionVisitor? =
-            c.applySingleExtension(type) {
+            applySingleExtension(type) {
                 writeConstructorExtensions(type, t, c)
             }
 
@@ -177,7 +164,7 @@ private fun writeFunction(c: WriteContext, flags: Flags, name: String, output: (
             writeContract(c) { t.contract = it.build() }
 
         override fun visitExtensions(type: KmExtensionType): KmFunctionExtensionVisitor? =
-            c.applySingleExtension(type) {
+            applySingleExtension(type) {
                 writeFunctionExtensions(type, t, c)
             }
 
@@ -211,7 +198,7 @@ fun writeProperty(
         writeVersionRequirement(c) { t.addVersionRequirement(it) }
 
     override fun visitExtensions(type: KmExtensionType): KmPropertyExtensionVisitor? =
-        c.applySingleExtension(type) {
+        applySingleExtension(type) {
             writePropertyExtensions(type, t, c)
         }
 
@@ -457,7 +444,7 @@ open class ClassWriter(stringTable: StringTable) : KmClassVisitor() {
         writeVersionRequirement(c) { t.addVersionRequirement(it) }
 
     override fun visitExtensions(type: KmExtensionType): KmClassExtensionVisitor? =
-        c.applySingleExtension(type) {
+        applySingleExtension(type) {
             writeClassExtensions(type, t, c)
         }
 
@@ -482,7 +469,7 @@ open class PackageWriter(stringTable: StringTable) : KmPackageVisitor() {
         writeTypeAlias(c, flags, name) { t.addTypeAlias(it) }
 
     override fun visitExtensions(type: KmExtensionType): KmPackageExtensionVisitor? =
-        c.applySingleExtension(type) {
+        applySingleExtension(type) {
             writePackageExtensions(type, t, c)
         }
 

libraries/kotlinx-metadata/src/kotlinx/metadata/nodes.kt

@@ -0,0 +1,518 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+@file:Suppress("MemberVisibilityCanBePrivate")
+
+package kotlinx.metadata
+
+import kotlinx.metadata.impl.extensions.*
+
+interface KmDeclarationContainer {
+    val functions: MutableList<KmFunction>
+    val properties: MutableList<KmProperty>
+    val typeAliases: MutableList<KmTypeAlias>
+}
+
+class KmClass : KmClassVisitor(), KmDeclarationContainer {
+    var flags: Flags = flagsOf()
+    lateinit var name: ClassName
+    val typeParameters: MutableList<KmTypeParameter> = ArrayList(0)
+    val supertypes: MutableList<KmType> = ArrayList(1)
+    override val functions: MutableList<KmFunction> = ArrayList()
+    override val properties: MutableList<KmProperty> = ArrayList()
+    override val typeAliases: MutableList<KmTypeAlias> = ArrayList(0)
+    val constructors: MutableList<KmConstructor> = ArrayList(1)
+    var companionObject: String? = null
+    val nestedClasses: MutableList<String> = ArrayList(0)
+    val enumEntries: MutableList<String> = ArrayList(0)
+    val sealedSubclasses: MutableList<ClassName> = ArrayList(0)
+    val versionRequirements: MutableList<KmVersionRequirement> = ArrayList(0)
+
+    private val extensions: List<KmClassExtension> =
+        MetadataExtensions.INSTANCES.map(MetadataExtensions::createClassExtension)
+
+    override fun visit(flags: Flags, name: ClassName) {
+        this.flags = flags
+        this.name = name
+    }
+
+    override fun visitTypeParameter(flags: Flags, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor =
+        KmTypeParameter(flags, name, id, variance).addTo(typeParameters)
+
+    override fun visitSupertype(flags: Flags): KmTypeVisitor =
+        KmType(flags).addTo(supertypes)
+
+    override fun visitFunction(flags: Flags, name: String): KmFunctionVisitor =
+        KmFunction(flags, name).addTo(functions)
+
+    override fun visitProperty(flags: Flags, name: String, getterFlags: Flags, setterFlags: Flags): KmPropertyVisitor =
+        KmProperty(flags, name, getterFlags, setterFlags).addTo(properties)
+
+    override fun visitTypeAlias(flags: Flags, name: String): KmTypeAliasVisitor =
+        KmTypeAlias(flags, name).addTo(typeAliases)
+
+    override fun visitConstructor(flags: Flags): KmConstructorVisitor =
+        KmConstructor(flags).addTo(constructors)
+
+    override fun visitCompanionObject(name: String) {
+        this.companionObject = name
+    }
+
+    override fun visitNestedClass(name: String) {
+        nestedClasses.add(name)
+    }
+
+    override fun visitEnumEntry(name: String) {
+        enumEntries.add(name)
+    }
+
+    override fun visitSealedSubclass(name: ClassName) {
+        sealedSubclasses.add(name)
+    }
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor =
+        KmVersionRequirement().addTo(versionRequirements)
+
+    override fun visitExtensions(type: KmExtensionType): KmClassExtensionVisitor =
+        extensions.singleOfType(type)
+
+    fun accept(visitor: KmClassVisitor) {
+        visitor.visit(flags, name)
+        typeParameters.forEach { visitor.visitTypeParameter(it.flags, it.name, it.id, it.variance)?.let(it::accept) }
+        supertypes.forEach { visitor.visitSupertype(it.flags)?.let(it::accept) }
+        functions.forEach { visitor.visitFunction(it.flags, it.name)?.let(it::accept) }
+        properties.forEach { visitor.visitProperty(it.flags, it.name, it.getterFlags, it.setterFlags)?.let(it::accept) }
+        typeAliases.forEach { visitor.visitTypeAlias(it.flags, it.name)?.let(it::accept) }
+        constructors.forEach { visitor.visitConstructor(it.flags)?.let(it::accept) }
+        companionObject?.let(visitor::visitCompanionObject)
+        nestedClasses.forEach(visitor::visitNestedClass)
+        enumEntries.forEach(visitor::visitEnumEntry)
+        sealedSubclasses.forEach(visitor::visitSealedSubclass)
+        versionRequirements.forEach { visitor.visitVersionRequirement()?.let(it::accept) }
+        extensions.forEach { visitor.visitExtensions(it.type)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmPackage : KmPackageVisitor(), KmDeclarationContainer {
+    override val functions: MutableList<KmFunction> = ArrayList()
+    override val properties: MutableList<KmProperty> = ArrayList()
+    override val typeAliases: MutableList<KmTypeAlias> = ArrayList(0)
+
+    private val extensions: List<KmPackageExtension> =
+        MetadataExtensions.INSTANCES.map(MetadataExtensions::createPackageExtension)
+
+    override fun visitFunction(flags: Flags, name: String): KmFunctionVisitor =
+        KmFunction(flags, name).addTo(functions)
+
+    override fun visitProperty(flags: Flags, name: String, getterFlags: Flags, setterFlags: Flags): KmPropertyVisitor =
+        KmProperty(flags, name, getterFlags, setterFlags).addTo(properties)
+
+    override fun visitTypeAlias(flags: Flags, name: String): KmTypeAliasVisitor =
+        KmTypeAlias(flags, name).addTo(typeAliases)
+
+    override fun visitExtensions(type: KmExtensionType): KmPackageExtensionVisitor =
+        extensions.singleOfType(type)
+
+    fun accept(visitor: KmPackageVisitor) {
+        functions.forEach { visitor.visitFunction(it.flags, it.name)?.let(it::accept) }
+        properties.forEach { visitor.visitProperty(it.flags, it.name, it.getterFlags, it.setterFlags)?.let(it::accept) }
+        typeAliases.forEach { visitor.visitTypeAlias(it.flags, it.name)?.let(it::accept) }
+        extensions.forEach { visitor.visitExtensions(it.type)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmLambda : KmLambdaVisitor() {
+    lateinit var function: KmFunction
+
+    override fun visitFunction(flags: Flags, name: String): KmFunctionVisitor =
+        KmFunction(flags, name).also { function = it }
+
+    fun accept(visitor: KmLambdaVisitor) {
+        visitor.visitFunction(function.flags, function.name)?.let(function::accept)
+        visitor.visitEnd()
+    }
+}
+
+class KmConstructor(var flags: Flags) : KmConstructorVisitor() {
+    val valueParameters: MutableList<KmValueParameter> = ArrayList()
+    val versionRequirements: MutableList<KmVersionRequirement> = ArrayList(0)
+
+    private val extensions: List<KmConstructorExtension> =
+        MetadataExtensions.INSTANCES.map(MetadataExtensions::createConstructorExtension)
+
+    override fun visitValueParameter(flags: Flags, name: String): KmValueParameterVisitor =
+        KmValueParameter(flags, name).addTo(valueParameters)
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor =
+        KmVersionRequirement().addTo(versionRequirements)
+
+    override fun visitExtensions(type: KmExtensionType): KmConstructorExtensionVisitor =
+        extensions.singleOfType(type)
+
+    fun accept(visitor: KmConstructorVisitor) {
+        valueParameters.forEach { visitor.visitValueParameter(it.flags, it.name)?.let(it::accept) }
+        versionRequirements.forEach { visitor.visitVersionRequirement()?.let(it::accept) }
+        extensions.forEach { visitor.visitExtensions(it.type)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmFunction(
+    var flags: Flags,
+    var name: String
+) : KmFunctionVisitor() {
+    val typeParameters: MutableList<KmTypeParameter> = ArrayList(0)
+    var receiverParameterType: KmType? = null
+    val valueParameters: MutableList<KmValueParameter> = ArrayList()
+    lateinit var returnType: KmType
+    val versionRequirements: MutableList<KmVersionRequirement> = ArrayList(0)
+    var contract: KmContract? = null
+
+    private val extensions: List<KmFunctionExtension> =
+        MetadataExtensions.INSTANCES.map(MetadataExtensions::createFunctionExtension)
+
+    override fun visitTypeParameter(flags: Flags, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor =
+        KmTypeParameter(flags, name, id, variance).addTo(typeParameters)
+
+    override fun visitReceiverParameterType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { receiverParameterType = it }
+
+    override fun visitValueParameter(flags: Flags, name: String): KmValueParameterVisitor =
+        KmValueParameter(flags, name).addTo(valueParameters)
+
+    override fun visitReturnType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { returnType = it }
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor =
+        KmVersionRequirement().addTo(versionRequirements)
+
+    override fun visitContract(): KmContractVisitor =
+        KmContract().also { contract = it }
+
+    override fun visitExtensions(type: KmExtensionType): KmFunctionExtensionVisitor =
+        extensions.singleOfType(type)
+
+    fun accept(visitor: KmFunctionVisitor) {
+        typeParameters.forEach { visitor.visitTypeParameter(it.flags, it.name, it.id, it.variance)?.let(it::accept) }
+        receiverParameterType?.let { visitor.visitReceiverParameterType(it.flags)?.let(it::accept) }
+        valueParameters.forEach { visitor.visitValueParameter(it.flags, it.name)?.let(it::accept) }
+        visitor.visitReturnType(returnType.flags)?.let(returnType::accept)
+        versionRequirements.forEach { visitor.visitVersionRequirement()?.let(it::accept) }
+        contract?.let { visitor.visitContract()?.let(it::accept) }
+        extensions.forEach { visitor.visitExtensions(it.type)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmProperty(
+    var flags: Flags,
+    var name: String,
+    var getterFlags: Flags,
+    var setterFlags: Flags
+) : KmPropertyVisitor() {
+    val typeParameters: MutableList<KmTypeParameter> = ArrayList(0)
+    var receiverParameterType: KmType? = null
+    var setterParameter: KmValueParameter? = null
+    lateinit var returnType: KmType
+    val versionRequirements: MutableList<KmVersionRequirement> = ArrayList(0)
+
+    private val extensions: List<KmPropertyExtension> =
+        MetadataExtensions.INSTANCES.map(MetadataExtensions::createPropertyExtension)
+
+    override fun visitTypeParameter(flags: Flags, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor =
+        KmTypeParameter(flags, name, id, variance).addTo(typeParameters)
+
+    override fun visitReceiverParameterType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { receiverParameterType = it }
+
+    override fun visitSetterParameter(flags: Flags, name: String): KmValueParameterVisitor =
+        KmValueParameter(flags, name).also { setterParameter = it }
+
+    override fun visitReturnType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { returnType = it }
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor =
+        KmVersionRequirement().addTo(versionRequirements)
+
+    override fun visitExtensions(type: KmExtensionType): KmPropertyExtensionVisitor =
+        extensions.singleOfType(type)
+
+    fun accept(visitor: KmPropertyVisitor) {
+        typeParameters.forEach { visitor.visitTypeParameter(it.flags, it.name, it.id, it.variance)?.let(it::accept) }
+        receiverParameterType?.let { visitor.visitReceiverParameterType(it.flags)?.let(it::accept) }
+        setterParameter?.let { visitor.visitSetterParameter(it.flags, it.name)?.let(it::accept) }
+        visitor.visitReturnType(returnType.flags)?.let(returnType::accept)
+        versionRequirements.forEach { visitor.visitVersionRequirement()?.let(it::accept) }
+        extensions.forEach { visitor.visitExtensions(it.type)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmTypeAlias(
+    var flags: Flags,
+    var name: String
+) : KmTypeAliasVisitor() {
+    val typeParameters: MutableList<KmTypeParameter> = ArrayList(0)
+    lateinit var underlyingType: KmType
+    lateinit var expandedType: KmType
+    val annotations: MutableList<KmAnnotation> = ArrayList(0)
+    val versionRequirements: MutableList<KmVersionRequirement> = ArrayList(0)
+
+    override fun visitTypeParameter(flags: Flags, name: String, id: Int, variance: KmVariance): KmTypeParameterVisitor =
+        KmTypeParameter(flags, name, id, variance).addTo(typeParameters)
+
+    override fun visitUnderlyingType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { underlyingType = it }
+
+    override fun visitExpandedType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { expandedType = it }
+
+    override fun visitAnnotation(annotation: KmAnnotation) {
+        annotations.add(annotation)
+    }
+
+    override fun visitVersionRequirement(): KmVersionRequirementVisitor =
+        KmVersionRequirement().addTo(versionRequirements)
+
+    fun accept(visitor: KmTypeAliasVisitor) {
+        typeParameters.forEach { visitor.visitTypeParameter(it.flags, it.name, it.id, it.variance)?.let(it::accept) }
+        visitor.visitUnderlyingType(underlyingType.flags)?.let(underlyingType::accept)
+        visitor.visitExpandedType(expandedType.flags)?.let(expandedType::accept)
+        annotations.forEach(visitor::visitAnnotation)
+        versionRequirements.forEach { visitor.visitVersionRequirement()?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmValueParameter(
+    var flags: Flags,
+    var name: String
+) : KmValueParameterVisitor() {
+    var type: KmType? = null
+    var varargElementType: KmType? = null
+
+    override fun visitType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { type = it }
+
+    override fun visitVarargElementType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { varargElementType = it }
+
+    fun accept(visitor: KmValueParameterVisitor) {
+        type?.let { visitor.visitType(it.flags)?.let(it::accept) }
+        varargElementType?.let { visitor.visitVarargElementType(it.flags)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmTypeParameter(
+    var flags: Flags,
+    var name: String,
+    var id: Int,
+    var variance: KmVariance
+) : KmTypeParameterVisitor() {
+    val upperBounds: MutableList<KmType> = ArrayList(1)
+
+    private val extensions: List<KmTypeParameterExtension> =
+        MetadataExtensions.INSTANCES.map(MetadataExtensions::createTypeParameterExtension)
+
+    override fun visitUpperBound(flags: Flags): KmTypeVisitor =
+        KmType(flags).addTo(upperBounds)
+
+    override fun visitExtensions(type: KmExtensionType): KmTypeParameterExtensionVisitor? =
+        extensions.singleOfType(type)
+
+    fun accept(visitor: KmTypeParameterVisitor) {
+        upperBounds.forEach { visitor.visitUpperBound(it.flags)?.let(it::accept) }
+        extensions.forEach { visitor.visitExtensions(it.type)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmType(var flags: Flags) : KmTypeVisitor() {
+    lateinit var classifier: KmClassifier
+    val arguments: MutableList<KmTypeProjection> = ArrayList(0)
+    var abbreviatedType: KmType? = null
+    var outerType: KmType? = null
+    var flexibleTypeUpperBound: KmFlexibleTypeUpperBound? = null
+
+    private val extensions: List<KmTypeExtension> =
+        MetadataExtensions.INSTANCES.map(MetadataExtensions::createTypeExtension)
+
+    override fun visitClass(name: ClassName) {
+        classifier = KmClassifier.Class(name)
+    }
+
+    override fun visitTypeAlias(name: ClassName) {
+        classifier = KmClassifier.TypeAlias(name)
+    }
+
+    override fun visitTypeParameter(id: Int) {
+        classifier = KmClassifier.TypeParameter(id)
+    }
+
+    override fun visitArgument(flags: Flags, variance: KmVariance): KmTypeVisitor =
+        KmType(flags).also { arguments.add(KmTypeProjection(variance, it)) }
+
+    override fun visitStarProjection() {
+        arguments.add(KmTypeProjection.STAR)
+    }
+
+    override fun visitAbbreviatedType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { abbreviatedType = it }
+
+    override fun visitOuterType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { outerType = it }
+
+    override fun visitFlexibleTypeUpperBound(flags: Flags, typeFlexibilityId: String?): KmTypeVisitor =
+        KmType(flags).also { flexibleTypeUpperBound = KmFlexibleTypeUpperBound(it, typeFlexibilityId) }
+
+    override fun visitExtensions(type: KmExtensionType): KmTypeExtension =
+        extensions.singleOfType(type)
+
+    fun accept(visitor: KmTypeVisitor) {
+        when (val classifier = classifier) {
+            is KmClassifier.Class -> visitor.visitClass(classifier.name)
+            is KmClassifier.TypeParameter -> visitor.visitTypeParameter(classifier.id)
+            is KmClassifier.TypeAlias -> visitor.visitTypeAlias(classifier.name)
+        }
+        arguments.forEach { argument ->
+            if (argument == KmTypeProjection.STAR) visitor.visitStarProjection()
+            else {
+                val (variance, type) = argument
+                if (variance == null || type == null)
+                    throw InconsistentKotlinMetadataException("Variance and type must be set for non-star type projection")
+                visitor.visitArgument(type.flags, variance)?.let(type::accept)
+            }
+        }
+        abbreviatedType?.let { visitor.visitAbbreviatedType(it.flags)?.let(it::accept) }
+        outerType?.let { visitor.visitOuterType(it.flags)?.let(it::accept) }
+        flexibleTypeUpperBound?.let { visitor.visitFlexibleTypeUpperBound(it.type.flags, it.typeFlexibilityId)?.let(it.type::accept) }
+        extensions.forEach { visitor.visitExtensions(it.type)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmVersionRequirement : KmVersionRequirementVisitor() {
+    lateinit var kind: KmVersionRequirementVersionKind
+    lateinit var level: KmVersionRequirementLevel
+    var errorCode: Int? = null
+    var message: String? = null
+    lateinit var version: KmVersion
+
+    override fun visit(kind: KmVersionRequirementVersionKind, level: KmVersionRequirementLevel, errorCode: Int?, message: String?) {
+        this.kind = kind
+        this.level = level
+        this.errorCode = errorCode
+        this.message = message
+    }
+
+    override fun visitVersion(major: Int, minor: Int, patch: Int) {
+        this.version = KmVersion(major, minor, patch)
+    }
+
+    fun accept(visitor: KmVersionRequirementVisitor) {
+        visitor.visit(kind, level, errorCode, message)
+        visitor.visitVersion(version.major, version.minor, version.patch)
+        visitor.visitEnd()
+    }
+}
+
+class KmContract : KmContractVisitor() {
+    val effects: MutableList<KmEffect> = ArrayList(1)
+
+    override fun visitEffect(type: KmEffectType, invocationKind: KmEffectInvocationKind?): KmEffectVisitor =
+        KmEffect(type, invocationKind).addTo(effects)
+
+    fun accept(visitor: KmContractVisitor) {
+        effects.forEach { visitor.visitEffect(it.type, it.invocationKind)?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmEffect(
+    var type: KmEffectType,
+    var invocationKind: KmEffectInvocationKind?
+) : KmEffectVisitor() {
+    val constructorArguments: MutableList<KmEffectExpression> = ArrayList(1)
+    var conclusion: KmEffectExpression? = null
+
+    override fun visitConstructorArgument(): KmEffectExpressionVisitor =
+        KmEffectExpression().addTo(constructorArguments)
+
+    override fun visitConclusionOfConditionalEffect(): KmEffectExpressionVisitor =
+        KmEffectExpression().also { conclusion = it }
+
+    fun accept(visitor: KmEffectVisitor) {
+        constructorArguments.forEach { visitor.visitConstructorArgument()?.let(it::accept) }
+        conclusion?.let { visitor.visitConclusionOfConditionalEffect()?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+class KmEffectExpression : KmEffectExpressionVisitor() {
+    var flags: Flags = flagsOf()
+    var parameterIndex: Int? = null
+    var constantValue: KmConstantValue? = null
+    var isInstanceType: KmType? = null
+    val andArguments: MutableList<KmEffectExpression> = ArrayList(0)
+    val orArguments: MutableList<KmEffectExpression> = ArrayList(0)
+
+    override fun visit(flags: Flags, parameterIndex: Int?) {
+        this.flags = flags
+        this.parameterIndex = parameterIndex
+    }
+
+    override fun visitConstantValue(value: Any?) {
+        constantValue = KmConstantValue(value)
+    }
+
+    override fun visitIsInstanceType(flags: Flags): KmTypeVisitor =
+        KmType(flags).also { isInstanceType = it }
+
+    override fun visitAndArgument(): KmEffectExpressionVisitor =
+        KmEffectExpression().addTo(andArguments)
+
+    override fun visitOrArgument(): KmEffectExpressionVisitor =
+        KmEffectExpression().addTo(orArguments)
+
+    fun accept(visitor: KmEffectExpressionVisitor) {
+        visitor.visit(flags, parameterIndex)
+        constantValue?.let { visitor.visitConstantValue(it.value) }
+        isInstanceType?.let { visitor.visitIsInstanceType(it.flags)?.let(it::accept) }
+        andArguments.forEach { visitor.visitAndArgument()?.let(it::accept) }
+        orArguments.forEach { visitor.visitOrArgument()?.let(it::accept) }
+        visitor.visitEnd()
+    }
+}
+
+sealed class KmClassifier {
+    data class Class(val name: ClassName) : KmClassifier()
+
+    data class TypeParameter(val id: Int) : KmClassifier()
+
+    data class TypeAlias(val name: ClassName) : KmClassifier()
+}
+
+data class KmTypeProjection(var variance: KmVariance?, var type: KmType?) {
+    companion object {
+        @JvmField
+        val STAR = KmTypeProjection(null, null)
+    }
+}
+
+data class KmFlexibleTypeUpperBound(var type: KmType, var typeFlexibilityId: String?)
+
+data class KmVersion(val major: Int, val minor: Int, val patch: Int) {
+    override fun toString(): String = "$major.$minor.$patch"
+}
+
+data class KmConstantValue(val value: Any?)
+
+internal fun <T> T.addTo(collection: MutableCollection<T>): T {
+    collection.add(this)
+    return this
+}

libraries/kotlinx-metadata/src/kotlinx/metadata/util.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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
+
+/**
+ * Assigns the contents of [collection] into this list. After this function completes, the list and [collection]
+ * will have the same size and will consist of the same elements in the same order.
+ *
+ * Useful for rewriting contents of any repeated Km* element represented by a [MutableList], e.g.:
+ *
+ *     val klass: KmClass
+ *     klass.functions.assignFrom(...)
+ *
+ */
+fun <T> MutableList<T>.assignFrom(collection: Collection<T>) {
+    clear()
+    addAll(collection)
+}

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

@@ -23,8 +23,9 @@ import java.io.File
 import kotlin.test.fail
 
 fun compileAndPrintAllFiles(file: File, disposable: Disposable, tmpdir: File, compareWithTxt: Boolean, readWriteAndCompare: Boolean) {
-    val read = StringBuilder()
+    val main = StringBuilder()
-    val readWriteRead = StringBuilder()
+    val afterVisitors = StringBuilder()
+    val afterNodes = StringBuilder()
 
     val kotlinp = Kotlinp(KotlinpSettings(isVerbose = true))
 
@@ -32,18 +33,24 @@ fun compileAndPrintAllFiles(file: File, disposable: Disposable, tmpdir: File, co
         when (outputFile.extension) {
             "kotlin_module" -> {
                 val moduleFile = kotlinp.readModuleFile(outputFile)!!
-                val moduleFile2 = transformModuleFile(moduleFile)
+                val transformedWithVisitors = transformModuleFileWithReadWriteVisitors(moduleFile)
+                val transformedWithNodes = transformModuleFileWithNodes(moduleFile)
 
-                for ((sb, moduleFileToRender) in listOf(read to moduleFile, readWriteRead to moduleFile2)) {
+                for ((sb, moduleFileToRender) in listOf(
+                    main to moduleFile, afterVisitors to transformedWithVisitors, afterNodes to transformedWithNodes
+                )) {
                     sb.appendFileName(outputFile.relativeTo(tmpdir))
                     sb.append(kotlinp.renderModuleFile(moduleFileToRender))
                 }
             }
             "class" -> {
                 val classFile = kotlinp.readClassFile(outputFile)!!
-                val classFile2 = transformClassFile(classFile)
+                val classFile2 = transformClassFileWithReadWriteVisitors(classFile)
+                val classFile3 = transformClassFileWithNodes(classFile)
 
-                for ((sb, classFileToRender) in listOf(read to classFile, readWriteRead to classFile2)) {
+                for ((sb, classFileToRender) in listOf(
+                    main to classFile, afterVisitors to classFile2, afterNodes to classFile3
+                )) {
                     sb.appendFileName(outputFile.relativeTo(tmpdir))
                     sb.append(kotlinp.renderClassFile(classFileToRender))
                 }
@@ -53,11 +60,12 @@ fun compileAndPrintAllFiles(file: File, disposable: Disposable, tmpdir: File, co
     }
 
     if (compareWithTxt) {
-        KotlinTestUtils.assertEqualsToFile(File(file.path.replace(".kt", ".txt")), read.toString())
+        KotlinTestUtils.assertEqualsToFile(File(file.path.replace(".kt", ".txt")), main.toString())
     }
 
     if (readWriteAndCompare) {
-        assertEquals(read.toString(), readWriteRead.toString())
+        assertEquals("Metadata is different after transformation with visitors.", main.toString(), afterVisitors.toString())
+        assertEquals("Metadata is different after transformation with nodes.", main.toString(), afterNodes.toString())
     }
 }
 
@@ -85,7 +93,7 @@ private fun StringBuilder.appendFileName(file: File) {
 // 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 =
+private fun transformClassFileWithReadWriteVisitors(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()
@@ -102,5 +110,22 @@ private fun transformClassFile(classFile: KotlinClassMetadata): KotlinClassMetad
         else -> classFile
     }
 
-private fun transformModuleFile(moduleFile: KotlinModuleMetadata): KotlinModuleMetadata =
+// Reads the class file and writes it back with KmClass/KmFunction/... elements.
+private fun transformClassFileWithNodes(classFile: KotlinClassMetadata): KotlinClassMetadata =
+    when (classFile) {
+        is KotlinClassMetadata.Class ->
+            KotlinClassMetadata.Class.Writer().apply(classFile.toKmClass()::accept).write()
+        is KotlinClassMetadata.FileFacade ->
+            KotlinClassMetadata.FileFacade.Writer().apply(classFile.toKmPackage()::accept).write()
+        is KotlinClassMetadata.SyntheticClass ->
+            KotlinClassMetadata.SyntheticClass.Writer().apply { classFile.toKmLambda()?.accept(this) }.write()
+        is KotlinClassMetadata.MultiFileClassPart ->
+            KotlinClassMetadata.MultiFileClassPart.Writer().apply(classFile.toKmPackage()::accept).write(classFile.facadeClassName)
+        else -> classFile
+    }
+
+private fun transformModuleFileWithReadWriteVisitors(moduleFile: KotlinModuleMetadata): KotlinModuleMetadata =
     KotlinModuleMetadata.Writer().apply(moduleFile::accept).write()
+
+private fun transformModuleFileWithNodes(moduleFile: KotlinModuleMetadata): KotlinModuleMetadata =
+    KotlinModuleMetadata.Writer().apply(moduleFile.toKmModule()::accept).write()