Lambda without parameters is inlined

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

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlin.backend.konan.lower
 
 import org.jetbrains.kotlin.backend.konan.Context
+import org.jetbrains.kotlin.backend.konan.descriptors.isFunctionInvoke
 import org.jetbrains.kotlin.backend.konan.ir.IrInlineFunctionBody
 import org.jetbrains.kotlin.backend.konan.ir.getArguments
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
@@ -11,10 +12,8 @@ import org.jetbrains.kotlin.ir.builders.Scope
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrContainerExpressionBase
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
-import org.jetbrains.kotlin.ir.expressions.impl.inlineStatement
 import org.jetbrains.kotlin.ir.util.DeepCopyIrTree
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.types.KotlinType
@@ -45,24 +44,32 @@ internal class FunctionInlining(val context: Context): IrElementTransformerVoid(
 
     //-------------------------------------------------------------------------//
 
+    fun getLambdaReturnType(value: IrExpression) : KotlinType {
+        val statements = (value as IrContainerExpressionBase).statements
+        val lambdaFunction = statements[0] as IrFunction
+        return lambdaFunction.descriptor.returnType!!
+    }
+
+    //-------------------------------------------------------------------------//
+
     fun evaluateParameters(callExpression: IrCall,
         statements: MutableList<IrStatement>): List<Pair<ParameterDescriptor, IrExpression>> {
 
-        val scope         = Scope(callExpression.descriptor as FunctionDescriptor)
+        val scope          = Scope(callExpression.descriptor as FunctionDescriptor)
-        val parametersOld = callExpression.getArguments()                                   // Create map inline_function_parameter -> containing_function_expression.
+        val parametersOld  = callExpression.getArguments()                                  // Create map inline_function_parameter -> containing_function_expression.
-        val parametersNew = parametersOld.map {
+        val parametersNew  = parametersOld.map {
-            val parameter = it.first
+            val parameter  = it.first
-            val value     = it.second
+            val expression = it.second
-            if (value is IrGetValue) return@map it                                          // If value is not an expression - do nothing.
+            if (expression is IrGetValue) return@map it                                     // There is nothing to evaluate.
-            if (isLambda(value)) {
+            if (isLambda(expression)) {                                                     // The expression is lambda.
-                val lambdaStatements = getLambdaStatements(value)
+                val inlineFunctionBody = inlineLambda(expression)                           // Create IrInlineFunctionBody to replace this parameter.
-                return@map it
+                return@map parameter to inlineFunctionBody
-            }                     // If value is lambda - d
+            }
 
-            val newVar = scope.createTemporaryVariable(value, "inline", false)              // Create new variable and init it with the expression.
+            val newVar = scope.createTemporaryVariable(expression, "inline", false)         // Create new variable and init it with the expression.
             statements.add(0, newVar)                                                       // Add initialization of the new variable in statement list.
 
-            val getVal = IrGetValueImpl(0, 0, newVar.descriptor)                            // Create new IR element representing access the new variable
+            val getVal = IrGetValueImpl(0, 0, newVar.descriptor)                            // Create new IR element representing access the new variable.
             parameter to getVal                                                             // Parameter will be replaced with the new variable.
         }
         return parametersNew
@@ -83,15 +90,26 @@ internal class FunctionInlining(val context: Context): IrElementTransformerVoid(
         val inlineBody  = copyFuncDeclaration.body!! as IrBlockBody
         val statements  = inlineBody.statements
         val parameters  = evaluateParameters(irCall, statements)                            // Evaluate parameters representing expression.
+        val irBlock     = IrInlineFunctionBody(startOffset, endOffset, returnType, null, statements)
 
         val transformer = ParametersTransformer(parameters)
-        val irBlock     = IrInlineFunctionBody(startOffset, endOffset, returnType, null, statements)
         irBlock.accept(transformer, null)                                                   // Replace parameters with expression.
         return irBlock
     }
 
     //-------------------------------------------------------------------------//
 
+    fun inlineLambda(value: IrExpression): IrBlock {
+
+        val lambdaStatements = getLambdaStatements(value)
+        val lambdaReturnType = getLambdaReturnType(value)
+        val irBlock     = IrInlineFunctionBody(0, 0, lambdaReturnType, null, lambdaStatements)
+
+        return irBlock
+    }
+
+    //-------------------------------------------------------------------------//
+
     override fun visitCall(expression: IrCall): IrExpression {
         val functionDescriptor = expression.descriptor as FunctionDescriptor
         if (!functionDescriptor.isInline) return super.visitCall(expression)                // Function is not to be inlined - do nothing.
@@ -114,6 +132,19 @@ internal class ParametersTransformer(val parameterToExpression: List <Pair<Param
 
     //-------------------------------------------------------------------------//
 
+    override fun visitCall(expression: IrCall): IrExpression {
+        if ((expression.descriptor as FunctionDescriptor).isFunctionInvoke) {               // If it is lambda call.
+            if (expression.dispatchReceiver !is IrGetValue) super.visitCall(expression)     // Do not process such dispatch receiver.
+            val getValue = expression.dispatchReceiver as IrGetValue                        //
+            val parExpr = parameterToExpression.find { it.first == getValue.descriptor }    // Find expression to replace this parameter.
+            if (parExpr == null) super.visitCall(expression)                                // It is not function parameter.
+            return parExpr!!.second                                                         // Replace call site with InlineFunctionBody.
+        }
+        return super.visitCall(expression)                                                  // Function is declared in another module.
+    }
+
+    //-------------------------------------------------------------------------//
+
     override fun visitGetValue(expression: IrGetValue): IrExpression {
         val descriptor = expression.descriptor
         if (descriptor !is ParameterDescriptor) {                                           // TODO do we need this check?
@@ -122,6 +153,6 @@ internal class ParametersTransformer(val parameterToExpression: List <Pair<Param
         val parExp = parameterToExpression.find { it.first == descriptor }                  // Find expression to replace this parameter.
         if (parExp == null) return super.visitGetValue(expression)
 
-        return  parExp.second             // TODO should we proceed with IR iteration here?
+        return  parExp.second                                                               // TODO should we proceed with IR iteration here?
     }
 }