[K/N] Support non-LazyIR for cached libraries

Use the same lazy deserialization mechanism (inline function bodies, class fields)
for both LazyIR and usual IR

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

@@ -577,6 +577,16 @@ class IrDeclarationDeserializer(
         }
     }
 
+    fun <T : IrFunction> T.withDeserializeBodies(block: T.() -> Unit) {
+        val oldBodiesPolicy = deserializeBodies
+        try {
+            deserializeBodies = true
+            usingParent { block() }
+        } finally {
+            deserializeBodies = oldBodiesPolicy
+        }
+    }
+
     internal fun deserializeIrFunction(proto: ProtoFunction, setParent: Boolean = true): IrSimpleFunction =
         withDeserializedIrFunctionBase<IrSimpleFunctionSymbol, IrSimpleFunction>(
             proto.base,

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

@@ -102,7 +102,7 @@ abstract class KotlinIrLinker(
 
     protected abstract fun isBuiltInModule(moduleDescriptor: ModuleDescriptor): Boolean
 
-    private fun deserializeAllReachableTopLevels() {
+    fun deserializeAllReachableTopLevels() {
         while (modulesWithReachableTopLevels.isNotEmpty()) {
             val moduleDeserializer = modulesWithReachableTopLevels.first()
             modulesWithReachableTopLevels.remove(moduleDeserializer)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/PsiToIr.kt

@@ -175,10 +175,8 @@ internal fun PsiToIrContext.psiToIr(
                     val kotlinLibrary = (dependency.getCapability(KlibModuleOrigin.CAPABILITY) as? DeserializedKlibModuleOrigin)?.library
                     val isFullyCachedLibrary = kotlinLibrary != null &&
                             config.cachedLibraries.isLibraryCached(kotlinLibrary) && kotlinLibrary != config.libraryToCache?.klib
-                    if (isProducingLibrary || (config.lazyIrForCaches && isFullyCachedLibrary))
+                    if (isProducingLibrary || isFullyCachedLibrary)
                         linker.deserializeOnlyHeaderModule(dependency, kotlinLibrary)
-                    else if (isFullyCachedLibrary)
-                        linker.deserializeHeadersWithInlineBodies(dependency, kotlinLibrary!!)
                     else
                         linker.deserializeIrModuleHeader(dependency, kotlinLibrary, dependency.name.asString())
                 }

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/descriptors/ClassLayoutBuilder.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.backend.konan.llvm.computeFunctionName
 import org.jetbrains.kotlin.backend.konan.llvm.toLLVMType
 import org.jetbrains.kotlin.backend.konan.llvm.localHash
 import org.jetbrains.kotlin.backend.konan.lower.bridgeTarget
+import org.jetbrains.kotlin.backend.konan.serialization.KonanIrLinker
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
@@ -507,17 +508,10 @@ internal class ClassLayoutBuilder(val irClass: IrClass, val context: Context) {
         val outerThisField = if (irClass.isInner)
             context.innerClassesSupport.getOuterThisField(irClass)
         else null
-        val packageFragment = irClass.getPackageFragment()
+
-        if (packageFragment is IrExternalPackageFragment) {
+        val moduleDeserializer = context.irLinker.getCachedDeclarationModuleDeserializer(irClass)
-            val moduleDescriptor = packageFragment.packageFragmentDescriptor.containingDeclaration
+        if (moduleDeserializer != null)
-            if (moduleDescriptor.isFromInteropLibrary()) return emptyList()
-            val moduleDeserializer = context.irLinker.moduleDeserializers[moduleDescriptor]
-                    ?: error("No module deserializer for ${irClass.render()}")
-            require(context.config.cachedLibraries.isLibraryCached(moduleDeserializer.klib)) {
-                "No IR and no cache for ${irClass.render()}"
-            }
             return moduleDeserializer.deserializeClassFields(irClass, outerThisField?.toFieldInfo(llvm))
-        }
 
         val declarations = irClass.declarations.toMutableList()
         outerThisField?.let {

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CacheInfoBuilder.kt

@@ -93,7 +93,7 @@ internal class CacheInfoBuilder(
             }
 
             private fun processFunction(function: IrFunction) {
-                if (function.getPackageFragment() !is IrExternalPackageFragment) {
+                if (generationState.context.irLinker.getCachedDeclarationModuleDeserializer(function) == null) {
                     generationState.calledFromExportedInlineFunctions.add(function)
                     (function as? IrConstructor)?.constructedClass?.let {
                         generationState.constructedFromExportedInlineFunctions.add(it)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeInlineFunctionResolver.kt

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.backend.common.DeclarationTransformer
 import org.jetbrains.kotlin.backend.common.lower.*
 import org.jetbrains.kotlin.backend.common.lower.inline.*
 import org.jetbrains.kotlin.backend.konan.*
-import org.jetbrains.kotlin.ir.declarations.IrExternalPackageFragment
+import org.jetbrains.kotlin.backend.konan.serialization.KonanIrLinker
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.deepCopyWithVariables
@@ -38,9 +38,15 @@ internal class NativeInlineFunctionResolver(override val context: Context, val g
         generationState.inlineFunctionOrigins[function]?.let { return it.irFunction }
 
         val packageFragment = function.getPackageFragment()
-        val functionIsNotFromLazyIr = packageFragment !is IrExternalPackageFragment
+        val moduleDeserializer = context.irLinker.getCachedDeclarationModuleDeserializer(function)
         val irFile: IrFile
-        val (possiblyLoweredFunction, shouldLower) = if (functionIsNotFromLazyIr) {
+        val (possiblyLoweredFunction, shouldLower) = if (moduleDeserializer != null) {
+            // The function is cached, get its body from the IR linker.
+            val (firstAccess, deserializedInlineFunction) = moduleDeserializer.deserializeInlineFunction(function)
+            generationState.inlineFunctionOrigins[function] = deserializedInlineFunction
+            irFile = deserializedInlineFunction.irFile
+            function to firstAccess
+        } else {
             irFile = packageFragment as IrFile
             val partiallyLoweredFunction = context.inlineFunctionsSupport.getPartiallyLoweredInlineFunction(function)
             if (partiallyLoweredFunction == null)
@@ -50,23 +56,12 @@ internal class NativeInlineFunctionResolver(override val context: Context, val g
                         InlineFunctionOriginInfo(partiallyLoweredFunction, irFile, function.startOffset, function.endOffset)
                 partiallyLoweredFunction to false
             }
-        } else {
-            // The function is from Lazy IR, get its body from the IR linker.
-            val moduleDescriptor = packageFragment.packageFragmentDescriptor.containingDeclaration
-            val moduleDeserializer = context.irLinker.moduleDeserializers[moduleDescriptor]
-                    ?: error("No module deserializer for ${function.render()}")
-            require(context.config.cachedLibraries.isLibraryCached(moduleDeserializer.klib)) {
-                "No IR and no cache for ${function.render()}"
-            }
-            val (firstAccess, deserializedInlineFunction) = moduleDeserializer.deserializeInlineFunction(function)
-            generationState.inlineFunctionOrigins[function] = deserializedInlineFunction
-            irFile = deserializedInlineFunction.irFile
-            function to firstAccess
         }
 
         if (shouldLower) {
-            lower(possiblyLoweredFunction, irFile, functionIsNotFromLazyIr)
+            val functionIsCached = moduleDeserializer != null
-            if (functionIsNotFromLazyIr) {
+            lower(possiblyLoweredFunction, irFile, functionIsCached)
+            if (!functionIsCached) {
                 generationState.inlineFunctionOrigins[function] =
                         InlineFunctionOriginInfo(context.inlineFunctionsSupport.savePartiallyLoweredInlineFunction(possiblyLoweredFunction),
                                 irFile, function.startOffset, function.endOffset)
@@ -75,7 +70,7 @@ internal class NativeInlineFunctionResolver(override val context: Context, val g
         return possiblyLoweredFunction
     }
 
-    private fun lower(function: IrFunction, irFile: IrFile, functionIsNotFromLazyIr: Boolean) {
+    private fun lower(function: IrFunction, irFile: IrFile, functionIsCached: Boolean) {
         val body = function.body ?: return
 
         PreInlineLowering(context).lower(body, function, irFile)
@@ -92,7 +87,7 @@ internal class NativeInlineFunctionResolver(override val context: Context, val g
 
         LocalClassesInInlineLambdasLowering(context).lower(body, function)
 
-        if (!context.config.produce.isCache && functionIsNotFromLazyIr) {
+        if (!(context.config.produce.isCache || functionIsCached)) {
             // Do not extract local classes off of inline functions from cached libraries.
             LocalClassesInInlineFunctionsLowering(context).lower(body, function)
             LocalClassesExtractionFromInlineFunctionsLowering(context).lower(body, function)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/serialization/KonanIrlinker.kt

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.backend.konan.descriptors.findPackage
 import org.jetbrains.kotlin.backend.konan.descriptors.isFromInteropLibrary
 import org.jetbrains.kotlin.backend.konan.descriptors.isInteropLibrary
 import org.jetbrains.kotlin.backend.konan.ir.interop.IrProviderForCEnumAndCStructStubs
+import org.jetbrains.kotlin.backend.konan.ir.konanLibrary
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.konan.isNativeStdlib
 import org.jetbrains.kotlin.fir.lazy.Fir2IrLazyClass
@@ -491,6 +492,18 @@ internal class KonanIrLinker(
     val moduleDeserializers = mutableMapOf<ModuleDescriptor, KonanPartialModuleDeserializer>()
     val klibToModuleDeserializerMap = mutableMapOf<KotlinLibrary, KonanPartialModuleDeserializer>()
 
+    fun getCachedDeclarationModuleDeserializer(declaration: IrDeclaration): KonanPartialModuleDeserializer? {
+        val packageFragment = declaration.getPackageFragment()
+        val moduleDescriptor = packageFragment.packageFragmentDescriptor.containingDeclaration
+        val klib = packageFragment.konanLibrary
+        val declarationBeingCached = packageFragment is IrFile && klib != null && libraryBeingCached?.klib == klib
+                && libraryBeingCached.strategy.contains(packageFragment.path)
+        return if (klib != null && !moduleDescriptor.isFromInteropLibrary()
+                && cachedLibraries.isLibraryCached(klib) && !declarationBeingCached)
+            moduleDeserializers[moduleDescriptor] ?: error("No module deserializer for ${declaration.render()}")
+        else null
+    }
+
     override fun createModuleDeserializer(moduleDescriptor: ModuleDescriptor, klib: KotlinLibrary?, strategyResolver: (String) -> DeserializationStrategy) =
             when {
                 moduleDescriptor === forwardModuleDescriptor -> {
@@ -839,76 +852,84 @@ internal class KonanIrLinker(
         }
 
         private fun deserializeInlineFunctionInternal(function: IrFunction): InlineFunctionOriginInfo {
-            val packageFragment = function.getPackageFragment() as? IrExternalPackageFragment
+            val packageFragment = function.getPackageFragment()
-                    ?: error("Expected an external package fragment for ${function.render()}")
             if (function.parents.any { (it as? IrFunction)?.isInline == true }) {
                 // Already deserialized by the top-most inline function.
                 return InlineFunctionOriginInfo(
                         function,
-                        inlineFunctionFiles[packageFragment]
+                        packageFragment as? IrFile
+                                ?: inlineFunctionFiles[packageFragment as IrExternalPackageFragment]
                                 ?: error("${function.render()} should've been deserialized along with its parent"),
                         function.startOffset, function.endOffset
                 )
             }
 
-            val signature = function.symbol.signature ?: descriptorSignatures[function.descriptor]
+            val signature = function.symbol.signature
-            ?: error("No signature for ${function.render()}")
+                    ?: descriptorSignatures[function.descriptor]
+                    ?: error("No signature for ${function.render()}")
             val inlineFunctionReference = inlineFunctionReferences[signature]
                     ?: error("No inline function reference for ${function.render()}, sig = ${signature.render()}")
             val fileDeserializationState = inlineFunctionReference.file.deserializationState
             val declarationDeserializer = fileDeserializationState.declarationDeserializer
-            val symbolDeserializer = declarationDeserializer.symbolDeserializer
 
-            inlineFunctionFiles[packageFragment]?.let {
+            if (packageFragment is IrExternalPackageFragment) {
-                require(it == fileDeserializationState.file) {
+                val symbolDeserializer = declarationDeserializer.symbolDeserializer
-                    "Different files ${it.fileEntry.name} and ${fileDeserializationState.file.fileEntry.name} have the same $packageFragment"
+
+                inlineFunctionFiles[packageFragment]?.let {
+                    require(it == fileDeserializationState.file) {
+                        "Different files ${it.fileEntry.name} and ${fileDeserializationState.file.fileEntry.name} have the same $packageFragment"
+                    }
                 }
-            }
+                inlineFunctionFiles[packageFragment] = fileDeserializationState.file
-            inlineFunctionFiles[packageFragment] = fileDeserializationState.file
 
-            val outerClasses = (function.parent as? IrClass)?.getOuterClasses(takeOnlyInner = true) ?: emptyList()
+                val outerClasses = (function.parent as? IrClass)?.getOuterClasses(takeOnlyInner = true) ?: emptyList()
-            require((outerClasses.getOrNull(0)?.firstNonClassParent ?: function.parent) is IrExternalPackageFragment) {
+                require((outerClasses.getOrNull(0)?.firstNonClassParent ?: function.parent) is IrPackageFragment) {
-                "Local inline functions are not supported: ${function.render()}"
+                    "Local inline functions are not supported: ${function.render()}"
-            }
+                }
 
-            var endToEndTypeParameterIndex = 0
+                var endToEndTypeParameterIndex = 0
-            outerClasses.forEach { outerClass ->
+                outerClasses.forEach { outerClass ->
-                outerClass.typeParameters.forEach { parameter ->
+                    outerClass.typeParameters.forEach { parameter ->
+                        val sigIndex = inlineFunctionReference.typeParameterSigs[endToEndTypeParameterIndex++]
+                        referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                    }
+                }
+                function.typeParameters.forEach { parameter ->
                     val sigIndex = inlineFunctionReference.typeParameterSigs[endToEndTypeParameterIndex++]
                     referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
                 }
-            }
+                function.valueParameters.forEachIndexed { index, parameter ->
-            function.typeParameters.forEach { parameter ->
+                    val sigIndex = inlineFunctionReference.valueParameterSigs[index]
-                val sigIndex = inlineFunctionReference.typeParameterSigs[endToEndTypeParameterIndex++]
+                    referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
-                referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                }
-            }
+                function.extensionReceiverParameter?.let { parameter ->
-            function.valueParameters.forEachIndexed { index, parameter ->
+                    val sigIndex = inlineFunctionReference.extensionReceiverSig
-                val sigIndex = inlineFunctionReference.valueParameterSigs[index]
+                    require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the extension receiver for ${function.render()}" }
-                referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                    referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
-            }
+                }
-            function.extensionReceiverParameter?.let { parameter ->
+                function.dispatchReceiverParameter?.let { parameter ->
-                val sigIndex = inlineFunctionReference.extensionReceiverSig
+                    val sigIndex = inlineFunctionReference.dispatchReceiverSig
-                require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the extension receiver for ${function.render()}" }
+                    require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the dispatch receiver for ${function.render()}" }
-                referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                    referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
-            }
+                }
-            function.dispatchReceiverParameter?.let { parameter ->
+                for (index in 0 until outerClasses.size - 1) {
-                val sigIndex = inlineFunctionReference.dispatchReceiverSig
+                    val sigIndex = inlineFunctionReference.outerReceiverSigs[index]
-                require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the dispatch receiver for ${function.render()}" }
+                    referenceIrSymbol(symbolDeserializer, sigIndex, outerClasses[index].thisReceiver!!.symbol)
-                referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                }
-            }
-            for (index in 0 until outerClasses.size - 1) {
-                val sigIndex = inlineFunctionReference.outerReceiverSigs[index]
-                referenceIrSymbol(symbolDeserializer, sigIndex, outerClasses[index].thisReceiver!!.symbol)
             }
 
             with(declarationDeserializer) {
-                function.body = (deserializeStatementBody(inlineFunctionReference.body) as IrBody).setDeclarationsParent(function)
+                function.withDeserializeBodies {
-                function.valueParameters.forEachIndexed { index, parameter ->
+                    body = (deserializeStatementBody(inlineFunctionReference.body) as IrBody)
-                    val defaultValueIndex = inlineFunctionReference.defaultValues[index]
+                    valueParameters.forEachIndexed { index, parameter ->
-                    if (defaultValueIndex != InvalidIndex)
+                        val defaultValueIndex = inlineFunctionReference.defaultValues[index]
-                        parameter.defaultValue = deserializeExpressionBody(defaultValueIndex)?.setDeclarationsParent(function)
+                        if (defaultValueIndex != InvalidIndex)
+                            parameter.defaultValue = deserializeExpressionBody(defaultValueIndex)
+                    }
                 }
             }
+            if (packageFragment is IrFile)
+                deserializeAllReachableTopLevels()
 
             partialLinkageSupport.exploreClassifiers(fakeOverrideBuilder)
             partialLinkageSupport.exploreClassifiersInInlineLazyIrFunction(function)
@@ -935,8 +956,6 @@ internal class KonanIrLinker(
         private val lock = Any()
 
         fun deserializeClassFields(irClass: IrClass, outerThisFieldInfo: ClassLayoutBuilder.FieldInfo?): List<ClassLayoutBuilder.FieldInfo> = synchronized(lock) {
-            irClass.getPackageFragment() as? IrExternalPackageFragment
-                    ?: error("Expected an external package fragment for ${irClass.render()}")
             val signature = irClass.symbol.signature
                     ?: error("No signature for ${irClass.render()}")
             val serializedClassFields = classesFields[signature]
@@ -945,20 +964,22 @@ internal class KonanIrLinker(
             val declarationDeserializer = fileDeserializationState.declarationDeserializer
             val symbolDeserializer = declarationDeserializer.symbolDeserializer
 
-            val outerClasses = irClass.getOuterClasses(takeOnlyInner = true)
+            if (irClass.getPackageFragment() is IrExternalPackageFragment) {
-            require(outerClasses.first().firstNonClassParent is IrExternalPackageFragment) {
+                val outerClasses = irClass.getOuterClasses(takeOnlyInner = true)
-                "Local classes are not supported: ${irClass.render()}"
+                require(outerClasses.first().firstNonClassParent is IrExternalPackageFragment) {
-            }
+                    "Local classes are not supported: ${irClass.render()}"
-
+                }
-            var endToEndTypeParameterIndex = 0
+
-            outerClasses.forEach { outerClass ->
+                var endToEndTypeParameterIndex = 0
-                outerClass.typeParameters.forEach { parameter ->
+                outerClasses.forEach { outerClass ->
-                    val sigIndex = serializedClassFields.typeParameterSigs[endToEndTypeParameterIndex++]
+                    outerClass.typeParameters.forEach { parameter ->
-                    referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                        val sigIndex = serializedClassFields.typeParameterSigs[endToEndTypeParameterIndex++]
+                        referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                    }
+                }
+                require(endToEndTypeParameterIndex == serializedClassFields.typeParameterSigs.size) {
+                    "Not all type parameters have been referenced"
                 }
-            }
-            require(endToEndTypeParameterIndex == serializedClassFields.typeParameterSigs.size) {
-                "Not all type parameters have been referenced"
             }
 
             fun getByClassId(classId: ClassId): IrClassSymbol {