KLIB: serialization support for global signatures and lowered IR features

Including:
- returnable blocks
- raw function references
- standalone type and value parameters
- file symbols

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/IrBodyDeserializer.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.backend.common.serialization
 import org.jetbrains.kotlin.backend.common.serialization.encodings.*
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrConst.ValueCase.*
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrOperation.OperationCase.*
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrReturnableBlockReturn
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrStatement.StatementCase
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrVarargElement.VarargElementCase
 import org.jetbrains.kotlin.ir.IrElement
@@ -17,6 +18,7 @@ import org.jetbrains.kotlin.ir.descriptors.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.*
+import org.jetbrains.kotlin.ir.symbols.impl.IrReturnableBlockSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.*
 import org.jetbrains.kotlin.ir.util.parentAsClass
@@ -41,6 +43,8 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.IrErrorExpression
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrExpression as ProtoExpression
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFunctionExpression as ProtoFunctionExpression
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFunctionReference as ProtoFunctionReference
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrRawFunctionReference as ProtoRawFunctionReference
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrReturnableBlock as ProtoReturnableBlock
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrGetClass as ProtoGetClass
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrGetEnumValue as ProtoGetEnumValue
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrGetField as ProtoGetField
@@ -75,6 +79,8 @@ class IrBodyDeserializer(
     private val irFactory: IrFactory,
     private val fileReader: IrLibraryFile,
     private val declarationDeserializer: IrDeclarationDeserializer,
+    private val statementOriginIndex: Map<String, IrStatementOrigin>,
+    private val allowErrorStatementOrigins: Boolean,
 ) {
 
     private val fileLoops = mutableMapOf<Int, IrLoop>()
@@ -156,6 +162,37 @@ class IrBodyDeserializer(
         return IrBlockImpl(start, end, type, origin, statements)
     }
 
+    private val returnableBlockStack = ArrayDeque<IrReturnableBlockSymbol>()
+
+    private inline fun <T> withReturnableBlock(symbol: IrReturnableBlockSymbol, fn: () -> T): T {
+        returnableBlockStack.addLast(symbol)
+        try {
+            return fn()
+        } finally {
+            returnableBlockStack.removeLast()
+        }
+    }
+
+    private fun cntToReturnableBlockSymbol(upCnt: Int) = returnableBlockStack[returnableBlockStack.size - upCnt]
+
+    private fun deserializeReturnableBlock(proto: ProtoReturnableBlock, start: Int, end: Int, type: IrType): IrReturnableBlock {
+
+        val symbol = IrReturnableBlockSymbolImpl()
+
+        val inlineFunctionSymbol = if (proto.hasInlineFunctionSymbol()) {
+            declarationDeserializer.deserializeIrSymbolAndRemap(proto.inlineFunctionSymbol) as IrFunctionSymbol
+        } else null
+
+        val statements = withReturnableBlock(symbol) {
+            proto.statementList.map {
+                deserializeStatement(it) as IrStatement
+            }
+        }
+
+        val origin = if (proto.hasOriginName()) deserializeIrStatementOrigin(proto.originName) else null
+        return IrReturnableBlockImpl(start, end, type, symbol, origin, statements, inlineFunctionSymbol)
+    }
+
     private fun deserializeMemberAccessCommon(access: IrMemberAccessExpression<*>, proto: ProtoMemberAccessCommon) {
 
         proto.valueArgumentList.mapIndexed { i, arg ->
@@ -387,6 +424,14 @@ class IrBodyDeserializer(
         return callable
     }
 
+    private fun deserializeRawFunctionReference(
+        proto: ProtoRawFunctionReference,
+        start: Int, end: Int, type: IrType
+    ): IrRawFunctionReference {
+        val symbol = declarationDeserializer.deserializeIrSymbolAndRemap(proto.symbol) as IrFunctionSymbol
+        return IrRawFunctionReferenceImpl(start, end, type, symbol)
+    }
+
     private fun deserializeGetClass(proto: ProtoGetClass, start: Int, end: Int, type: IrType): IrGetClass {
         val argument = deserializeExpression(proto.argument)
         return IrGetClassImpl(start, end, type, argument)
@@ -494,6 +539,12 @@ class IrBodyDeserializer(
         return IrReturnImpl(start, end, builtIns.nothingType, symbol, value)
     }
 
+    private fun deserializeReturnableBlockReturn(proto: IrReturnableBlockReturn, start: Int, end: Int): IrReturn {
+        val symbol = cntToReturnableBlockSymbol(proto.upCnt)
+        val value = deserializeExpression(proto.value)
+        return IrReturnImpl(start, end, builtIns.nothingType, symbol, value)
+    }
+
     private fun deserializeSetField(proto: ProtoSetField, start: Int, end: Int): IrSetField {
         val access = proto.fieldAccess
         val symbol = declarationDeserializer.deserializeIrSymbolAndRemap(access.symbol) as IrFieldSymbol
@@ -803,6 +854,9 @@ class IrBodyDeserializer(
             ERROR_EXPRESSION -> deserializeErrorExpression(proto.errorExpression, start, end, type)
             ERROR_CALL_EXPRESSION -> deserializeErrorCallExpression(proto.errorCallExpression, start, end, type)
             OPERATION_NOT_SET -> error("Expression deserialization not implemented: ${proto.operationCase}")
+            RAW_FUNCTION_REFERENCE -> deserializeRawFunctionReference(proto.rawFunctionReference, start, end, type)
+            RETURNABLE_BLOCK -> deserializeReturnableBlock(proto.returnableBlock, start, end, type)
+            RETURNABLE_BLOCK_RETURN -> deserializeReturnableBlockReturn(proto.returnableBlockReturn, start, end)
         }
 
     fun deserializeExpression(proto: ProtoExpression): IrExpression {
@@ -823,16 +877,8 @@ class IrBodyDeserializer(
                 it.startsWith(componentPrefix) -> {
                     IrStatementOrigin.COMPONENT_N.withIndex(it.removePrefix(componentPrefix).toInt())
                 }
-                else -> statementOriginIndex[it] ?: error("Unexpected statement origin: $it")
+                else -> statementOriginIndex[it] ?: if (allowErrorStatementOrigins) object : IrStatementOriginImpl(it) {} else error("Unexpected statement origin: $it")
             }
         }
     }
-
-    companion object {
-
-        private val allKnownStatementOrigins = IrStatementOrigin::class.nestedClasses.toList()
-
-        private val statementOriginIndex =
-            allKnownStatementOrigins.mapNotNull { it.objectInstance as? IrStatementOriginImpl }.associateBy { it.debugName }
-    }
 }

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/IrDeclarationDeserializer.kt

@@ -19,6 +19,8 @@ import org.jetbrains.kotlin.ir.expressions.IrBlockBody
 import org.jetbrains.kotlin.ir.expressions.IrBody
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.IrStatementOriginImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.symbols.impl.IrPublicSymbolBase
@@ -56,23 +58,24 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.IrTypeParameter a
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrValueParameter as ProtoValueParameter
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrVariable as ProtoVariable
 
-
 class IrDeclarationDeserializer(
     builtIns: IrBuiltIns,
     private val symbolTable: SymbolTable,
-    private val irFactory: IrFactory,
+    val irFactory: IrFactory,
     private val fileReader: IrLibraryFile,
-    file: IrFile,
-    private val allowErrorNodes: Boolean,
+    private val file: IrFile,
+    val allowErrorNodes: Boolean,
     private val deserializeInlineFunctions: Boolean,
     private var deserializeBodies: Boolean,
     val symbolDeserializer: IrSymbolDeserializer,
     private val platformFakeOverrideClassFilter: FakeOverrideClassFilter,
     private val fakeOverrideBuilder: FakeOverrideBuilder,
     private val skipMutableState: Boolean = false,
+    additionalStatementOriginIndex: Map<String, IrStatementOrigin> = emptyMap(),
+    allowErrorStatementOrigins: Boolean = false,
 ) {
 
-    val bodyDeserializer = IrBodyDeserializer(builtIns, allowErrorNodes, irFactory, fileReader, this)
+    val bodyDeserializer = IrBodyDeserializer(builtIns, allowErrorNodes, irFactory, fileReader, this, statementOriginIndex + additionalStatementOriginIndex, allowErrorStatementOrigins)
 
     private fun deserializeName(index: Int): Name {
         val name = fileReader.deserializeString(index)
@@ -699,6 +702,11 @@ class IrDeclarationDeserializer(
         private val allKnownDeclarationOrigins = IrDeclarationOrigin::class.nestedClasses.toList()
         private val declarationOriginIndex =
             allKnownDeclarationOrigins.map { it.objectInstance as IrDeclarationOriginImpl }.associateBy { it.name }
+
+
+        private val allKnownStatementOrigins = IrStatementOrigin::class.nestedClasses.toList()
+        private val statementOriginIndex =
+            allKnownStatementOrigins.mapNotNull { it.objectInstance as? IrStatementOriginImpl }.associateBy { it.debugName }
     }
 
     fun deserializeIrDeclarationOrigin(protoName: Int): IrDeclarationOriginImpl {
@@ -714,9 +722,9 @@ class IrDeclarationDeserializer(
             IR_CLASS -> deserializeIrClass(proto.irClass)
             IR_FUNCTION -> deserializeIrFunction(proto.irFunction)
             IR_PROPERTY -> deserializeIrProperty(proto.irProperty)
-            IR_TYPE_PARAMETER -> error("Unreachable execution Type Parameter") // deserializeIrTypeParameter(proto.irTypeParameter)
+            IR_TYPE_PARAMETER -> deserializeIrTypeParameter(proto.irTypeParameter, proto.irTypeParameter.index, proto.irTypeParameter.isGlobal)
             IR_VARIABLE -> deserializeIrVariable(proto.irVariable)
-            IR_VALUE_PARAMETER -> error("Unreachable execution Value Parameter") // deserializeIrValueParameter(proto.irValueParameter)
+            IR_VALUE_PARAMETER -> deserializeIrValueParameter(proto.irValueParameter, proto.irValueParameter.index)
             IR_ENUM_ENTRY -> deserializeIrEnumEntry(proto.irEnumEntry)
             IR_LOCAL_DELEGATED_PROPERTY -> deserializeIrLocalDelegatedProperty(proto.irLocalDelegatedProperty)
             IR_TYPE_ALIAS -> deserializeIrTypeAlias(proto.irTypeAlias)

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/IrFileDeserializer.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
+import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.symbols.impl.IrFileSymbolImpl
 import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.ir.util.NaiveSourceBasedFileEntryImpl
@@ -70,7 +71,12 @@ class FileDeserializationState(
 
     val symbolDeserializer =
         IrSymbolDeserializer(
-            linker.symbolTable, fileReader, fileProto.actualList, ::addIdSignature, linker::handleExpectActualMapping
+            linker.symbolTable,
+            fileReader,
+            file.symbol,
+            fileProto.actualList,
+            ::addIdSignature,
+            linker::handleExpectActualMapping,
         ) { idSig, symbolKind ->
             assert(idSig.isPublic)
 

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/IrFileSerializer.kt

@@ -33,6 +33,8 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.Actual as ProtoAc
 import org.jetbrains.kotlin.backend.common.serialization.proto.FieldAccessCommon as ProtoFieldAccessCommon
 import org.jetbrains.kotlin.backend.common.serialization.proto.FileEntry as ProtoFileEntry
 import org.jetbrains.kotlin.backend.common.serialization.proto.FileLocalIdSignature as ProtoFileLocalIdSignature
+import org.jetbrains.kotlin.backend.common.serialization.proto.ScopeLocalIdSignature as ProtoScopeLocalIdSignature
+import org.jetbrains.kotlin.backend.common.serialization.proto.FileSignature as ProtoFileSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.IdSignature as ProtoIdSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrAnonymousInit as ProtoAnonymousInit
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrBlock as ProtoBlock
@@ -68,6 +70,8 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.IrFunction as Pro
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFunctionBase as ProtoFunctionBase
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFunctionExpression as ProtoFunctionExpression
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFunctionReference as ProtoFunctionReference
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrRawFunctionReference as ProtoRawFunctionReference
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrReturnableBlock as ProtoReturnableBlock
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrGetClass as ProtoGetClass
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrGetEnumValue as ProtoGetEnumValue
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrGetField as ProtoGetField
@@ -81,6 +85,7 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.IrOperation as Pr
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrProperty as ProtoProperty
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrPropertyReference as ProtoPropertyReference
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrReturn as ProtoReturn
+import org.jetbrains.kotlin.backend.common.serialization.proto.IrReturnableBlockReturn as ProtoReturnableBlockReturn
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrSetField as ProtoSetField
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrSetValue as ProtoSetValue
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrSimpleType as ProtoSimpleType
@@ -107,6 +112,7 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.Loop as ProtoLoop
 import org.jetbrains.kotlin.backend.common.serialization.proto.MemberAccessCommon as ProtoMemberAccessCommon
 import org.jetbrains.kotlin.backend.common.serialization.proto.NullableIrExpression as ProtoNullableIrExpression
 import org.jetbrains.kotlin.backend.common.serialization.proto.PublicIdSignature as ProtoPublicIdSignature
+import org.jetbrains.kotlin.backend.common.serialization.proto.LoweredIdSignature as ProtoLoweredIdSignature
 
 open class IrFileSerializer(
     val messageLogger: IrMessageLogger,
@@ -114,7 +120,10 @@ open class IrFileSerializer(
     private val expectDescriptorToSymbol: MutableMap<DeclarationDescriptor, IrSymbol>,
     private val bodiesOnlyForInlines: Boolean = false,
     private val skipExpects: Boolean = false,
-    private val skipMutableState: Boolean = false, // required for JS IC caches
+    // required for JS IC caches
+    private val skipMutableState: Boolean = false,
+    private val allowNullTypes: Boolean = false,
+    private val allowErrorStatementOrigins: Boolean = false, // TODO: support InlinerExpressionLocationHint
 ) {
     private val loopIndex = mutableMapOf<IrLoop, Int>()
     private var currentLoopIndex = 0
@@ -174,7 +183,8 @@ open class IrFileSerializer(
 
     fun serializeIrDeclarationOrigin(origin: IrDeclarationOrigin): Int = serializeString((origin as IrDeclarationOriginImpl).name)
 
-    private fun serializeIrStatementOrigin(origin: IrStatementOrigin): Int = serializeString((origin as IrStatementOriginImpl).debugName)
+    private fun serializeIrStatementOrigin(origin: IrStatementOrigin): Int =
+        serializeString((origin as? IrStatementOriginImpl)?.debugName ?: if (allowErrorStatementOrigins) "error_origin" else error("Unable to serialize origin ${origin.javaClass.name}"))
 
     private fun serializeCoordinates(start: Int, end: Int): Long = BinaryCoordinates.encode(start, end)
 
@@ -228,13 +238,50 @@ open class IrFileSerializer(
 
     private fun serializeScopeLocalSignature(signature: IdSignature.ScopeLocalDeclaration): Int = signature.id
 
-    private fun serializeIdSignature(idSignature: IdSignature): ProtoIdSignature {
+    private fun serializePrivateSignature(signature: IdSignature.GlobalFileLocalSignature): ProtoFileLocalIdSignature {
+        val proto = ProtoFileLocalIdSignature.newBuilder()
+
+        proto.container = protoIdSignature(signature.container)
+        proto.localId = signature.id
+        proto.file = serializeString(signature.filePath)
+
+        return proto.build()
+    }
+
+    private fun serializeScopeLocalSignature(signature: IdSignature.GlobalScopeLocalDeclaration): ProtoScopeLocalIdSignature {
+        val proto = ProtoScopeLocalIdSignature.newBuilder()
+
+        proto.id = signature.id
+        proto.file = serializeString(signature.filePath)
+
+        return proto.build()
+    }
+
+    private fun serializeFileSignature(): ProtoFileSignature {
+        return ProtoFileSignature.newBuilder().build()
+    }
+
+    private fun serializeLoweredDeclarationSignature(signature: IdSignature.LoweredDeclarationSignature): ProtoLoweredIdSignature {
+        val proto = ProtoLoweredIdSignature.newBuilder()
+
+        proto.parentSignature = protoIdSignature(signature.original)
+        proto.stage = signature.stage
+        proto.index = signature.index
+
+        return proto.build()
+    }
+
+    fun serializeIdSignature(idSignature: IdSignature): ProtoIdSignature {
         val proto = ProtoIdSignature.newBuilder()
         when (idSignature) {
             is IdSignature.PublicSignature -> proto.publicSig = serializePublicSignature(idSignature)
             is IdSignature.AccessorSignature -> proto.accessorSig = serializeAccessorSignature(idSignature)
             is IdSignature.FileLocalSignature -> proto.privateSig = serializePrivateSignature(idSignature)
             is IdSignature.ScopeLocalDeclaration -> proto.scopedLocalSig = serializeScopeLocalSignature(idSignature)
+            is IdSignature.LoweredDeclarationSignature -> proto.icSig = serializeLoweredDeclarationSignature(idSignature)
+            is IdSignature.FileSignature -> proto.fileSig = serializeFileSignature()
+            is IdSignature.GlobalFileLocalSignature -> proto.privateSig = serializePrivateSignature(idSignature)
+            is IdSignature.GlobalScopeLocalDeclaration -> proto.externalScopedLocalSig = serializeScopeLocalSignature(idSignature)
         }
         return proto.build()
     }
@@ -287,18 +334,25 @@ open class IrFileSerializer(
                 BinarySymbolData.SymbolKind.LOCAL_DELEGATED_PROPERTY_SYMBOL
             is IrTypeAliasSymbol ->
                 BinarySymbolData.SymbolKind.TYPEALIAS_SYMBOL
+            is IrFileSymbol ->
+                BinarySymbolData.SymbolKind.FILE_SYMBOL
             else ->
                 TODO("Unexpected symbol kind: $symbol")
         }
     }
 
     fun serializeIrSymbol(symbol: IrSymbol): Long {
-        val declaration = symbol.owner as? IrDeclaration ?: error("Expected IrDeclaration: ${symbol.owner.render()}")
-
         val symbolKind = protoSymbolKind(symbol)
-        val signatureId = protoIdSignature(declaration)
-        return BinarySymbolData
-            .encode(symbolKind, signatureId)
+
+        val signatureId = when {
+            symbol is IrFileSymbol -> protoIdSignature(IdSignature.FileSignature(symbol)) // TODO: special signature for files?
+            else -> {
+                val declaration = symbol.owner as? IrDeclaration ?: error("Expected IrDeclaration: ${symbol.owner.render()}")
+                protoIdSignature(declaration)
+            }
+        }
+
+        return BinarySymbolData.encode(symbolKind, signatureId)
     }
 
     /* ------- IrTypes ---------------------------------------------------------- */
@@ -422,7 +476,7 @@ open class IrFileSerializer(
             type = (this as? IrTypeProjection)?.type?.toIrTypeKey
         )
 
-    fun serializeIrType(type: IrType) = protoTypeMap.getOrPut(type.toIrTypeKey) {
+    fun serializeIrType(type: IrType?) = if (type == null) -1 else protoTypeMap.getOrPut(type.toIrTypeKey) {
         protoTypeArray.add(serializeIrTypeData(type))
         protoTypeArray.size - 1
     }
@@ -446,6 +500,33 @@ open class IrFileSerializer(
         return proto.build()
     }
 
+    private val returnableBlockStack = ArrayDeque<IrBlock>()
+
+    private fun <T> withBlock(block: IrReturnableBlock, fn: () -> T): T {
+        returnableBlockStack.addLast(block)
+        try {
+            return fn()
+        } finally {
+            returnableBlockStack.removeLast()
+        }
+    }
+
+    private fun serializeReturnableBlock(returnableBlock: IrReturnableBlock): ProtoReturnableBlock = withBlock(returnableBlock) {
+        val proto = ProtoReturnableBlock.newBuilder()
+
+        returnableBlock.origin?.let { proto.setOriginName(serializeIrStatementOrigin(it)) }
+
+        returnableBlock.statements.forEach {
+            proto.addStatement(serializeStatement(it))
+        }
+
+        returnableBlock.inlineFunctionSymbol?.let {
+            proto.inlineFunctionSymbol = serializeIrSymbol(it)
+        }
+
+        proto.build()
+    }
+
     private fun serializeBlock(block: IrBlock): ProtoBlock {
         val proto = ProtoBlock.newBuilder()
 
@@ -493,7 +574,8 @@ open class IrFileSerializer(
         }
 
         for (index in 0 until call.typeArgumentsCount) {
-            val typeArgument = call.getTypeArgument(index) ?: error("Expected type argument at $index in ${call.render()}")
+            val typeArgument = call.getTypeArgument(index)
+            if (typeArgument == null && !allowNullTypes) error("Expected type argument at $index in ${call.render()}")
             proto.addTypeArgument(serializeIrType(typeArgument))
         }
 
@@ -554,6 +636,13 @@ open class IrFileSerializer(
         return proto.build()
     }
 
+    private fun serializeRawFunctionReference(rawReference: IrRawFunctionReference): ProtoRawFunctionReference {
+        val proto = ProtoRawFunctionReference.newBuilder()
+            .setSymbol(serializeIrSymbol(rawReference.symbol))
+
+        return proto.build()
+    }
+
     private fun serializeIrLocalDelegatedPropertyReference(
         callable: IrLocalDelegatedPropertyReference
     ): ProtoLocalDelegatedPropertyReference {
@@ -672,11 +761,23 @@ open class IrFileSerializer(
         return proto.build()
     }
 
-    private fun serializeReturn(expression: IrReturn): ProtoReturn {
-        val proto = ProtoReturn.newBuilder()
-            .setReturnTarget(serializeIrSymbol(expression.returnTargetSymbol))
-            .setValue(serializeExpression(expression.value))
-        return proto.build()
+    private fun serializeReturn(operation: ProtoOperation.Builder, expression: IrReturn) {
+        if (expression.returnTargetSymbol is IrReturnableBlockSymbol) {
+            val targetBlock = expression.returnTargetSymbol.owner
+
+            var t = returnableBlockStack.size - 1
+            while (t >= 0 && returnableBlockStack[t] != targetBlock) --t
+
+            val proto = ProtoReturnableBlockReturn.newBuilder()
+                .setUpCnt(returnableBlockStack.size - t)
+                .setValue(serializeExpression(expression.value))
+            operation.returnableBlockReturn = proto.build()
+        } else {
+            val proto = ProtoReturn.newBuilder()
+                .setReturnTarget(serializeIrSymbol(expression.returnTargetSymbol))
+                .setValue(serializeExpression(expression.value))
+            operation.`return` = proto.build()
+        }
     }
 
     private fun serializeSetField(expression: IrSetField): ProtoSetField =
@@ -913,7 +1014,7 @@ open class IrFileSerializer(
         expression.label?.let {
             proto.label = serializeString(it)
         }
-        val loopId = loopIndex[expression.loop]!!
+        val loopId = loopIndex[expression.loop] ?: -1
         proto.loopId = loopId
 
         return proto.build()
@@ -940,6 +1041,7 @@ open class IrFileSerializer(
 
         // TODO: make me a visitor.
         when (expression) {
+            is IrReturnableBlock -> operationProto.returnableBlock = serializeReturnableBlock(expression)
             is IrBlock -> operationProto.block = serializeBlock(expression)
             is IrBreak -> operationProto.`break` = serializeBreak(expression)
             is IrClassReference -> operationProto.classReference = serializeClassReference(expression)
@@ -961,7 +1063,7 @@ open class IrFileSerializer(
             is IrInstanceInitializerCall -> operationProto.instanceInitializerCall = serializeInstanceInitializerCall(expression)
             is IrLocalDelegatedPropertyReference -> operationProto.localDelegatedPropertyReference = serializeIrLocalDelegatedPropertyReference(expression)
             is IrPropertyReference -> operationProto.propertyReference = serializePropertyReference(expression)
-            is IrReturn -> operationProto.`return` = serializeReturn(expression)
+            is IrReturn -> serializeReturn(operationProto, expression)
             is IrSetField -> operationProto.setField = serializeSetField(expression)
             is IrSetValue -> operationProto.setValue = serializeSetValue(expression)
             is IrStringConcatenation -> operationProto.stringConcat = serializeStringConcat(expression)
@@ -975,6 +1077,7 @@ open class IrFileSerializer(
             is IrDynamicOperatorExpression -> operationProto.dynamicOperator = serializeDynamicOperatorExpression(expression)
             is IrErrorCallExpression -> operationProto.errorCallExpression = serializeErrorCallExpression(expression)
             is IrErrorExpression -> operationProto.errorExpression = serializeErrorExpression(expression)
+            is IrRawFunctionReference -> operationProto.rawFunctionReference = serializeRawFunctionReference(expression)
             else -> {
                 TODO("Expression serialization not implemented yet: ${expression.render()}.")
             }
@@ -1009,6 +1112,9 @@ open class IrFileSerializer(
             is IrSyntheticBody -> {
                 proto.syntheticBody = serializeSyntheticBody(statement)
             }
+            is IrExpressionBody -> {
+                proto.expression = serializeExpression(statement.expression)
+            }
             else -> {
                 TODO("Statement not implemented yet: ${statement.render()}")
             }
@@ -1042,6 +1148,8 @@ open class IrFileSerializer(
         if (!skipMutableState) {
             parameter.varargElementType?.let { proto.setVarargElementType(serializeIrType(it)) }
             parameter.defaultValue?.let { proto.setDefaultValue(serializeIrExpressionBody(it.expression)) }
+        } else {
+            proto.index = parameter.index
         }
 
         return proto.build()
@@ -1054,6 +1162,12 @@ open class IrFileSerializer(
         parameter.superTypes.forEach {
             proto.addSuperType(serializeIrType(it))
         }
+
+        if (skipMutableState) {
+            proto.index = parameter.index
+            proto.isGlobal = true // parameter.parent is IrClass
+        }
+
         return proto.build()
     }
 

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/IrSymbolDeserializer.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.backend.common.serialization
 import org.jetbrains.kotlin.backend.common.serialization.encodings.*
 import org.jetbrains.kotlin.backend.common.serialization.proto.Actual
 import org.jetbrains.kotlin.backend.common.serialization.proto.IdSignature.IdsigCase.*
+import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.symbols.impl.IrAnonymousInitializerSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrLocalDelegatedPropertySymbolImpl
@@ -18,20 +19,24 @@ 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.FileLocalIdSignature as ProtoFileLocalIdSignature
+import org.jetbrains.kotlin.backend.common.serialization.proto.ScopeLocalIdSignature as ProtoScopeLocalIdSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.IdSignature as ProtoIdSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.PublicIdSignature as ProtoPublicIdSignature
+import org.jetbrains.kotlin.backend.common.serialization.proto.LoweredIdSignature as ProtoLoweredIdSignature
 
 class IrSymbolDeserializer(
     val symbolTable: ReferenceSymbolTable,
     val fileReader: IrLibraryFile,
+    val fileSymbol: IrFileSymbol,
     val actuals: List<Actual>,
     val enqueueLocalTopLevelDeclaration: (IdSignature) -> Unit,
     val handleExpectActualMapping: (IdSignature, IrSymbol) -> IrSymbol,
+    private val enqueueAllDeclarations: Boolean = false,
+    private val useGlobalSignatures: Boolean = false,
+    val deserializedSymbols: MutableMap<IdSignature, IrSymbol> = mutableMapOf(), // Per-file signature cache. TODO: do we really need it?
     val deserializePublicSymbol: (IdSignature, BinarySymbolData.SymbolKind) -> IrSymbol,
 ) {
 
-    val deserializedSymbols = mutableMapOf<IdSignature, IrSymbol>()
-
     fun deserializeIrSymbol(idSig: IdSignature, symbolKind: BinarySymbolData.SymbolKind): IrSymbol {
         return deserializedSymbols.getOrPut(idSig) {
             val symbol = referenceDeserializedSymbol(symbolKind, idSig)
@@ -57,13 +62,15 @@ class IrSymbolDeserializer(
             BinarySymbolData.SymbolKind.RECEIVER_PARAMETER_SYMBOL -> IrValueParameterSymbolImpl()
             BinarySymbolData.SymbolKind.LOCAL_DELEGATED_PROPERTY_SYMBOL ->
                 IrLocalDelegatedPropertySymbolImpl()
+            BinarySymbolData.SymbolKind.FILE_SYMBOL -> (idSig as IdSignature.FileSignature).symbol
             else -> error("Unexpected classifier symbol kind: $symbolKind for signature $idSig")
         }
     }
 
     fun referenceLocalIrSymbol(symbol: IrSymbol, signature: IdSignature) {
-        assert(signature.isLocal)
+//        assert(signature.isLocal)
         deserializedSymbols.putIfAbsent(signature, symbol)
+        if (enqueueAllDeclarations) enqueueLocalTopLevelDeclaration(signature)
     }
 
     fun referenceSimpleFunctionByLocalSignature(idSignature: IdSignature) : IrSimpleFunctionSymbol =
@@ -74,10 +81,12 @@ class IrSymbolDeserializer(
 
     private fun deserializeIrSymbolData(idSignature: IdSignature, symbolKind: BinarySymbolData.SymbolKind): IrSymbol {
         if (idSignature.isLocal) {
-            if (idSignature.hasTopLevel) {
-                enqueueLocalTopLevelDeclaration(idSignature.topLevelSignature())
-            }
             return deserializedSymbols.getOrPut(idSignature) {
+                if (enqueueAllDeclarations) {
+                    enqueueLocalTopLevelDeclaration(idSignature)
+                } else if (idSignature.hasTopLevel) {
+                    enqueueLocalTopLevelDeclaration(idSignature.topLevelSignature())
+                }
                 referenceDeserializedSymbol(symbolKind, idSignature)
             }
         }
@@ -136,12 +145,33 @@ class IrSymbolDeserializer(
         return IdSignature.AccessorSignature(propertySignature, accessorSignature)
     }
 
-    private fun deserializeFileLocalIdSignature(proto: ProtoFileLocalIdSignature): IdSignature.FileLocalSignature {
-        return IdSignature.FileLocalSignature(deserializeIdSignature(proto.container), proto.localId)
+    private fun deserializeFileLocalIdSignature(proto: ProtoFileLocalIdSignature): IdSignature {
+        if (useGlobalSignatures) {
+            val fp = if (proto.hasFile()) fileReader.deserializeString(proto.file) else fileSymbol.owner.path
+            return IdSignature.GlobalFileLocalSignature(deserializeIdSignature(proto.container), proto.localId, fp)
+        } else {
+            return IdSignature.FileLocalSignature(deserializeIdSignature(proto.container), proto.localId)
+        }
     }
 
-    private fun deserializeScopeLocalIdSignature(proto: Int): IdSignature.ScopeLocalDeclaration {
-        return IdSignature.ScopeLocalDeclaration(proto)
+    private fun deserializeScopeLocalIdSignature(proto: Int): IdSignature {
+        if (useGlobalSignatures) {
+            return IdSignature.GlobalScopeLocalDeclaration(proto, filePath = fileSymbol.owner.path)
+        } else {
+            return IdSignature.ScopeLocalDeclaration(proto)
+        }
+    }
+
+    private fun deserializeExternalScopeLocalIdSignature(proto: ProtoScopeLocalIdSignature): IdSignature {
+        if (useGlobalSignatures) {
+            return IdSignature.GlobalScopeLocalDeclaration(proto.id, filePath = fileReader.deserializeString(proto.file))
+        } else {
+            return IdSignature.ScopeLocalDeclaration(proto.id)
+        }
+    }
+
+    private fun deserializeLoweredDeclarationSignature(proto: ProtoLoweredIdSignature): IdSignature.LoweredDeclarationSignature {
+        return IdSignature.LoweredDeclarationSignature(deserializeIdSignature(proto.parentSignature), proto.stage, proto.index)
     }
 
     fun deserializeSignatureData(proto: ProtoIdSignature): IdSignature {
@@ -150,6 +180,9 @@ class IrSymbolDeserializer(
             ACCESSOR_SIG -> deserializeAccessorIdSignature(proto.accessorSig)
             PRIVATE_SIG -> deserializeFileLocalIdSignature(proto.privateSig)
             SCOPED_LOCAL_SIG -> deserializeScopeLocalIdSignature(proto.scopedLocalSig)
+            IC_SIG -> deserializeLoweredDeclarationSignature(proto.icSig)
+            FILE_SIG -> IdSignature.FileSignature(fileSymbol)
+            EXTERNAL_SCOPED_LOCAL_SIG -> deserializeExternalScopeLocalIdSignature(proto.externalScopedLocalSig)
             else -> error("Unexpected IdSignature kind: ${proto.idsigCase}")
         }
     }

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/encodings/BinarySymbolData.kt

@@ -22,7 +22,8 @@ value class BinarySymbolData(val code: Long) {
         RECEIVER_PARAMETER_SYMBOL, // ReceiverParameterDescriptor rather than ValueParameterDescriptor.
         PROPERTY_SYMBOL,
         LOCAL_DELEGATED_PROPERTY_SYMBOL,
-        TYPEALIAS_SYMBOL;
+        TYPEALIAS_SYMBOL,
+        FILE_SYMBOL;
     }
 
     private fun symbolKindId(): Int = (code and 0xFF).toInt()