[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
2023-06-13 12:28:30 +03:00
parent 6318da0f02
commit cc9dd48c09
7 changed files with 112 additions and 94 deletions
@@ -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,
@@ -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)
@@ -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())
                }
@@ -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 moduleDescriptor = packageFragment.packageFragmentDescriptor.containingDeclaration
-            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()}"
-            }
+
+        val moduleDeserializer = context.irLinker.getCachedDeclarationModuleDeserializer(irClass)
+        if (moduleDeserializer != null)
            return moduleDeserializer.deserializeClassFields(irClass, outerThisField?.toFieldInfo(llvm))
-        }

        val declarations = irClass.declarations.toMutableList()
        outerThisField?.let {
@@ -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)
@@ -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)
-            if (functionIsNotFromLazyIr) {
+            val functionIsCached = moduleDeserializer != null
+            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)
@@ -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
-                    ?: error("Expected an external package fragment for ${function.render()}")
+            val packageFragment = function.getPackageFragment()
            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]
-            ?: error("No signature for ${function.render()}")
+            val signature = function.symbol.signature
+                    ?: 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 {
-                require(it == fileDeserializationState.file) {
-                    "Different files ${it.fileEntry.name} and ${fileDeserializationState.file.fileEntry.name} have the same $packageFragment"
+            if (packageFragment is IrExternalPackageFragment) {
+                val symbolDeserializer = declarationDeserializer.symbolDeserializer
+
+                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()
-            require((outerClasses.getOrNull(0)?.firstNonClassParent ?: function.parent) is IrExternalPackageFragment) {
-                "Local inline functions are not supported: ${function.render()}"
-            }
+                val outerClasses = (function.parent as? IrClass)?.getOuterClasses(takeOnlyInner = true) ?: emptyList()
+                require((outerClasses.getOrNull(0)?.firstNonClassParent ?: function.parent) is IrPackageFragment) {
+                    "Local inline functions are not supported: ${function.render()}"
+                }

-            var endToEndTypeParameterIndex = 0
-            outerClasses.forEach { outerClass ->
-                outerClass.typeParameters.forEach { parameter ->
+                var endToEndTypeParameterIndex = 0
+                outerClasses.forEach { outerClass ->
+                    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.typeParameters.forEach { parameter ->
-                val sigIndex = inlineFunctionReference.typeParameterSigs[endToEndTypeParameterIndex++]
-                referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
-            }
-            function.valueParameters.forEachIndexed { index, parameter ->
-                val sigIndex = inlineFunctionReference.valueParameterSigs[index]
-                referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
-            }
-            function.extensionReceiverParameter?.let { parameter ->
-                val sigIndex = inlineFunctionReference.extensionReceiverSig
-                require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the extension receiver for ${function.render()}" }
-                referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
-            }
-            function.dispatchReceiverParameter?.let { parameter ->
-                val sigIndex = inlineFunctionReference.dispatchReceiverSig
-                require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the dispatch receiver for ${function.render()}" }
-                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)
+                function.valueParameters.forEachIndexed { index, parameter ->
+                    val sigIndex = inlineFunctionReference.valueParameterSigs[index]
+                    referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                }
+                function.extensionReceiverParameter?.let { parameter ->
+                    val sigIndex = inlineFunctionReference.extensionReceiverSig
+                    require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the extension receiver for ${function.render()}" }
+                    referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+                }
+                function.dispatchReceiverParameter?.let { parameter ->
+                    val sigIndex = inlineFunctionReference.dispatchReceiverSig
+                    require(sigIndex != InvalidIndex) { "Expected a valid sig reference to the dispatch receiver for ${function.render()}" }
+                    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.valueParameters.forEachIndexed { index, parameter ->
-                    val defaultValueIndex = inlineFunctionReference.defaultValues[index]
-                    if (defaultValueIndex != InvalidIndex)
-                        parameter.defaultValue = deserializeExpressionBody(defaultValueIndex)?.setDeclarationsParent(function)
+                function.withDeserializeBodies {
+                    body = (deserializeStatementBody(inlineFunctionReference.body) as IrBody)
+                    valueParameters.forEachIndexed { index, parameter ->
+                        val defaultValueIndex = inlineFunctionReference.defaultValues[index]
+                        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)
-            require(outerClasses.first().firstNonClassParent is IrExternalPackageFragment) {
-                "Local classes are not supported: ${irClass.render()}"
-            }
-
-            var endToEndTypeParameterIndex = 0
-            outerClasses.forEach { outerClass ->
-                outerClass.typeParameters.forEach { parameter ->
-                    val sigIndex = serializedClassFields.typeParameterSigs[endToEndTypeParameterIndex++]
-                    referenceIrSymbol(symbolDeserializer, sigIndex, parameter.symbol)
+            if (irClass.getPackageFragment() is IrExternalPackageFragment) {
+                val outerClasses = irClass.getOuterClasses(takeOnlyInner = true)
+                require(outerClasses.first().firstNonClassParent is IrExternalPackageFragment) {
+                    "Local classes are not supported: ${irClass.render()}"
+                }
+
+                var endToEndTypeParameterIndex = 0
+                outerClasses.forEach { outerClass ->
+                    outerClass.typeParameters.forEach { parameter ->
+                        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 {