Inlining enabled for functions with arguments

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

@@ -1,38 +1,75 @@
 package org.jetbrains.kotlin.backend.konan.lower
 
 import org.jetbrains.kotlin.backend.konan.Context
+import org.jetbrains.kotlin.backend.konan.ir.getArguments
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.ParameterDescriptor
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
+import org.jetbrains.kotlin.ir.util.DeepCopyIrTree
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 
-internal class FunctionInlining(val context: Context) {
+//-----------------------------------------------------------------------------//
-    fun inline(irFile: IrFile) {
-        irFile.accept(object : IrElementTransformerVoid() {
-            override fun visitElement(element: IrElement) = element.accept(this, null)
 
-            override fun visitCall(expression: IrCall): IrExpression {
+internal class FunctionInlining(val context: Context): IrElementTransformerVoid() {
-                val functionDescriptor = expression.descriptor as FunctionDescriptor                //
-                if (functionDescriptor.isInline == false) return super.visitCall(expression)        // function is not to be inlined - do nothing
 
-                if (!functionDescriptor.name.asString().contains("foo")) return super.visitCall(expression)
+    fun inline(irFile: IrFile) = irFile.accept(this, null)
 
-                val functionDeclaration = context.ir.moduleIndex.functions[functionDescriptor]!!    // get FunctionDeclaration by FunctionDescriptor
+    //-------------------------------------------------------------------------//
-                val body = functionDeclaration.body!! as IrBlockBodyImpl
 
-                val startOffset = functionDeclaration.startOffset
+    fun removeReturn(statements: MutableList<IrStatement>) =
-                val endOffset   = functionDeclaration.endOffset
+        statements.map { statement ->
-                val returnType  = functionDeclaration.descriptor.returnType!!
+            (statement as? IrReturn)?.value ?: statement
-                val statements  = body.statements
+        } as MutableList<IrStatement>
-                val irBlock = IrBlockImpl(startOffset, endOffset, returnType, null, statements)     // create
-                return irBlock
-            }
-        }, null)
 
+    //-------------------------------------------------------------------------//
+
+    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
+
+        val functionDeclaration = context.ir.moduleIndex.functions[functionDescriptor]      // get FunctionDeclaration by FunctionDescriptor
+        if (functionDeclaration == null) return super.visitCall(expression)                 // TODO what if we do not have declaration?
+        val copyFuncDeclaration = functionDeclaration.accept(DeepCopyIrTree(), null) as IrFunction
+
+        val body        = copyFuncDeclaration.body!! as IrBlockBody
+        val statements  = removeReturn(body.statements)
+        val startOffset = copyFuncDeclaration.startOffset
+        val endOffset   = copyFuncDeclaration.endOffset
+        val returnType  = copyFuncDeclaration.descriptor.returnType!!
+        val irBlock     = IrBlockImpl(startOffset, endOffset, returnType, null, statements) // create
+
+        val parameterToExpression = expression.getArguments()
+        val parametersTransformer = ParametersTransformer(parameterToExpression)
+        irBlock.accept(parametersTransformer, null)
+
+        return irBlock
     }
+
+    //-------------------------------------------------------------------------//
+
+    override fun visitElement(element: IrElement) = element.accept(this, null)
 }
+
+//-----------------------------------------------------------------------------//
+
+internal class ParametersTransformer(val parameterToExpression: List <Pair<ParameterDescriptor, IrExpression>>): IrElementTransformerVoid() {
+
+    override fun visitElement(element: IrElement) = element.accept(this, null)
+
+    //-------------------------------------------------------------------------//
+
+    override fun visitGetValue(expression: IrGetValue): IrExpression {
+        val descriptor = expression.descriptor
+        if (descriptor !is ParameterDescriptor) {
+            return super.visitGetValue(expression)
+        }
+        val parExp = parameterToExpression.find { it.first == descriptor }
+        return parExp?.let { parExp.second } ?: super.visitGetValue(expression)
+    }
+}