[K/N][IR] Reworked saving inline function bodies for cross-file inlining

#KT-57053 Fixed #KT-57029 Fixed
2023-03-02 18:07:11 +02:00
parent 315d9089b0
commit 7f0a4ad737
5 changed files with 92 additions and 79 deletions
@@ -13,41 +13,24 @@ import org.jetbrains.kotlin.backend.konan.cexport.CAdapterExportedElements
 import org.jetbrains.kotlin.backend.konan.descriptors.BridgeDirections
 import org.jetbrains.kotlin.backend.konan.descriptors.ClassLayoutBuilder
 import org.jetbrains.kotlin.backend.konan.descriptors.GlobalHierarchyAnalysisResult
-import org.jetbrains.kotlin.backend.konan.driver.phases.PsiToIrContext
-import org.jetbrains.kotlin.backend.konan.driver.phases.PsiToIrOutput
 import org.jetbrains.kotlin.backend.konan.ir.KonanIr
 import org.jetbrains.kotlin.backend.konan.ir.KonanSymbols
 import org.jetbrains.kotlin.backend.konan.llvm.CodegenClassMetadata
-import org.jetbrains.kotlin.backend.konan.llvm.Lifetime
-import org.jetbrains.kotlin.backend.konan.llvm.coverage.CoverageManager
 import org.jetbrains.kotlin.backend.konan.lower.*
 import org.jetbrains.kotlin.backend.konan.objcexport.ObjCExportCodeSpec
 import org.jetbrains.kotlin.backend.konan.objcexport.ObjCExportedInterface
-import org.jetbrains.kotlin.backend.konan.optimizations.DevirtualizationAnalysis
-import org.jetbrains.kotlin.backend.konan.optimizations.ModuleDFG
 import org.jetbrains.kotlin.backend.konan.serialization.KonanIrLinker
 import org.jetbrains.kotlin.builtins.konan.KonanBuiltIns
-import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
-import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
 import org.jetbrains.kotlin.ir.IrBuiltIns
-import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.lazy.IrLazyClass
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
-import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.types.IrTypeSystemContext
 import org.jetbrains.kotlin.ir.types.IrTypeSystemContextImpl
-import org.jetbrains.kotlin.ir.util.SymbolTable
-import org.jetbrains.kotlin.konan.library.KonanLibraryLayout
-import org.jetbrains.kotlin.konan.target.Architecture
-import org.jetbrains.kotlin.konan.target.KonanTarget
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
-import kotlin.LazyThreadSafetyMode.PUBLICATION

 internal class NativeMapping : DefaultMapping() {
    data class BridgeKey(val target: IrSimpleFunction, val bridgeDirections: BridgeDirections)
@@ -60,7 +43,7 @@ internal class NativeMapping : DefaultMapping() {
    val enumValueGetters = DefaultDelegateFactory.newDeclarationToDeclarationMapping<IrClass, IrFunction>()
    val enumEntriesMaps = mutableMapOf<IrClass, Map<Name, LoweredEnumEntryDescription>>()
    val bridges = mutableMapOf<BridgeKey, IrSimpleFunction>()
-    val notLoweredInlineFunctions = mutableMapOf<IrFunctionSymbol, IrFunction>()
+    val partiallyLoweredInlineFunctions = mutableMapOf<IrFunctionSymbol, IrFunction>()
    val outerThisCacheAccessors = DefaultDelegateFactory.newDeclarationToDeclarationMapping<IrClass, IrSimpleFunction>()
    val lateinitPropertyCacheAccessors = DefaultDelegateFactory.newDeclarationToDeclarationMapping<IrProperty, IrSimpleFunction>()
    val objectInstanceGetter = DefaultDelegateFactory.newDeclarationToDeclarationMapping<IrClass, IrSimpleFunction>()
@@ -69,7 +69,7 @@ internal class NativeGenerationState(
    val eagerInitializedFiles = mutableListOf<SerializedEagerInitializedFile>()
    val calledFromExportedInlineFunctions = mutableSetOf<IrFunction>()
    val constructedFromExportedInlineFunctions = mutableSetOf<IrClass>()
-    val loweredInlineFunctions = mutableMapOf<IrFunction, InlineFunctionOriginInfo>()
+    val inlineFunctionOrigins = mutableMapOf<IrFunction, InlineFunctionOriginInfo>()

    private val localClassNames = mutableMapOf<IrAttributeContainer, String>()
    fun getLocalClassName(container: IrAttributeContainer): String? = localClassNames[container.attributeOwnerId]
@@ -33,7 +33,6 @@ import org.jetbrains.kotlin.backend.konan.lower.RedundantCoercionsCleaner
 import org.jetbrains.kotlin.backend.konan.lower.ReturnsInsertionLowering
 import org.jetbrains.kotlin.backend.konan.lower.UnboxInlineLowering
 import org.jetbrains.kotlin.backend.konan.optimizations.KonanBCEForLoopBodyTransformer
-import org.jetbrains.kotlin.ir.util.isLocal
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
@@ -105,7 +104,7 @@ private val sharedVariablesPhase = createFileLoweringPhase(
        prerequisite = setOf(lateinitPhase)
 )

-private val extractLocalClassesFromInlineBodies = createFileLoweringPhase(
+private val lowerOuterThisInInlineFunctionsPhase = createFileLoweringPhase(
        { context, irFile ->
            irFile.acceptChildrenVoid(object : IrElementVisitorVoid {
                override fun visitElement(element: IrElement) {
@@ -113,15 +112,24 @@ private val extractLocalClassesFromInlineBodies = createFileLoweringPhase(
                }

                override fun visitFunction(declaration: IrFunction) {
-                    if (declaration.isInline && !declaration.isLocal)
-                        context.inlineFunctionsSupport.saveNonLoweredInlineFunction(declaration)
                    declaration.acceptChildrenVoid(this)
+
+                    if (declaration.isInline)
+                        OuterThisLowering(context).lower(declaration)
                }
            })
+        },
+        name = "LowerOuterThisInInlineFunctions",
+        description = "Lower outer this in inline functions"
+)

+private val extractLocalClassesFromInlineBodies = createFileLoweringPhase(
+        { context, irFile ->
            LocalClassesInInlineLambdasLowering(context).lower(irFile)
-            LocalClassesInInlineFunctionsLowering(context).lower(irFile)
-            LocalClassesExtractionFromInlineFunctionsLowering(context).lower(irFile)
+            if (!context.config.produce.isCache) {
+                LocalClassesInInlineFunctionsLowering(context).lower(irFile)
+                LocalClassesExtractionFromInlineFunctionsLowering(context).lower(irFile)
+            }
        },
        name = "ExtractLocalClassesFromInlineBodies",
        description = "Extraction of local classes from inline bodies",
@@ -322,6 +330,18 @@ private val inlinePhase = createFileLoweringPhase(
        lowering = { context: NativeGenerationState ->
            object : FileLoweringPass {
                override fun lower(irFile: IrFile) {
+                    irFile.acceptChildrenVoid(object : IrElementVisitorVoid {
+                        override fun visitElement(element: IrElement) {
+                            element.acceptChildrenVoid(this)
+                        }
+
+                        override fun visitFunction(declaration: IrFunction) {
+                            if (declaration.isInline)
+                                context.context.inlineFunctionsSupport.savePartiallyLoweredInlineFunction(declaration)
+                            declaration.acceptChildrenVoid(this)
+                        }
+                    })
+
                    FunctionInlining(
                            context.context,
                            NativeInlineFunctionResolver(context.context, context),
@@ -493,6 +513,7 @@ private fun PhaseEngine<NativeGenerationState>.getAllLowerings() = listOfNotNull
        arrayConstructorPhase,
        lateinitPhase,
        sharedVariablesPhase,
+        lowerOuterThisInInlineFunctionsPhase,
        inventNamesForLocalClasses,
        extractLocalClassesFromInlineBodies,
        wrapInlineDeclarationsWithReifiedTypeParametersLowering,
@@ -1960,7 +1960,7 @@ internal class CodeGeneratorVisitor(
    //-------------------------------------------------------------------------//
    private inner class ReturnableBlockScope(val returnableBlock: IrReturnableBlock, val resultSlot: LLVMValueRef?) :
            FileScope(returnableBlock.inlineFunctionSymbol?.owner?.let {
-                generationState.loweredInlineFunctions[it]?.irFile ?: it.fileOrNull
+                generationState.inlineFunctionOrigins[it]?.irFile ?: it.fileOrNull
            }
                    ?: (currentCodeContext.fileScope() as? FileScope)?.file
                    ?: error("returnable block should belong to current file at least")) {
@@ -1969,13 +1969,13 @@ internal class CodeGeneratorVisitor(
        var resultPhi : LLVMValueRef? = null
        private val functionScope by lazy {
            returnableBlock.inlineFunctionSymbol?.owner?.let {
-                it.scope(file().fileEntry.line(generationState.loweredInlineFunctions[it]?.startOffset ?: it.startOffset))
+                it.scope(file().fileEntry.line(generationState.inlineFunctionOrigins[it]?.startOffset ?: it.startOffset))
            }
        }

        private fun getExit(): LLVMBasicBlockRef {
            val location = returnableBlock.inlineFunctionSymbol?.owner?.let {
-                location(generationState.loweredInlineFunctions[it]?.endOffset ?: it.endOffset)
+                location(generationState.inlineFunctionOrigins[it]?.endOffset ?: it.endOffset)
            } ?: returnableBlock.statements.lastOrNull()?.let {
                location(it.endOffset)
            }
@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.backend.konan.lower
 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.Context
-import org.jetbrains.kotlin.backend.konan.InlineFunctionOriginInfo
-import org.jetbrains.kotlin.backend.konan.NativeGenerationState
-import org.jetbrains.kotlin.backend.konan.NativeMapping
+import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.ir.declarations.IrExternalPackageFragment
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
@@ -20,21 +17,17 @@ import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.util.*

 internal class InlineFunctionsSupport(mapping: NativeMapping) {
-    private val notLoweredInlineFunctions = mapping.notLoweredInlineFunctions
+    // Inline functions lowered up to just before the inliner.
+    private val partiallyLoweredInlineFunctions = mapping.partiallyLoweredInlineFunctions

-    fun saveNonLoweredInlineFunction(function: IrFunction) {
-        getNonLoweredInlineFunction(function, copy = false)
-    }
+    fun savePartiallyLoweredInlineFunction(function: IrFunction) =
+            function.deepCopyWithVariables().also {
+                it.patchDeclarationParents(function.parent)
+                partiallyLoweredInlineFunctions[function.symbol] = it
+            }

-    fun getNonLoweredInlineFunction(function: IrFunction, copy: Boolean): IrFunction {
-        val notLoweredInlineFunction = notLoweredInlineFunctions.getOrPut(function.symbol) {
-            function.deepCopyWithVariables().also { it.patchDeclarationParents(function.parent) }
-        }
-        return if (copy)
-            notLoweredInlineFunction.deepCopyWithVariables().also { it.patchDeclarationParents(function.parent) }
-        else
-            notLoweredInlineFunction
-    }
+    fun getPartiallyLoweredInlineFunction(function: IrFunction) =
+            partiallyLoweredInlineFunctions[function.symbol]
 }

 // TODO: This is a bit hacky. Think about adopting persistent IR ideas.
@@ -42,14 +35,21 @@ internal class NativeInlineFunctionResolver(override val context: Context, val g
    override fun getFunctionDeclaration(symbol: IrFunctionSymbol): IrFunction {
        val function = super.getFunctionDeclaration(symbol)

-        generationState.loweredInlineFunctions[function]?.let { return it.irFunction }
+        generationState.inlineFunctionOrigins[function]?.let { return it.irFunction }

        val packageFragment = function.getPackageFragment()
-        val (possiblyLoweredFunction, shouldLower) = if (packageFragment !is IrExternalPackageFragment) {
-            context.inlineFunctionsSupport.getNonLoweredInlineFunction(function, copy = context.config.producePerFileCache).also {
-                generationState.loweredInlineFunctions[function] =
-                        InlineFunctionOriginInfo(it, packageFragment as IrFile, function.startOffset, function.endOffset)
-            } to true
+        val functionIsNotFromLazyIr = packageFragment !is IrExternalPackageFragment
+        val irFile: IrFile
+        val (possiblyLoweredFunction, shouldLower) = if (functionIsNotFromLazyIr) {
+            irFile = packageFragment as IrFile
+            val partiallyLoweredFunction = context.inlineFunctionsSupport.getPartiallyLoweredInlineFunction(function)
+            if (partiallyLoweredFunction == null)
+                function to true
+            else {
+                generationState.inlineFunctionOrigins[function] =
+                        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
@@ -58,40 +58,49 @@ internal class NativeInlineFunctionResolver(override val context: Context, val g
            require(context.config.cachedLibraries.isLibraryCached(moduleDeserializer.klib)) {
                "No IR and no cache for ${function.render()}"
            }
-            val (shouldLower, deserializedInlineFunction) = moduleDeserializer.deserializeInlineFunction(function)
-            generationState.loweredInlineFunctions[function] = deserializedInlineFunction
-            function to shouldLower
+            val (firstAccess, deserializedInlineFunction) = moduleDeserializer.deserializeInlineFunction(function)
+            generationState.inlineFunctionOrigins[function] = deserializedInlineFunction
+            irFile = deserializedInlineFunction.irFile
+            function to firstAccess
        }

-        if (!shouldLower) return possiblyLoweredFunction
-
-        val body = possiblyLoweredFunction.body ?: return possiblyLoweredFunction
-
-        PreInlineLowering(context).lower(body, possiblyLoweredFunction, generationState.loweredInlineFunctions[function]!!.irFile)
-
-        ArrayConstructorLowering(context).lower(body, possiblyLoweredFunction)
-
-        NullableFieldsForLateinitCreationLowering(context).lowerWithLocalDeclarations(possiblyLoweredFunction)
-        NullableFieldsDeclarationLowering(context).lowerWithLocalDeclarations(possiblyLoweredFunction)
-        LateinitUsageLowering(context).lower(body, possiblyLoweredFunction)
-
-        SharedVariablesLowering(context).lower(body, possiblyLoweredFunction)
-
-        OuterThisLowering(context).lower(possiblyLoweredFunction)
-
-        LocalClassesInInlineLambdasLowering(context).lower(body, possiblyLoweredFunction)
-
-        if (generationState.llvmModuleSpecification.containsDeclaration(function)) {
-            // Do not extract local classes off of inline functions from cached libraries.
-            LocalClassesInInlineFunctionsLowering(context).lower(body, possiblyLoweredFunction)
-            LocalClassesExtractionFromInlineFunctionsLowering(context).lower(body, possiblyLoweredFunction)
+        if (shouldLower) {
+            lower(possiblyLoweredFunction, irFile, functionIsNotFromLazyIr)
+            if (functionIsNotFromLazyIr) {
+                generationState.inlineFunctionOrigins[function] =
+                        InlineFunctionOriginInfo(context.inlineFunctionsSupport.savePartiallyLoweredInlineFunction(possiblyLoweredFunction),
+                                irFile, function.startOffset, function.endOffset)
+            }
        }
-
-        WrapInlineDeclarationsWithReifiedTypeParametersLowering(context).lower(body, possiblyLoweredFunction)
-
        return possiblyLoweredFunction
    }

+    private fun lower(function: IrFunction, irFile: IrFile, functionIsNotFromLazyIr: Boolean) {
+        val body = function.body ?: return
+
+        PreInlineLowering(context).lower(body, function, irFile)
+
+        ArrayConstructorLowering(context).lower(body, function)
+
+        NullableFieldsForLateinitCreationLowering(context).lowerWithLocalDeclarations(function)
+        NullableFieldsDeclarationLowering(context).lowerWithLocalDeclarations(function)
+        LateinitUsageLowering(context).lower(body, function)
+
+        SharedVariablesLowering(context).lower(body, function)
+
+        OuterThisLowering(context).lower(function)
+
+        LocalClassesInInlineLambdasLowering(context).lower(body, function)
+
+        if (!context.config.produce.isCache && functionIsNotFromLazyIr) {
+            // Do not extract local classes off of inline functions from cached libraries.
+            LocalClassesInInlineFunctionsLowering(context).lower(body, function)
+            LocalClassesExtractionFromInlineFunctionsLowering(context).lower(body, function)
+        }
+
+        WrapInlineDeclarationsWithReifiedTypeParametersLowering(context).lower(body, function)
+    }
+
    private fun DeclarationTransformer.lowerWithLocalDeclarations(function: IrFunction) {
        if (transformFlat(function) != null)
            error("Unexpected transformation of function ${function.dump()}")