JVM_IR: avoid double encoding of byte strings

IrLibraryFile, ingerited from Klib code, needed types, bodies, strings, signatures encoded as byte strings. When we store this data as class annotations, it is better to store it as protobuf structs, to avoid re-encoding byte streams twice.
2021-09-27 19:41:49 +03:00
parent b61389f6f9
commit 8a459821d0
13 changed files with 686 additions and 326 deletions
@@ -6,7 +6,8 @@
 package org.jetbrains.kotlin.ir.backend.js.ic

 import org.jetbrains.kotlin.backend.common.serialization.IdSignatureDeserializer
-import org.jetbrains.kotlin.backend.common.serialization.IrLibraryFile
+import org.jetbrains.kotlin.backend.common.serialization.IrLibraryBytesSource
+import org.jetbrains.kotlin.backend.common.serialization.IrLibraryFileFromBytes
 import org.jetbrains.kotlin.backend.common.serialization.codedInputStream
 import org.jetbrains.kotlin.backend.common.serialization.signature.IdSignatureDescriptor
 import org.jetbrains.kotlin.config.CompilerConfiguration
@@ -121,7 +122,7 @@ private fun KotlinLibrary.filesAndSigReaders(): List<Pair<String, IdSignatureDes
    for (i in 0 until fileSize) {
        val fileStream = file(i).codedInputStream
        val fileProto = ProtoFile.parseFrom(fileStream, extReg)
-        val sigReader = IdSignatureDeserializer(object : IrLibraryFile() {
+        val sigReader = IdSignatureDeserializer(IrLibraryFileFromBytes(object : IrLibraryBytesSource() {
            private fun err(): Nothing = error("Not supported")
            override fun irDeclaration(index: Int): ByteArray = err()

@@ -134,7 +135,7 @@ private fun KotlinLibrary.filesAndSigReaders(): List<Pair<String, IdSignatureDes
            override fun body(index: Int): ByteArray = err()

            override fun debugInfo(index: Int): ByteArray? = null
-        }, null)
+        }), null)

        result.add(fileProto.fileEntry.name to sigReader)
    }
@@ -104,7 +104,7 @@ abstract class BasicIrModuleDeserializer(

    private fun deserializeIrFile(fileProto: ProtoFile, fileIndex: Int, moduleDeserializer: IrModuleDeserializer, allowErrorNodes: Boolean): IrFile {

-        val fileReader = IrLibraryFileFromKlib(moduleDeserializer.klib, fileIndex)
+        val fileReader = IrLibraryFileFromBytes(IrKlibBytesSource(moduleDeserializer.klib, fileIndex))
        val file = fileReader.createFile(moduleFragment, fileProto)
        val fileStrategy = strategyResolver(file.fileEntry.name)

@@ -8,8 +8,6 @@ package org.jetbrains.kotlin.backend.common.serialization
 import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.ir.util.IdSignature.FileSignature
-import org.jetbrains.kotlin.protobuf.CodedInputStream
-import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
 import org.jetbrains.kotlin.backend.common.serialization.proto.AccessorIdSignature as ProtoAccessorIdSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.CommonIdSignature as ProtoCommonIdSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.CompositeSignature as ProtoCompositeSignature
@@ -19,14 +17,12 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.IdSignature as Pr
 import org.jetbrains.kotlin.backend.common.serialization.proto.LocalSignature as ProtoLocalSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.LoweredIdSignature as ProtoLoweredIdSignature

-class IdSignatureDeserializer(private val fileReader: IrLibraryFile, fileSymbol: IrFileSymbol?) {
-
-    private fun readSignature(index: Int): CodedInputStream = fileReader.signature(index).codedInputStream
+class IdSignatureDeserializer(private val libraryFile: IrLibraryFile, fileSymbol: IrFileSymbol?) {

    private val fileSignature: FileSignature? = fileSymbol?.let { FileSignature(it) }

    private fun loadSignatureProto(index: Int): ProtoIdSignature {
-        return ProtoIdSignature.parseFrom(readSignature(index), extensionRegistryLite)
+        return libraryFile.signature(index)
    }

    private val signatureCache = HashMap<Int, IdSignature>()
@@ -39,8 +35,8 @@ class IdSignatureDeserializer(private val fileReader: IrLibraryFile, fileSymbol:
    }

    private fun deserializePublicIdSignature(proto: ProtoCommonIdSignature): IdSignature.CommonSignature {
-        val pkg = fileReader.deserializeFqName(proto.packageFqNameList)
-        val cls = fileReader.deserializeFqName(proto.declarationFqNameList)
+        val pkg = libraryFile.deserializeFqName(proto.packageFqNameList)
+        val cls = libraryFile.deserializeFqName(proto.declarationFqNameList)
        val memberId = if (proto.hasMemberUniqId()) proto.memberUniqId else null

        return IdSignature.CommonSignature(pkg, cls, memberId, proto.flags)
@@ -49,7 +45,7 @@ class IdSignatureDeserializer(private val fileReader: IrLibraryFile, fileSymbol:
    private fun deserializeAccessorIdSignature(proto: ProtoAccessorIdSignature): IdSignature.AccessorSignature {
        val propertySignature = deserializeIdSignature(proto.propertySignature)
        require(propertySignature is IdSignature.CommonSignature) { "For public accessor corresponding property supposed to be public as well" }
-        val name = fileReader.deserializeString(proto.name)
+        val name = libraryFile.string(proto.name)
        val hash = proto.accessorHashId
        val mask = proto.flags

@@ -78,9 +74,9 @@ class IdSignatureDeserializer(private val fileReader: IrLibraryFile, fileSymbol:
    }

    private fun deserializeLocalIdSignature(proto: ProtoLocalSignature): IdSignature.LocalSignature {
-        val localFqn = fileReader.deserializeFqName(proto.localFqNameList)
+        val localFqn = libraryFile.deserializeFqName(proto.localFqNameList)
        val localHash = if (proto.hasLocalHash()) proto.localHash else null
-        val description = if (proto.hasDebugInfo()) fileReader.deserializeDebugInfo(proto.debugInfo) else null
+        val description = if (proto.hasDebugInfo()) libraryFile.debugInfo(proto.debugInfo) else null
        return IdSignature.LocalSignature(localFqn, localHash, description)
    }

@@ -105,8 +101,4 @@ class IdSignatureDeserializer(private val fileReader: IrLibraryFile, fileSymbol:
    fun signatureToIndexMapping(): Map<IdSignature, Int> {
        return signatureCache.entries.associate { it.value to it.key }
    }
-
-    companion object {
-        private val extensionRegistryLite = ExtensionRegistryLite.newInstance()
-    }
 }
@@ -78,7 +78,7 @@ class IrBodyDeserializer(
    private val builtIns: IrBuiltIns,
    private val allowErrorNodes: Boolean,
    private val irFactory: IrFactory,
-    private val fileReader: IrLibraryFile,
+    private val libraryFile: IrLibraryFile,
    private val declarationDeserializer: IrDeclarationDeserializer,
    private val statementOriginIndex: Map<String, IrStatementOrigin>,
    private val allowErrorStatementOrigins: Boolean, // TODO: support InlinerExpressionLocationHint
@@ -382,9 +382,9 @@ class IrBodyDeserializer(
        start: Int, end: Int, type: IrType
    ): IrErrorExpression {
        require(allowErrorNodes) {
-            "IrErrorExpression($start, $end, \"${fileReader.deserializeString(proto.description)}\") found but error code is not allowed"
+            "IrErrorExpression($start, $end, \"${libraryFile.string(proto.description)}\") found but error code is not allowed"
        }
-        return IrErrorExpressionImpl(start, end, type, fileReader.deserializeString(proto.description))
+        return IrErrorExpressionImpl(start, end, type, libraryFile.string(proto.description))
    }

    private fun deserializeErrorCallExpression(
@@ -392,9 +392,9 @@ class IrBodyDeserializer(
        start: Int, end: Int, type: IrType
    ): IrErrorCallExpression {
        require(allowErrorNodes) {
-            "IrErrorCallExpressionImpl($start, $end, \"${fileReader.deserializeString(proto.description)}\") found but error code is not allowed"
+            "IrErrorCallExpressionImpl($start, $end, \"${libraryFile.string(proto.description)}\") found but error code is not allowed"
        }
-        return IrErrorCallExpressionImpl(start, end, type, fileReader.deserializeString(proto.description)).apply {
+        return IrErrorCallExpressionImpl(start, end, type, libraryFile.string(proto.description)).apply {
            if (proto.hasReceiver()) {
                explicitReceiver = deserializeExpression(proto.receiver)
            }
@@ -672,7 +672,7 @@ class IrBodyDeserializer(
    }

    private fun deserializeLoop(proto: ProtoLoop, loop: IrLoop): IrLoop {
-        val label = if (proto.hasLabel()) fileReader.deserializeString(proto.label) else null
+        val label = if (proto.hasLabel()) libraryFile.string(proto.label) else null
        val body = if (proto.hasBody()) deserializeExpression(proto.body) else null
        val condition = deserializeExpression(proto.condition)

@@ -713,7 +713,7 @@ class IrBodyDeserializer(
            start,
            end,
            type,
-            fileReader.deserializeString(proto.memberName),
+            libraryFile.string(proto.memberName),
            deserializeExpression(proto.receiver)
        )

@@ -774,7 +774,7 @@ class IrBodyDeserializer(
        }

    private fun deserializeBreak(proto: ProtoBreak, start: Int, end: Int, type: IrType): IrBreak {
-        val label = if (proto.hasLabel()) fileReader.deserializeString(proto.label) else null
+        val label = if (proto.hasLabel()) libraryFile.string(proto.label) else null
        val loopId = proto.loopId
        val loop = deserializeLoopHeader(loopId) {
            if (allowErrorLoopIndices) {
@@ -790,7 +790,7 @@ class IrBodyDeserializer(
    }

    private fun deserializeContinue(proto: ProtoContinue, start: Int, end: Int, type: IrType): IrContinue {
-        val label = if (proto.hasLabel()) fileReader.deserializeString(proto.label) else null
+        val label = if (proto.hasLabel()) libraryFile.string(proto.label) else null
        val loopId = proto.loopId
        val loop = deserializeLoopHeader(loopId) {
            if (allowErrorLoopIndices) {
@@ -822,7 +822,7 @@ class IrBodyDeserializer(
            LONG
            -> IrConstImpl.long(start, end, type, proto.long)
            STRING
-            -> IrConstImpl.string(start, end, type, fileReader.deserializeString(proto.string))
+            -> IrConstImpl.string(start, end, type, libraryFile.string(proto.string))
            FLOAT_BITS
            -> IrConstImpl.float(start, end, type, Float.fromBits(proto.floatBits))
            DOUBLE_BITS
@@ -891,7 +891,7 @@ class IrBodyDeserializer(
    }

    fun deserializeIrStatementOrigin(protoName: Int): IrStatementOrigin {
-        return fileReader.deserializeString(protoName).let {
+        return libraryFile.string(protoName).let {
            val componentPrefix = "COMPONENT_"
            when {
                it.startsWith(componentPrefix) -> {
@@ -27,8 +27,6 @@ import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.protobuf.CodedInputStream
-import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
 import org.jetbrains.kotlin.types.Variance
 import kotlin.collections.set
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrAnonymousInit as ProtoAnonymousInit
@@ -61,7 +59,7 @@ class IrDeclarationDeserializer(
    builtIns: IrBuiltIns,
    private val symbolTable: SymbolTable,
    val irFactory: IrFactory,
-    private val fileReader: IrLibraryFile,
+    private val libraryFile: IrLibraryFile,
    private val parent: IrDeclarationParent,
    val allowErrorNodes: Boolean,
    private val deserializeInlineFunctions: Boolean,
@@ -81,7 +79,7 @@ class IrDeclarationDeserializer(
        builtIns,
        allowErrorNodes,
        irFactory,
-        fileReader,
+        libraryFile,
        this,
        statementOriginIndex + additionalStatementOriginIndex,
        allowErrorStatementOrigins,
@@ -89,17 +87,14 @@ class IrDeclarationDeserializer(
    )

    private fun deserializeName(index: Int): Name {
-        val name = fileReader.deserializeString(index)
+        val name = libraryFile.string(index)
        return Name.guessByFirstCharacter(name)
    }

    private val irTypeCache = mutableMapOf<Int, IrType>()

-    private fun readType(index: Int): CodedInputStream =
-        fileReader.type(index).codedInputStream
-
    private fun loadTypeProto(index: Int): ProtoType {
-        return ProtoType.parseFrom(readType(index), extensionRegistry)
+        return libraryFile.type(index)
    }

    fun deserializeNullableIrType(index: Int): IrType? = if (index == -1) null else deserializeIrType(index)
@@ -485,15 +480,12 @@ class IrDeclarationDeserializer(
        }
    }

-    private fun readBody(index: Int): CodedInputStream =
-        fileReader.body(index).codedInputStream
-
    private fun loadStatementBodyProto(index: Int): ProtoStatement {
-        return ProtoStatement.parseFrom(readBody(index), extensionRegistry)
+        return libraryFile.statementBody(index)
    }

    private fun loadExpressionBodyProto(index: Int): ProtoExpression {
-        return ProtoExpression.parseFrom(readBody(index), extensionRegistry)
+        return libraryFile.expressionBody(index)
    }

    fun deserializeExpressionBody(index: Int): IrExpressionBody? {
@@ -748,12 +740,10 @@ class IrDeclarationDeserializer(
        private val allKnownStatementOrigins = IrStatementOrigin::class.nestedClasses.toList()
        private val statementOriginIndex =
            allKnownStatementOrigins.mapNotNull { it.objectInstance as? IrStatementOriginImpl }.associateBy { it.debugName }
-
-        private val extensionRegistry = ExtensionRegistryLite.newInstance()
    }

    fun deserializeIrDeclarationOrigin(protoName: Int): IrDeclarationOriginImpl {
-        val originName = fileReader.deserializeString(protoName)
+        val originName = libraryFile.string(protoName)
        return declarationOriginIndex[originName] ?: object : IrDeclarationOriginImpl(originName) {}
    }

@@ -16,15 +16,18 @@ import org.jetbrains.kotlin.ir.util.NaiveSourceBasedFileEntryImpl
 import org.jetbrains.kotlin.library.IrLibrary
 import org.jetbrains.kotlin.library.encodings.WobblyTF8
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.protobuf.CodedInputStream
 import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
+import org.jetbrains.kotlin.backend.common.serialization.proto.IdSignature as ProtoIdSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrConstructorCall as ProtoConstructorCall
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrDeclaration as ProtoDeclaration
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrExpression as ProtoExpression
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFile as ProtoFile
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrStatement as ProtoStatement
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrType as ProtoType

 class IrFileDeserializer(
    val file: IrFile,
-    private val fileReader: IrLibraryFile,
+    private val libraryFile: IrLibraryFile,
    fileProto: ProtoFile,
    val symbolDeserializer: IrSymbolDeserializer,
    val declarationDeserializer: IrDeclarationDeserializer,
@@ -39,12 +42,9 @@ class IrFileDeserializer(
        }
    }

-    private fun readDeclaration(index: Int): CodedInputStream =
-        fileReader.irDeclaration(index).codedInputStream
-
    private fun loadTopLevelDeclarationProto(idSig: IdSignature): ProtoDeclaration {
        val idSigIndex = reversedSignatureIndex[idSig] ?: error("Not found Idx for $idSig")
-        return ProtoDeclaration.parseFrom(readDeclaration(idSigIndex), ExtensionRegistryLite.newInstance())
+        return libraryFile.declaration(idSigIndex)
    }

    fun deserializeFileImplicitDataIfFirstUse() {
@@ -58,7 +58,7 @@ class IrFileDeserializer(
 class FileDeserializationState(
    val linker: KotlinIrLinker,
    file: IrFile,
-    fileReader: IrLibraryFile,
+    fileReader: IrLibraryFileFromBytes,
    fileProto: ProtoFile,
    deserializeBodies: Boolean,
    allowErrorNodes: Boolean,
@@ -139,6 +139,16 @@ class FileDeserializationState(
 }

 abstract class IrLibraryFile {
+    abstract fun declaration(index: Int): ProtoDeclaration
+    abstract fun type(index: Int): ProtoType
+    abstract fun signature(index: Int): ProtoIdSignature
+    abstract fun string(index: Int): String
+    abstract fun expressionBody(index: Int): ProtoExpression
+    abstract fun statementBody(index: Int): ProtoStatement
+    abstract fun debugInfo(index: Int): String?
+}
+
+abstract class IrLibraryBytesSource {
    abstract fun irDeclaration(index: Int): ByteArray
    abstract fun type(index: Int): ByteArray
    abstract fun signature(index: Int): ByteArray
@@ -147,7 +157,32 @@ abstract class IrLibraryFile {
    abstract fun debugInfo(index: Int): ByteArray?
 }

-class IrLibraryFileFromKlib(private val klib: IrLibrary, private val fileIndex: Int) : IrLibraryFile() {
+class IrLibraryFileFromBytes(private val bytesSource: IrLibraryBytesSource) : IrLibraryFile() {
+
+    override fun declaration(index: Int): ProtoDeclaration =
+        ProtoDeclaration.parseFrom(bytesSource.irDeclaration(index).codedInputStream, extensionRegistryLite)
+
+    override fun type(index: Int): ProtoType = ProtoType.parseFrom(bytesSource.type(index).codedInputStream, extensionRegistryLite)
+
+    override fun signature(index: Int): ProtoIdSignature =
+        ProtoIdSignature.parseFrom(bytesSource.signature(index).codedInputStream, extensionRegistryLite)
+
+    override fun string(index: Int): String = WobblyTF8.decode(bytesSource.string(index))
+
+    override fun expressionBody(index: Int): ProtoExpression =
+        ProtoExpression.parseFrom(bytesSource.body(index).codedInputStream, extensionRegistryLite)
+
+    override fun statementBody(index: Int): ProtoStatement =
+        ProtoStatement.parseFrom(bytesSource.body(index).codedInputStream, extensionRegistryLite)
+
+    override fun debugInfo(index: Int): String? = bytesSource.debugInfo(index)?.let { WobblyTF8.decode(it) }
+
+    companion object {
+        val extensionRegistryLite: ExtensionRegistryLite = ExtensionRegistryLite.newInstance()
+    }
+}
+
+class IrKlibBytesSource(private val klib: IrLibrary, private val fileIndex: Int) : IrLibraryBytesSource() {
    override fun irDeclaration(index: Int): ByteArray = klib.irDeclaration(index, fileIndex)
    override fun type(index: Int): ByteArray = klib.type(index, fileIndex)
    override fun signature(index: Int): ByteArray = klib.signature(index, fileIndex)
@@ -156,11 +191,8 @@ class IrLibraryFileFromKlib(private val klib: IrLibrary, private val fileIndex:
    override fun debugInfo(index: Int): ByteArray? = klib.debugInfo(index, fileIndex)
 }

-internal fun IrLibraryFile.deserializeString(index: Int): String = WobblyTF8.decode(string(index))
-internal fun IrLibraryFile.deserializeDebugInfo(index: Int): String? = debugInfo(index)?.let { WobblyTF8.decode(it) }
-
 internal fun IrLibraryFile.deserializeFqName(fqn: List<Int>): String =
-    fqn.joinToString(".", transform = ::deserializeString)
+    fqn.joinToString(".", transform = ::string)

 fun IrLibraryFile.createFile(module: IrModuleFragment, fileProto: ProtoFile): IrFile {
    val fileName = fileProto.fileEntry.name
@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable

 class IrSymbolDeserializer(
    val symbolTable: ReferenceSymbolTable,
-    val fileReader: IrLibraryFile,
+    val libraryFile: IrLibraryFile,
    val fileSymbol: IrFileSymbol,
    val actuals: List<Actual>,
    val enqueueLocalTopLevelDeclaration: (IdSignature) -> Unit,
@@ -103,7 +103,7 @@ class IrSymbolDeserializer(
        }
    }

-    val signatureDeserializer = IdSignatureDeserializer(fileReader, fileSymbol)
+    val signatureDeserializer = IdSignatureDeserializer(libraryFile, fileSymbol)

    fun deserializeIdSignature(index: Int): IdSignature {
        return signatureDeserializer.deserializeIdSignature(index)
@@ -23,13 +23,12 @@ import org.jetbrains.kotlin.library.impl.DeclarationIrTableMemoryReader
 import org.jetbrains.kotlin.library.impl.IrArrayMemoryReader
 import org.jetbrains.kotlin.library.impl.IrLongArrayMemoryReader
 import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
-import org.jetbrains.kotlin.backend.common.serialization.proto.IrDeclaration as ProtoIrDeclaration
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFile as ProtoIrFile

 class IcFileDeserializer(
    val linker: JsIrLinker,
    private val file: IrFile,
-    originalFileReader: IrLibraryFile,
+    originalFileReader: IrLibraryFileFromBytes,
    fileProto: ProtoIrFile,
    deserializeBodies: Boolean,
    allowErrorNodes: Boolean,
@@ -115,7 +114,7 @@ class IcFileDeserializer(

    // IC data processing starts here

-    private val icFileReader = FileReaderFromSerializedIrFile(icFileData.file)
+    private val icFileReader = IrLibraryFileFromBytes(FileReaderFromSerializedIrFile(icFileData.file))

    val symbolDeserializer = IrSymbolDeserializer(
        linker.symbolTable,
@@ -195,8 +194,7 @@ class IcFileDeserializer(
        cachedDeclaration(idSig)?.let { return it }

        val idSigIndex = reversedSignatureIndex[idSig] ?: return null
-        val declarationStream = icFileReader.irDeclaration(idSigIndex).codedInputStream
-        val declarationProto = ProtoIrDeclaration.parseFrom(declarationStream, extensionRegistryLite)
+        val declarationProto = icFileReader.declaration(idSigIndex)
        return declarationDeserializer.deserializeDeclaration(declarationProto)
    }

@@ -221,8 +219,7 @@ class IcFileDeserializer(
        }

        reversedSignatureIndex[maybeTopLevel]?.let { idSigIndex ->
-            val declarationStream = icFileReader.irDeclaration(idSigIndex).codedInputStream
-            val declarationProto = ProtoIrDeclaration.parseFrom(declarationStream, extensionRegistryLite)
+            val declarationProto = icFileReader.declaration(idSigIndex)
            return declarationDeserializer.deserializeDeclaration(declarationProto)
        }

@@ -250,9 +247,7 @@ class IcFileDeserializer(
    }
 }

-private val extensionRegistryLite = ExtensionRegistryLite.newInstance()
-
-private class FileReaderFromSerializedIrFile(val irFile: SerializedIrFile) : IrLibraryFile() {
+private class FileReaderFromSerializedIrFile(val irFile: SerializedIrFile) : IrLibraryBytesSource() {
    private val declarationReader = DeclarationIrTableMemoryReader(irFile.declarations)
    private val typeReader = IrArrayMemoryReader(irFile.types)
    private val signatureReader = IrArrayMemoryReader(irFile.signatures)
@@ -127,7 +127,7 @@ class IcModuleDeserializer(
        allowErrorNodes: Boolean
    ): IrFile {

-        val fileReader = IrLibraryFileFromKlib(moduleDeserializer.klib, fileIndex)
+        val fileReader = IrLibraryFileFromBytes(IrKlibBytesSource(moduleDeserializer.klib, fileIndex))
        val file = fileReader.createFile(moduleFragment, fileProto)
        val fileStrategy = strategyResolver(file.fileEntry.name)

@@ -16,11 +16,11 @@ message XStatementOrExpression {

 message AuxTables {
  // TODO: optimize the representation.
-  repeated bytes type = 1;
-  repeated bytes signature = 2;
-  repeated bytes string = 3;
-  repeated bytes body = 4;
-  repeated bytes debug_info = 5;
+  repeated common.serialization.proto.IrType type = 1;
+  repeated common.serialization.proto.IdSignature signature = 2;
+  repeated string string = 3;
+  repeated XStatementOrExpression body = 4;
+  repeated string debug_info = 5;
 }

 message JvmIrFile {
@@ -54,11 +54,16 @@ class JvmIrSerializerSession(

    private fun serializeAuxTables(): JvmIr.AuxTables {
        val proto = JvmIr.AuxTables.newBuilder()
-        protoTypeArray.forEach { proto.addType(it.toByteString()) }
-        protoIdSignatureArray.forEach { proto.addSignature(it.toByteString()) }
-        protoStringArray.forEach { proto.addString(ByteString.copyFromUtf8(it)) }
-        protoBodyArray.forEach { proto.addBody(ByteString.copyFrom(it.toByteArray())) }
-        protoDebugInfoArray.forEach { proto.addDebugInfo(ByteString.copyFromUtf8(it)) }
+        protoTypeArray.forEach(proto::addType)
+        protoIdSignatureArray.forEach(proto::addSignature)
+        protoStringArray.forEach(proto::addString)
+        protoBodyArray.forEach { proto.addBody(it.toProto()) }
+        protoDebugInfoArray.forEach(proto::addDebugInfo)
        return proto.build()
    }
+
+    fun XStatementOrExpression.toProto(): JvmIr.XStatementOrExpression = when (this) {
+        is XStatementOrExpression.XStatement -> JvmIr.XStatementOrExpression.newBuilder().setStatement(toProtoStatement()).build()
+        is XStatementOrExpression.XExpression -> JvmIr.XStatementOrExpression.newBuilder().setExpression(toProtoExpression()).build()
+    }
 }
@@ -29,7 +29,11 @@ import org.jetbrains.kotlin.ir.types.IrTypeSystemContext
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
-import org.jetbrains.kotlin.protobuf.ByteString
+import org.jetbrains.kotlin.backend.common.serialization.proto.IdSignature as ProtoIdSignature
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrDeclaration as ProtoDeclaration
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrExpression as ProtoExpression
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrStatement as ProtoStatement
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrType as ProtoType

 fun deserializeClassFromByteArray(
    byteArray: ByteArray,
@@ -40,7 +44,7 @@ fun deserializeClassFromByteArray(
 ) {
    val irBuiltIns = stubGenerator.irBuiltIns
    val symbolTable = stubGenerator.symbolTable
-    val irProto = JvmIr.JvmIrClass.parseFrom(byteArray)
+    val irProto = JvmIr.JvmIrClass.parseFrom(byteArray.codedInputStream)
    val irLibraryFile = IrLibraryFileFromAnnotation(
        irProto.auxTables.typeList,
        irProto.auxTables.signatureList,
@@ -98,7 +102,7 @@ fun deserializeIrFileFromByteArray(
 ) {
    val irBuiltIns = stubGenerator.irBuiltIns
    val symbolTable = stubGenerator.symbolTable
-    val irProto = JvmIr.JvmIrFile.parseFrom(byteArray)
+    val irProto = JvmIr.JvmIrFile.parseFrom(byteArray.codedInputStream)
    val irLibraryFile = IrLibraryFileFromAnnotation(
        irProto.auxTables.typeList,
        irProto.auxTables.signatureList,
@@ -151,21 +155,26 @@ fun deserializeIrFileFromByteArray(
 }

 private class IrLibraryFileFromAnnotation(
-    private val types: List<ByteString>,
-    private val signatures: List<ByteString>,
-    private val strings: List<ByteString>,
-    private val bodies: List<ByteString>,
-    private val debugInfo: List<ByteString>
+    private val types: List<ProtoType>,
+    private val signatures: List<ProtoIdSignature>,
+    private val strings: List<String>,
+    private val bodies: List<JvmIr.XStatementOrExpression>,
+    private val debugInfo: List<String>
 ) : IrLibraryFile() {
-    override fun irDeclaration(index: Int): ByteArray {
+    override fun declaration(index: Int): ProtoDeclaration {
        error("This method is never supposed to be called")
    }

-    override fun type(index: Int): ByteArray = types[index].toByteArray()
-    override fun signature(index: Int): ByteArray = signatures[index].toByteArray()
-    override fun string(index: Int): ByteArray = strings[index].toByteArray()
-    override fun body(index: Int): ByteArray = bodies[index].toByteArray()
-    override fun debugInfo(index: Int): ByteArray = debugInfo[index].toByteArray()
+    override fun type(index: Int): ProtoType = types[index]
+    override fun signature(index: Int): ProtoIdSignature = signatures[index]
+    override fun string(index: Int): String = strings[index]
+    override fun debugInfo(index: Int): String = debugInfo[index]
+
+    override fun expressionBody(index: Int): ProtoExpression =
+        bodies[index].also { require(it.hasExpression()) }.expression
+
+    override fun statementBody(index: Int): ProtoStatement =
+        bodies[index].also { require(it.hasStatement()) }.statement
 }

 private fun referencePublicSymbol(