Optimized inliner: substitute leaking descriptors only once

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/common/DeepCopyIrTreeWithDescriptoros.kt

@@ -23,14 +23,12 @@ import org.jetbrains.kotlin.backend.konan.ir.IrReturnableBlockImpl
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.*
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
 import org.jetbrains.kotlin.ir.descriptors.IrTemporaryVariableDescriptorImpl
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.util.DeepCopyIrTree
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.name.Name
@@ -41,7 +39,7 @@ import org.jetbrains.kotlin.types.TypeSubstitutor
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.typeUtil.makeNullable
 
-internal class DeepCopyIrTreeWithDescriptors(val targetDescriptor: DeclarationDescriptor,
+internal class DeepCopyIrTreeWithDescriptors(val targetDescriptor: FunctionDescriptor,
                                              val context: Context) {
 
     private val descriptorSubstituteMap: MutableMap<DeclarationDescriptor, DeclarationDescriptor> = mutableMapOf()
@@ -51,7 +49,6 @@ internal class DeepCopyIrTreeWithDescriptors(val targetDescriptor: DeclarationDe
     //-------------------------------------------------------------------------//
 
     fun copy(irElement: IrElement, typeSubstitutor: TypeSubstitutor?): IrElement {
-
         this.typeSubstitutor = typeSubstitutor
         irElement.acceptChildrenVoid(DescriptorCollector())
         return irElement.accept(InlineCopyIr(), null)
@@ -550,51 +547,51 @@ internal class DeepCopyIrTreeWithDescriptors(val targetDescriptor: DeclarationDe
 
     //-------------------------------------------------------------------------//
 
-    val descriptorSubstitutorForExternalScope = object : IrElementTransformerVoid() {
+    fun addCurrentSubstituteMap(globalSubstituteMap: MutableMap<DeclarationDescriptor, SubstitutedDescriptor>) {
-
+        descriptorSubstituteMap.forEach { t, u ->
-        override fun visitFunction(declaration: IrFunction): IrStatement {
+            globalSubstituteMap.put(t, SubstitutedDescriptor(targetDescriptor, u))
-            if (declaration.descriptor == targetDescriptor)
-                return declaration
-            return super.visitFunction(declaration)
         }
+    }
 
-        override fun visitCall(expression: IrCall): IrExpression {
+}
-            val oldExpression = super.visitCall(expression) as IrCall
 
-            return when (oldExpression) {
+internal class SubstitutedDescriptor(val inlinedFunction: FunctionDescriptor, val descriptor: DeclarationDescriptor)
-                is IrCallImpl -> copyIrCallImpl(oldExpression)
+
-                is IrCallWithShallowCopy -> copyIrCallWithShallowCopy(oldExpression)
+internal class DescriptorSubstitutorForExternalScope(val globalSubstituteMap: MutableMap<DeclarationDescriptor, SubstitutedDescriptor>)
-                else -> oldExpression
+    : IrElementTransformerVoidWithContext() {
-            }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        val oldExpression = super.visitCall(expression) as IrCall
+
+        val substitutedDescriptor = globalSubstituteMap[expression.descriptor.original]
+                ?: return oldExpression
+        if (allScopes.any { it.scope.scopeOwner == substitutedDescriptor.inlinedFunction })
+            return oldExpression
+        return when (oldExpression) {
+            is IrCallImpl -> copyIrCallImpl(oldExpression, substitutedDescriptor)
+            is IrCallWithShallowCopy -> copyIrCallWithShallowCopy(oldExpression, substitutedDescriptor)
+            else -> oldExpression
         }
     }
 
     //-------------------------------------------------------------------------//
 
-    private fun copyIrCallImpl(oldExpression: IrCallImpl): IrCallImpl {
+    private fun copyIrCallImpl(oldExpression: IrCallImpl, substitutedDescriptor: SubstitutedDescriptor): IrCallImpl {
 
         val oldDescriptor = oldExpression.descriptor
-        val newDescriptor = (descriptorSubstituteMap[oldDescriptor.original] ?: oldDescriptor) as FunctionDescriptor
+        val newDescriptor = substitutedDescriptor.descriptor as FunctionDescriptor
 
-        val oldSuperQualifier = oldExpression.superQualifier
+        if (newDescriptor == oldDescriptor)
-        val newSuperQualifier = oldSuperQualifier?.let { (descriptorSubstituteMap[it] ?: it) as ClassDescriptor }
-
-        val newType = substituteType(oldExpression.type)!!
-        val newTypeArguments = substituteTypeArguments(oldExpression.typeArguments)
-
-        if (newDescriptor == oldDescriptor && newSuperQualifier == oldSuperQualifier && newType == oldExpression.type
-                && ((newTypeArguments == null && oldExpression.typeArguments == null)
-                    || (newTypeArguments!!.all { it.value == oldExpression.typeArguments!![it.key] })))
             return oldExpression
 
         val newExpression = IrCallImpl(
-            startOffset              = oldExpression.startOffset,
+                startOffset              = oldExpression.startOffset,
-            endOffset                = oldExpression.endOffset,
+                endOffset                = oldExpression.endOffset,
-            type                     = newType,
+                type                     = oldExpression.type,
-            calleeDescriptor         = newDescriptor,
+                calleeDescriptor         = newDescriptor,
-            typeArguments            = newTypeArguments,
+                typeArguments            = oldExpression.typeArguments,
-            origin                   = oldExpression.origin,
+                origin                   = oldExpression.origin,
-            superQualifierDescriptor = newSuperQualifier
+                superQualifierDescriptor = oldExpression.superQualifier
         ).apply {
             oldExpression.descriptor.valueParameters.forEach {
                 val valueArgument = oldExpression.getValueArgument(it)
@@ -609,19 +606,15 @@ internal class DeepCopyIrTreeWithDescriptors(val targetDescriptor: DeclarationDe
 
     //-------------------------------------------------------------------------//
 
-    private fun copyIrCallWithShallowCopy(oldExpression: IrCallWithShallowCopy): IrCall {
+    private fun copyIrCallWithShallowCopy(oldExpression: IrCallWithShallowCopy, substitutedDescriptor: SubstitutedDescriptor): IrCall {
 
         val oldDescriptor = oldExpression.descriptor
-        val newDescriptor = (descriptorSubstituteMap[oldDescriptor.original] ?: oldDescriptor) as FunctionDescriptor
+        val newDescriptor = substitutedDescriptor.descriptor as FunctionDescriptor
 
-        val oldSuperQualifier = oldExpression.superQualifier
+        if (newDescriptor == oldDescriptor)
-        val newSuperQualifier = oldSuperQualifier?.let { (descriptorSubstituteMap[it] ?: it) as ClassDescriptor }
-
-        if (newDescriptor == oldDescriptor && newSuperQualifier == oldSuperQualifier)
             return oldExpression
 
-        return oldExpression.shallowCopy(oldExpression.origin, newDescriptor, newSuperQualifier)
+        return oldExpression.shallowCopy(oldExpression.origin, newDescriptor, oldExpression.superQualifier)
     }
 }
 
-

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/common/IrElementTransformerVoidWithContext.kt

@@ -71,6 +71,7 @@ abstract internal class IrElementTransformerVoidWithContext : IrElementTransform
     protected val currentProperty get() = scopeStack.lastOrNull { it.scope.scopeOwner is PropertyDescriptor }
     protected val currentScope get() = scopeStack.peek()
     protected val parentScope get() = if (scopeStack.size < 2) null else scopeStack[scopeStack.size - 2]
+    protected val allScopes get() = scopeStack
 
     fun printScopeStack() {
         scopeStack.forEach { println(it.scope.scopeOwner) }

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

@@ -25,12 +25,12 @@ import org.jetbrains.kotlin.backend.konan.descriptors.needsInlining
 import org.jetbrains.kotlin.backend.konan.descriptors.resolveFakeOverride
 import org.jetbrains.kotlin.backend.konan.ir.DeserializerDriver
 import org.jetbrains.kotlin.backend.konan.ir.IrReturnableBlockImpl
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ValueDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
-import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.getDefault
@@ -50,10 +50,17 @@ import org.jetbrains.kotlin.types.TypeSubstitutor
 internal class FunctionInlining(val context: Context): IrElementTransformerVoidWithContext() {
 
     private val deserializer = DeserializerDriver(context)
+    private val globalSubstituteMap = mutableMapOf<DeclarationDescriptor, SubstitutedDescriptor>()
 
     //-------------------------------------------------------------------------//
 
-    fun inline(irModule: IrModuleFragment) = irModule.accept(this, null)
+    fun inline(irModule: IrModuleFragment): IrElement {
+        val transformedModule = irModule.accept(this, null)
+        transformedModule.transformChildrenVoid(
+                DescriptorSubstitutorForExternalScope(globalSubstituteMap)                  // Transform calls to object that might be returned from inline function call.
+        )
+        return transformedModule
+    }
 
     //-------------------------------------------------------------------------//
 
@@ -72,7 +79,7 @@ internal class FunctionInlining(val context: Context): IrElementTransformerVoidW
         }
 
         functionDeclaration.transformChildrenVoid(this)                                     // Process recursive inline.
-        val inliner = Inliner(currentFile, functionDeclaration, currentScope!!, context)    // Create inliner for this scope.
+        val inliner = Inliner(globalSubstituteMap, functionDeclaration, currentScope!!, context)    // Create inliner for this scope.
         return inliner.inline(irCall )                                  // Return newly created IrInlineBody instead of IrCall.
     }
 
@@ -95,7 +102,7 @@ internal class FunctionInlining(val context: Context): IrElementTransformerVoidW
 
 //-----------------------------------------------------------------------------//
 
-private class Inliner(val irFile: IrFile,
+private class Inliner(val globalSubstituteMap: MutableMap<DeclarationDescriptor, SubstitutedDescriptor>,
                       val functionDeclaration: IrFunction,                                  // Function to substitute.
                       val currentScope: ScopeWithIr,
                       val context: Context) {
@@ -107,8 +114,7 @@ private class Inliner(val irFile: IrFile,
 
     fun inline(irCall: IrCall): IrReturnableBlockImpl {                                     // Call to be substituted.
         val inlineFunctionBody = inlineFunction(irCall, functionDeclaration)
-        val descriptorSubstitutor = copyIrElement.descriptorSubstitutorForExternalScope
+        copyIrElement.addCurrentSubstituteMap(globalSubstituteMap)
-        irFile.transformChildrenVoid(descriptorSubstitutor)                                 // Transform calls to object that might be returned from inline function call.
         return inlineFunctionBody
     }
 

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

@@ -1166,7 +1166,7 @@ internal class IrDeserializer(val context: Context,
             (key,value) ->
         key to value}
 
-        val copyFunctionDeclaration = DeepCopyIrTreeWithDescriptors(rootFunction.containingDeclaration, context).copy(
+        val copyFunctionDeclaration = DeepCopyIrTreeWithDescriptors(rootFunction, context).copy(
             irElement       = declaration,      
             typeSubstitutor = TypeSubstitutor.create(substitutionContext)  
         ) as IrFunction