From baaecaa052ef396b6eda25ae534b8fadbb18370e Mon Sep 17 00:00:00 2001
From: Georgy Bronnikov <georgy.bronnikov@jetbrains.com>
Date: Mon, 1 Mar 2021 12:54:16 +0300
Subject: [PATCH] JVM_IR: write serialized IR to an annotation

---
 .../backend/jvm/codegen/ClassCodegen.kt       | 27 +++++++++++++++----
 .../kotlin/load/java/JvmAnnotationNames.java  |  1 +
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt
index 8b1f1754fca..78a22fd9298 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt
@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.ir.types.classifierOrFail
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
+import org.jetbrains.kotlin.metadata.jvm.deserialization.BitEncoding
 import org.jetbrains.kotlin.metadata.jvm.serialization.JvmStringTable
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.protobuf.MessageLite
@@ -241,6 +242,11 @@ class ClassCodegen private constructor(
             entry is MultifileFacadeFileEntry -> KotlinClassHeader.Kind.MULTIFILE_CLASS
             else -> KotlinClassHeader.Kind.SYNTHETIC_CLASS
         }
+        val serializedIr = when (metadata) {
+            is MetadataSource.Class -> metadata.serializedIr
+            is MetadataSource.File -> metadata.serializedIr
+            else -> null
+        }
 
         val isMultifileClassOrPart = kind == KotlinClassHeader.Kind.MULTIFILE_CLASS || kind == KotlinClassHeader.Kind.MULTIFILE_CLASS_PART
 
@@ -252,15 +258,15 @@ class ClassCodegen private constructor(
             extraFlags = extraFlags or JvmAnnotationNames.METADATA_SCRIPT_FLAG
         }
 
-        writeKotlinMetadata(visitor, state, kind, extraFlags) {
+        writeKotlinMetadata(visitor, state, kind, extraFlags) { av ->
             if (metadata != null) {
                 metadataSerializer.serialize(metadata)?.let { (proto, stringTable) ->
-                    DescriptorAsmUtil.writeAnnotationData(it, proto, stringTable)
+                    DescriptorAsmUtil.writeAnnotationData(av, proto, stringTable)
                 }
             }
 
             if (entry is MultifileFacadeFileEntry) {
-                val arv = it.visitArray(JvmAnnotationNames.METADATA_DATA_FIELD_NAME)
+                val arv = av.visitArray(JvmAnnotationNames.METADATA_DATA_FIELD_NAME)
                 for (partFile in entry.partFiles) {
                     val fileClass = partFile.declarations.singleOrNull { it.isFileClass } as IrClass?
                     if (fileClass != null) arv.visit(null, typeMapper.mapClass(fileClass).internalName)
@@ -269,14 +275,16 @@ class ClassCodegen private constructor(
             }
 
             if (facadeClassName != null) {
-                it.visit(JvmAnnotationNames.METADATA_MULTIFILE_CLASS_NAME_FIELD_NAME, facadeClassName.internalName)
+                av.visit(JvmAnnotationNames.METADATA_MULTIFILE_CLASS_NAME_FIELD_NAME, facadeClassName.internalName)
             }
 
             if (irClass in context.classNameOverride) {
                 val isFileClass = isMultifileClassOrPart || kind == KotlinClassHeader.Kind.FILE_FACADE
                 assert(isFileClass) { "JvmPackageName is not supported for classes: ${irClass.render()}" }
-                it.visit(JvmAnnotationNames.METADATA_PACKAGE_NAME_FIELD_NAME, irClass.fqNameWhenAvailable!!.parent().asString())
+                av.visit(JvmAnnotationNames.METADATA_PACKAGE_NAME_FIELD_NAME, irClass.fqNameWhenAvailable!!.parent().asString())
             }
+
+            serializedIr?.let { storeSerializedIr(av, it) }
         }
     }
 
@@ -562,3 +570,12 @@ private val Modality.flags: Int
 
 private val DescriptorVisibility.flags: Int
     get() = DescriptorAsmUtil.getVisibilityAccessFlag(this) ?: throw AssertionError("Unsupported visibility $this")
+
+private fun storeSerializedIr(av: AnnotationVisitor, serializedIr: ByteArray) {
+    val serializedIrParts = BitEncoding.encodeBytes(serializedIr)
+    val partsVisitor = av.visitArray(JvmAnnotationNames.METADATA_SERIALIZED_IR_FIELD_NAME)
+    for (part in serializedIrParts) {
+        partsVisitor.visit(null, part)
+    }
+    partsVisitor.visitEnd()
+}
diff --git a/core/compiler.common.jvm/src/org/jetbrains/kotlin/load/java/JvmAnnotationNames.java b/core/compiler.common.jvm/src/org/jetbrains/kotlin/load/java/JvmAnnotationNames.java
index 6ec7c6317ab..edcd3687bb0 100644
--- a/core/compiler.common.jvm/src/org/jetbrains/kotlin/load/java/JvmAnnotationNames.java
+++ b/core/compiler.common.jvm/src/org/jetbrains/kotlin/load/java/JvmAnnotationNames.java
@@ -37,6 +37,7 @@ public final class JvmAnnotationNames {
     public static final String METADATA_PACKAGE_NAME_FIELD_NAME = "pn";
     public static final String METADATA_MULTIFILE_CLASS_NAME_FIELD_NAME = METADATA_EXTRA_STRING_FIELD_NAME;
     public static final String METADATA_EXTRA_INT_FIELD_NAME = "xi";
+    public static final String METADATA_SERIALIZED_IR_FIELD_NAME = "si";
 
     public static final int METADATA_MULTIFILE_PARTS_INHERIT_FLAG = 1 << 0;
     public static final int METADATA_PRE_RELEASE_FLAG = 1 << 1;